From d84b4282b3e0d4bc67830ed5d8578f1425ccc8a4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 23 Jun 2026 14:16:20 -0700 Subject: [PATCH] feat(ocr): RustBridge with ocr+aocr; provider-gated rust path --- litellm/ocr/rust_bridge.py | 85 ++++++++++++++++++++++++++------------ 1 file changed, 58 insertions(+), 27 deletions(-) diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 61f9e8ca69a..81b22a8c622 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -1,56 +1,81 @@ """ Optional Rust-backed OCR path. -Enable with ``litellm.use_litellm_rust()``; the sync ``litellm.ocr()`` entrypoint -then routes supported Mistral calls through the compiled ``litellm_python_bridge`` +Enable with ``litellm.use_litellm_rust()``; ``litellm.ocr()`` / ``litellm.aocr()`` +then route supported providers through the compiled ``litellm_python_bridge`` extension, which performs the whole OCR call (URL, headers, HTTP, parse) in Rust. +The bridge exposes two entry points that mirror the Python API: + +* ``ocr(...)`` blocks on the async core with the GIL released (sync SDK callers). +* ``aocr(...)`` returns a Python awaitable driven by a Tokio runtime, so the + proxy can ``await`` it without tying up a thread-pool worker per request. + No module-level ``litellm`` imports keep this a leaf so ``litellm/ocr/main.py`` can import it statically without forming an import cycle. """ from __future__ import annotations -from typing import Final, Protocol, cast +from typing import Any, Awaitable, Final, Protocol, Union, cast -class RustOcr(Protocol): - """Signature of the compiled ``litellm_python_bridge.ocr`` entrypoint.""" +class RustBridge(Protocol): + """The compiled ``litellm_python_bridge`` surface used by the OCR path.""" - def __call__( + def ocr( self, + provider: str, model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: ... + document: dict[str, Any], + api_key: str | None = None, + api_base: str | None = None, + extra_headers: dict[str, str] | None = None, + timeout_seconds: float | None = None, + params: dict[str, Any] | None = None, + ) -> dict[str, Any]: ... + + def aocr( + self, + provider: str, + model: str, + document: dict[str, Any], + api_key: str | None = None, + api_base: str | None = None, + extra_headers: dict[str, str] | None = None, + timeout_seconds: float | None = None, + params: dict[str, Any] | None = None, + ) -> Awaitable[dict[str, Any]]: ... + + +# Providers whose full OCR call the Rust core can handle today. Grows one entry +# per provider PR; everything else stays on the Python path. +RUST_SUPPORTED_PROVIDERS: Final[frozenset[str]] = frozenset({"mistral"}) class _Unset: - """Sentinel type so ``ocr=None`` can clear a prior injection while omission preserves it.""" + """Sentinel so ``bridge=None`` can clear a prior injection while omission preserves it.""" _UNSET: Final[_Unset] = _Unset() _rust_ocr_enabled = False -_rust_ocr_impl: RustOcr | None = None +_rust_bridge: RustBridge | None = None def use_litellm_rust( - enabled: bool = True, *, ocr: RustOcr | None | _Unset = _UNSET + enabled: bool = True, *, bridge: Union[RustBridge, None, _Unset] = _UNSET ) -> None: """Route supported OCR calls through the Rust ``litellm_python_bridge`` extension. - ``ocr`` injects the bridge callable; when omitted the compiled extension is - loaded on demand and any previously injected bridge is preserved. Pass - ``ocr=None`` explicitly to clear a prior injection. + ``bridge`` injects the bridge object (mainly for tests); when omitted the + compiled extension is loaded on demand and any previously injected bridge is + preserved. Pass ``bridge=None`` explicitly to clear a prior injection. """ - global _rust_ocr_enabled, _rust_ocr_impl + global _rust_ocr_enabled, _rust_bridge _rust_ocr_enabled = enabled - if not isinstance(ocr, _Unset): - _rust_ocr_impl = ocr + if not isinstance(bridge, _Unset): + _rust_bridge = bridge def rust_ocr_enabled() -> bool: @@ -58,17 +83,23 @@ def rust_ocr_enabled() -> bool: return _rust_ocr_enabled -def load_rust_ocr() -> RustOcr | None: - """Return the Rust OCR callable, or ``None`` when no bridge is available. +def rust_supports(provider: str) -> bool: + """Whether the Rust core can handle this provider's OCR call end to end.""" + return provider in RUST_SUPPORTED_PROVIDERS - Prefers an injected implementation, otherwise loads the compiled + +def load_rust_bridge() -> RustBridge | None: + """Return the Rust bridge, or ``None`` when no bridge is available. + + Prefers an injected bridge, otherwise loads the compiled ``litellm_python_bridge`` extension; a missing extension yields ``None`` so the caller can fall back to the Python path instead of hard-failing. """ - if _rust_ocr_impl is not None: - return _rust_ocr_impl + if _rust_bridge is not None: + return _rust_bridge try: import litellm_python_bridge except ImportError: return None - return cast(RustOcr, litellm_python_bridge.ocr) + # The compiled module exposes ocr / aocr / gil_stats at module level. + return cast(RustBridge, litellm_python_bridge)