feat: make rust OCR async-first - #31253
Conversation
|
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
|
@greptile review |
6906840 to
2977e2b
Compare
|
@greptile review again |
Greptile SummaryThis PR makes the Rust OCR path async-first by converting
Confidence Score: 4/5The core async refactor is solid and well-tested, but the RustOcr protocol signature change breaks any caller who previously injected a custom sync bridge via use_litellm_rust(ocr=…). The new RustOcr protocol adds custom_llm_provider and extra_headers as required keyword arguments, inserted in the middle of the parameter list. Any user who injected their own callable via use_litellm_rust(ocr=my_bridge) matching the prior six-parameter signature will receive a runtime TypeError with no advance warning — the Protocol mismatch is invisible to static checkers until the bridge is actually called. litellm/ocr/rust_bridge.py — the RustOcr protocol signature change is the main concern; litellm/ocr/main.py is otherwise clean
|
| Filename | Overview |
|---|---|
| litellm/ocr/main.py | Major refactor: extracts _prepare_ocr_request / _prepare_rust_ocr_call helpers and adds async _run_rust_aocr path; exception handler now uses resolved model/provider for the common case; moves ocr() after convert_file_document_to_url_document in file order |
| litellm/ocr/rust_bridge.py | Adds RustAocr Protocol, _rust_aocr_impl global, load_rust_aocr(), and extends use_litellm_rust() with aocr= parameter; RustOcr protocol signature change is backwards-incompatible for existing injected bridges |
| litellm-rust/crates/providers/src/ocr.rs | Converts blocking run_ocr to async ocr(OcrRequest), adds OcrRequest struct, string_headers/has_authorization_header helpers, and ocr_config_for provider dispatch; well-tested with local loopback server |
| litellm-rust/crates/python-bridge/src/lib.rs | Adds aocr PyO3 export via future_into_py; refactors sync ocr() to block_on the shared Tokio runtime; extracts marshal_inputs helper to avoid duplication between sync and async paths |
| tests/test_litellm/ocr/test_rust_bridge.py | Adds RecordingAsyncBridge, RaisingBridge/RaisingAsyncBridge, build_prepared_request factory, and new async test coverage; all tests use injected fakes with no real network calls |
| litellm/llms/mistral/ocr/transformation.py | Extracts MISTRAL_OCR_API_KEY_ENV_VAR constant and adds get_api_key_env_var() to MistralOCRConfig; straightforward refactor |
| litellm/proxy/ocr_endpoints/endpoints.py | Moves top-level import of convert_file_document_to_url_document/get_mime_type into the function body to avoid circular import; no behaviour change |
| litellm/llms/base_llm/ocr/transformation.py | Adds default get_api_key_env_var() returning None to BaseOCRConfig; clean extension point |
Reviews (9): Last reviewed commit: "fix: avoid duplicate Rust OCR authorizat..." | Re-trigger Greptile
1 similar comment
|
@greptile review |
| @@ -23,9 +23,29 @@ def __call__( | |||
| document: dict[str, object], | |||
| api_key: str | None, | |||
| api_base: str | None, | |||
| custom_llm_provider: str, | |||
| extra_headers: dict[str, object] | None, | |||
| optional_params: dict[str, object], | |||
| timeout_seconds: float | None, | |||
| ) -> dict[str, object]: ... | |||
| ) -> dict[str, object]: | |||
| raise NotImplementedError | |||
There was a problem hiding this comment.
Breaking change to
RustOcr protocol signature
The RustOcr protocol now requires custom_llm_provider: str and extra_headers: dict[str, object] | None as keyword arguments (inserted between api_base and optional_params). Any callable injected via use_litellm_rust(ocr=my_bridge) that matched the previous 6-argument signature (model, document, api_key, api_base, optional_params, timeout_seconds) will raise TypeError: __call__() got an unexpected keyword argument 'custom_llm_provider' at the point _run_rust_ocr calls it. Since use_litellm_rust is part of the public API, existing users who injected custom bridges will silently break at runtime without any indication from the type checker (the Protocol mismatch only surfaces at call time).
Rule Used: What: avoid backwards-incompatible changes without... (source)
Relevant issues
N/A
Linear ticket
N/A
What is changing
This PR makes the Rust OCR path async-first and keeps the Python bridge compatible with both sync and async LiteLLM callers.
reqwest::blocking::Clientto asyncreqwest::Client.litellm_python_bridge.aocr(), backed by the shared Tokio runtime.litellm_python_bridge.ocr()wrapper for compatibility, but it now blocks on the async route while releasing the GIL.litellm.aocr()to await Rust OCR directly for Mistral instead of routing throughrun_in_executor(ocr)._prepare_rust_ocr_call()so sync and async paths share key resolution, environment validation, complete URL resolution, headers, and pre-call logging.provider_endpoints_support.jsoninstead of maintaining a Rust-only provider registry.default_credsmetadata for stable provider defaults, currently Mistral's default API base and API key env var.Why
The previous async Python OCR path used Rust through a sync bridge in an executor. That worked, but it kept two Rust transport modes alive and made future provider growth harder. This change gives Rust routes one async execution model while preserving the public sync API.
The provider metadata cleanup avoids a second provider list as Rust grows toward more providers. Rust now generates typed provider metadata from the same provider/endpoints support JSON used elsewhere in LiteLLM, with provider defaults kept in a small
default_credsmap.Pre-Submission checklist
make test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewCI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Screenshots / Proof of Fix
Local verification:
OCR function smoke:
Note: this environment does not have
MISTRAL_API_KEYset, so I could not complete a live Mistral OCR provider call. The native bridge and public sync/async OCR dispatch were verified locally up to the expected missing-key boundary.Type
🆕 New Feature
✅ Test
Review notes
reqwest::Clientbecause Rust cannot call Python'sBaseLLMHTTPHandlerdirectly. The route remains Mistral-only and opt-in behinduse_litellm_rust.