Skip to content

feat: make rust OCR async-first - #31253

Merged
ishaan-berri merged 16 commits into
mainfrom
litellm_hotfix_ocr_async_rust
Jun 25, 2026
Merged

feat: make rust OCR async-first#31253
ishaan-berri merged 16 commits into
mainfrom
litellm_hotfix_ocr_async_rust

Conversation

@ishaan-berri

Copy link
Copy Markdown
Contributor

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.

  • Converts Rust OCR provider I/O from reqwest::blocking::Client to async reqwest::Client.
  • Adds an awaitable PyO3 bridge export, litellm_python_bridge.aocr(), backed by the shared Tokio runtime.
  • Keeps the sync litellm_python_bridge.ocr() wrapper for compatibility, but it now blocks on the async route while releasing the GIL.
  • Updates litellm.aocr() to await Rust OCR directly for Mistral instead of routing through run_in_executor(ocr).
  • Extracts shared Python Rust-OCR setup into _prepare_rust_ocr_call() so sync and async paths share key resolution, environment validation, complete URL resolution, headers, and pre-call logging.
  • Moves Rust provider metadata generation to the existing repo-root provider_endpoints_support.json instead of maintaining a Rust-only provider registry.
  • Adds top-level default_creds metadata for stable provider defaults, currently Mistral's default API base and API key env var.
  • Adds provider docs guidance for Rust provider calls: route I/O should stay async/Tokio-based, with pure transform/reverse-transform/response-transform code around transport.

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_creds map.

Pre-Submission checklist

  • I have added meaningful tests
  • My PR passes all unit tests on make test-unit
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

CI (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:

cargo fmt --check && cargo clippy --workspace --all-targets --locked -- -D warnings && cargo test --workspace --locked
result: passed

uv run pytest tests/test_litellm/ocr/test_rust_bridge.py -q
result: 19 passed

uv run python tests/code_coverage_tests/check_provider_folders_documented.py
result: passed

python -m json.tool provider_endpoints_support.json
python -m json.tool litellm/provider_endpoints_support_backup.json
result: passed

git fetch upstream main && git rebase upstream/main
result: branch already up to date; no merge conflicts against PR base

OCR function smoke:

Compiled native bridge with:
PYO3_PYTHON=.venv/bin/python RUSTFLAGS='-C link-arg=-undefined -C link-arg=dynamic_lookup' cargo build -p litellm-python-bridge

Native bridge import/export check:
bridge_exports=ocr,aocr,gil_stats
sync_error_type=ValueError
sync_error_contains_invalid_provider=True
aocr_is_awaitable=True
async_error_type=ValueError
async_error_contains_invalid_provider=True

Top-level OCR smoke with injected Rust bridges:
sync_bridge_provider=mistral
sync_bridge_headers={'x-trace-id': 'trace-1'}
sync_response=ocr ok
async_bridge_provider=mistral
async_bridge_headers={'x-trace-id': 'trace-1'}
async_response=async ocr ok

Top-level OCR smoke with compiled native bridge and no provider key:
top_level_sync_native_error_type=APIConnectionError
top_level_sync_native_missing_key=True
top_level_async_native_error_type=APIConnectionError
top_level_async_native_missing_key=True

Note: this environment does not have MISTRAL_API_KEY set, 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

  • The Rust OCR HTTP path still uses a Rust reqwest::Client because Rust cannot call Python's BaseLLMHTTPHandler directly. The route remains Mistral-only and opt-in behind use_litellm_rust.
  • The Rust sync API is intentionally a compatibility wrapper. New Rust provider route I/O should be async, with sync behavior only at the Python bridge boundary.

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

@ishaan-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@codspeed-hq

codspeed-hq Bot commented Jun 24, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 16 untouched benchmarks


Comparing litellm_hotfix_ocr_async_rust (554cf88) with main (3818d64)

Open in CodSpeed

Comment thread litellm/ocr/main.py
Comment thread provider_endpoints_support.json
@codecov

codecov Bot commented Jun 25, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.54839% with 8 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
litellm/ocr/main.py 96.03% 4 Missing ⚠️
litellm/ocr/rust_bridge.py 87.50% 2 Missing ⚠️
litellm/llms/base_llm/ocr/transformation.py 50.00% 1 Missing ⚠️
litellm/llms/mistral/ocr/transformation.py 75.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@ishaan-berri

Copy link
Copy Markdown
Contributor Author

@greptile review

@ishaan-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@ishaan-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

Comment thread litellm/ocr/main.py Fixed
Comment thread litellm/ocr/main.py Fixed
Comment thread litellm/ocr/main.py Fixed
@ishaan-berri
ishaan-berri force-pushed the litellm_hotfix_ocr_async_rust branch from 6906840 to 2977e2b Compare June 25, 2026 00:27
@BerriAI BerriAI deleted a comment from greptile-apps Bot Jun 25, 2026
@BerriAI BerriAI deleted a comment from greptile-apps Bot Jun 25, 2026
@ishaan-berri

Copy link
Copy Markdown
Contributor Author

@greptile review again

@greptile-apps

greptile-apps Bot commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR makes the Rust OCR path async-first by converting reqwest::blocking::Client to the async reqwest::Client, adding an aocr() PyO3 export backed by the shared Tokio runtime, and refactoring litellm.aocr() to await the Rust bridge directly instead of routing through run_in_executor(ocr). A shared _prepare_ocr_request() helper deduplicates provider resolution, logging, and parameter mapping between the sync and async Python paths.

  • Async bridge: New litellm_providers::ocr::ocr is fully async; the sync ocr() bridge now uses block_on on the Tokio runtime while releasing the GIL, and a new aocr() export uses future_into_py to return an asyncio-compatible awaitable.
  • Shared preparation: _prepare_ocr_request and _prepare_rust_ocr_call centralize provider resolution, key lookup, header construction, and pre_call logging so sync and async Python callers walk the same setup path.
  • Exception context fix: completion_kwargs is now updated with the resolved model and custom_llm_provider after _prepare_ocr_request succeeds, so exceptions thrown by the bridge are classified against the correct provider.

Confidence Score: 4/5

The 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

Important Files Changed

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

@ishaan-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

1 similar comment
@ishaan-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@ishaan-berri

Copy link
Copy Markdown
Contributor Author

@greptile review

@ishaan-berri
ishaan-berri enabled auto-merge June 25, 2026 01:15
Comment on lines 17 to +31
@@ -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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 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)

@ishaan-berri
ishaan-berri merged commit bd2a165 into main Jun 25, 2026
119 of 123 checks passed
@ishaan-berri
ishaan-berri deleted the litellm_hotfix_ocr_async_rust branch June 25, 2026 01:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants