feat: package Rust OCR bridge in LiteLLM wheel

.circleci/config.yml

@@ -205,6 +205,24 @@ jobs:
           environment:
             UV_HTTP_TIMEOUT: "300"
           command: |
+            $rustupInit = Join-Path $env:TEMP "rustup-init.exe"
+            $rustupVersion = "1.28.2"
+            $rustupUrl = "https://static.rust-lang.org/rustup/archive/$rustupVersion/x86_64-pc-windows-msvc/rustup-init.exe"
+            Invoke-WebRequest -Uri $rustupUrl -OutFile $rustupInit
+            $rustupExpected = "88d8258dcf6ae4f7a80c7d1088e1f36fa7025a1cfd1343731b4ee6f385121fc0"
+            $rustupActual = (Get-FileHash -Path $rustupInit -Algorithm SHA256).Hash.ToLower()
+            if ($rustupActual -ne $rustupExpected) {
+              throw "rustup installer hash mismatch: expected $rustupExpected got $rustupActual"
+            }
+            & $rustupInit -y --profile minimal --default-toolchain stable
+            if ($LASTEXITCODE -ne 0) {
+              exit $LASTEXITCODE
+            }
+            Remove-Item $rustupInit
+            $cargoBin = Join-Path $HOME ".cargo\bin"
+            $env:Path = "$cargoBin;$env:Path"
+            rustc --version
+            cargo --version
             $installer = Join-Path $env:TEMP "uv-install.ps1"
             Invoke-WebRequest -Uri https://astral.sh/uv/0.10.9/install.ps1 -OutFile $installer
             $expected = "d43ffff8d28e7d1e7d1831a212465f12b24a43c7f87f386e95e2a5915aee5d7d"
@@ -222,6 +240,9 @@ jobs:
             if (-not (Select-String -Path $PROFILE -SimpleMatch $uvBin -Quiet)) {
               Add-Content -Path $PROFILE -Value "`$env:Path = `"$uvBin;`$env:Path`""
             }
+            if (-not (Select-String -Path $PROFILE -SimpleMatch $cargoBin -Quiet)) {
+              Add-Content -Path $PROFILE -Value "`$env:Path = `"$cargoBin;`$env:Path`""
+            }
             uv sync --frozen --group dev --python 3.11
       - run:
           name: Run Windows-specific test
@@ -232,6 +253,8 @@ jobs:
           environment:
             UV_HTTP_TIMEOUT: "300"
           command: |
+            $env:Path = "$HOME\.cargo\bin;$HOME\.local\bin;$env:Path"
+            cargo --version
             uv build --wheel --out-dir dist
             uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py
 

.dockerignore

@@ -49,6 +49,8 @@ build/
 *.egg-info/
 .DS_Store
 **/node_modules
+litellm-rust/target/
+litellm/rust_bridge/_native*.so
 *.log
 .env
 .env.local

Dockerfile

@@ -21,6 +21,7 @@ RUN apk add --no-cache \
     gcc \
     python3 \
     python3-dev \
+    rust \
     openssl \
     openssl-dev \
     nodejs \

docker/Dockerfile.non_root

@@ -19,6 +19,7 @@ RUN for i in 1 2 3; do \
         python3 \
         python3-dev \
         gcc \
+        rust \
         bash \
         coreutils \
         curl \

litellm-rust/crates/python-bridge/Cargo.toml

@@ -6,7 +6,7 @@ license.workspace = true
 repository.workspace = true
 
 [lib]
-name = "litellm_python_bridge"
+name = "_native"
 crate-type = ["cdylib"]
 
 [dependencies]

litellm-rust/crates/python-bridge/build.rs

@@ -0,0 +1,6 @@
+fn main() {
+    if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("macos") {
+        println!("cargo:rustc-cdylib-link-arg=-undefined");
+        println!("cargo:rustc-cdylib-link-arg=dynamic_lookup");
+    }
+}

litellm-rust/crates/python-bridge/src/lib.rs

@@ -93,7 +93,7 @@ fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
 }
 
 #[pymodule]
-fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> {
+fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add_function(wrap_pyfunction!(ocr, module)?)?;
     module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
     Ok(())

litellm/ocr/rust_bridge.py

@@ -2,7 +2,7 @@
 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``
+then routes supported Mistral calls through the compiled ``litellm.rust_bridge._native``
 extension, which performs the whole OCR call (URL, headers, HTTP, parse) in Rust.
 
 No module-level ``litellm`` imports keep this a leaf so ``litellm/ocr/main.py``
@@ -15,7 +15,7 @@
 
 
 class RustOcr(Protocol):
-    """Signature of the compiled ``litellm_python_bridge.ocr`` entrypoint."""
+    """Signature of the compiled Rust OCR entrypoint."""
 
     def __call__(
         self,
@@ -41,7 +41,7 @@ class _Unset:
 def use_litellm_rust(
     enabled: bool = True, *, ocr: RustOcr | None | _Unset = _UNSET
 ) -> None:
-    """Route supported OCR calls through the Rust ``litellm_python_bridge`` extension.
+    """Route supported OCR calls through the packaged Rust extension.
 
     ``ocr`` injects the bridge callable; when omitted the compiled extension is
     loaded on demand and any previously injected bridge is preserved. Pass
@@ -62,13 +62,14 @@ def load_rust_ocr() -> RustOcr | None:
     """Return the Rust OCR callable, or ``None`` when no bridge is available.
 
     Prefers an injected implementation, otherwise loads the compiled
-    ``litellm_python_bridge`` extension; a missing extension yields ``None`` so
+    ``litellm.rust_bridge._native`` 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
-    try:
-        import litellm_python_bridge
-    except ImportError:
+    from litellm.rust_bridge import get_native_bridge
+
+    native_bridge = get_native_bridge()
+    if native_bridge is None:
         return None
-    return cast(RustOcr, litellm_python_bridge.ocr)
+    return cast(RustOcr, native_bridge.ocr)

litellm/rust_bridge/__init__.py

@@ -0,0 +1,8 @@
+"""LiteLLM Rust bridge package."""
+
+from litellm.rust_bridge.loader import (
+    get_native_bridge,
+    native_bridge_available,
+)
+
+__all__ = ["get_native_bridge", "native_bridge_available"]

litellm/rust_bridge/loader.py

@@ -0,0 +1,28 @@
+"""Loader for the packaged LiteLLM Rust extension."""
+
+from __future__ import annotations
+
+from types import ModuleType
+
+_BRIDGE_SENTINEL = object()
+_cached_bridge: ModuleType | None | object = _BRIDGE_SENTINEL
+
+
+def get_native_bridge() -> ModuleType | None:
+    """Return the packaged Rust extension, or ``None`` when unavailable."""
+    global _cached_bridge
+    if _cached_bridge is not _BRIDGE_SENTINEL:
+        return _cached_bridge if isinstance(_cached_bridge, ModuleType) else None
+
+    try:
+        from litellm.rust_bridge import _native
+    except ImportError:
+        _cached_bridge = None
+        return None
+    _cached_bridge = _native
+    return _native
+
+
+def native_bridge_available() -> bool:
+    """Whether the packaged Rust extension is importable."""
+    return get_native_bridge() is not None

pyproject.toml

@@ -234,8 +234,24 @@ healthcheck = [
 ]
 
 [build-system]
-requires = ["uv_build==0.11.8"]
-build-backend = "uv_build"
+requires = ["maturin>=1.9.4,<2"]
+build-backend = "maturin"
+
+[tool.maturin]
+manifest-path = "litellm-rust/crates/python-bridge/Cargo.toml"
+module-name = "litellm.rust_bridge._native"
+python-source = "."
+bindings = "pyo3"
+exclude = [
+    "litellm/proxy/enterprise",
+    "litellm/proxy/enterprise/**",
+    "**/__pycache__",
+    "**/__pycache__/**",
+    "**/.pytest_cache",
+    "**/.pytest_cache/**",
+    "**/.ruff_cache",
+    "**/.ruff_cache/**",
+]
 
 [tool.uv]
 constraint-dependencies = [
@@ -253,18 +269,6 @@ litellm-enterprise = { workspace = true }
 [tool.uv.workspace]
 members = ["enterprise", "litellm-proxy-extras"]
 
-[tool.uv.build-backend]
-module-root = ""
-source-exclude = [
-    "litellm/proxy/enterprise",
-    "**/__pycache__",
-    "**/__pycache__/**",
-    "**/.pytest_cache",
-    "**/.pytest_cache/**",
-    "**/.ruff_cache",
-    "**/.ruff_cache/**",
-]
-
 [tool.isort]
 profile = "black"
 
@@ -328,4 +332,3 @@ pytest_add_cli_args = [
 [tool.coverage.run]
 source = ["litellm"]
 relative_files = true
-

tests/test_litellm/ocr/test_rust_bridge.py

@@ -1,7 +1,7 @@
 """Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``)."""
 
 import importlib
-import sys
+import builtins
 import types
 
 import httpx
@@ -15,6 +15,7 @@
 # explicitly via importlib rather than attribute traversal.
 ocr_main = importlib.import_module("litellm.ocr.main")
 rust_bridge = importlib.import_module("litellm.ocr.rust_bridge")
+rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader")
 
 MODEL = "mistral/mistral-ocr-latest"
 DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"}
@@ -80,8 +81,10 @@ def get_complete_url(self, *, api_base, model, optional_params, litellm_params):
 def _reset_rust_flag():
     """Keep the global toggle isolated between tests."""
     rust_bridge.use_litellm_rust(False, ocr=None)
+    rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
     yield
     rust_bridge.use_litellm_rust(False, ocr=None)
+    rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
 
 
 @pytest.fixture
@@ -106,6 +109,44 @@ def test_load_rust_ocr_returns_injected_impl():
     assert rust_bridge.load_rust_ocr() is bridge
 
 
+def test_native_bridge_loader_returns_none_when_extension_absent(monkeypatch):
+    real_import = builtins.__import__
+
+    def fake_import(name, globals=None, locals=None, fromlist=(), level=0):
+        if name == "litellm.rust_bridge" and "_native" in fromlist:
+            raise ImportError
+        return real_import(name, globals, locals, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", fake_import)
+
+    assert rust_bridge_loader.get_native_bridge() is None
+
+
+def test_native_bridge_loader_caches_absent_extension(monkeypatch):
+    real_import = builtins.__import__
+    attempts = 0
+
+    def fake_import(name, globals=None, locals=None, fromlist=(), level=0):
+        nonlocal attempts
+        if name == "litellm.rust_bridge" and "_native" in fromlist:
+            attempts += 1
+            raise ImportError
+        return real_import(name, globals, locals, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", fake_import)
+
+    assert rust_bridge_loader.get_native_bridge() is None
+    assert rust_bridge_loader.get_native_bridge() is None
+    assert attempts == 1
+
+
+def test_native_bridge_available_reflects_loader(monkeypatch):
+    fake_module = types.ModuleType("litellm.rust_bridge._native")
+    monkeypatch.setattr(rust_bridge_loader, "get_native_bridge", lambda: fake_module)
+
+    assert rust_bridge_loader.native_bridge_available() is True
+
+
 def test_toggle_without_ocr_arg_preserves_injected_impl():
     """Regression: routine enable/disable calls must not clobber a prior injection.
 
@@ -122,28 +163,42 @@ def test_toggle_without_ocr_arg_preserves_injected_impl():
     assert rust_bridge.load_rust_ocr() is bridge
 
 
-def test_explicit_ocr_none_clears_injected_impl():
+def test_explicit_ocr_none_clears_injected_impl(monkeypatch):
+    monkeypatch.setattr(
+        importlib.import_module("litellm.rust_bridge"),
+        "get_native_bridge",
+        lambda: None,
+    )
     bridge = RecordingBridge()
     litellm.use_litellm_rust(True, ocr=bridge)
 
     litellm.use_litellm_rust(True, ocr=None)
     assert rust_bridge.load_rust_ocr() is None
 
 
-def test_load_rust_ocr_none_when_extension_absent():
+def test_load_rust_ocr_none_when_extension_absent(monkeypatch):
     """With no injected impl and no compiled wheel, the loader returns None so the
     caller degrades to the Python path instead of raising ImportError."""
+    monkeypatch.setattr(
+        importlib.import_module("litellm.rust_bridge"),
+        "get_native_bridge",
+        lambda: None,
+    )
     litellm.use_litellm_rust(True)  # no impl injected; extension isn't built in CI
     assert rust_bridge.load_rust_ocr() is None
 
 
 def test_load_rust_ocr_uses_compiled_extension(monkeypatch):
-    """With no injected impl but a compiled ``litellm_python_bridge`` importable,
+    """With no injected impl but a packaged ``litellm.rust_bridge._native`` importable,
     the loader returns the extension's ``ocr`` callable. The native wheel isn't
-    built in CI, so stand in a fake module via ``sys.modules``."""
-    fake_module = types.ModuleType("litellm_python_bridge")
+    built in CI, so stand in a fake module via the bridge loader."""
+    fake_module = types.ModuleType("litellm.rust_bridge._native")
     fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE)  # type: ignore[attr-defined]
-    monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module)
+    monkeypatch.setattr(
+        importlib.import_module("litellm.rust_bridge"),
+        "get_native_bridge",
+        lambda: fake_module,
+    )
 
     litellm.use_litellm_rust(True)  # enabled, no impl injected -> import the extension
     assert rust_bridge.load_rust_ocr() is fake_module.ocr
@@ -317,6 +372,7 @@ def test_ocr_does_not_route_to_rust_when_disabled():
 def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch):
     """Rust enabled but no bridge available (no injected impl, no compiled wheel):
     ocr() must degrade to the Python HTTP handler instead of raising."""
+    monkeypatch.setattr(ocr_main, "load_rust_ocr", lambda: None)
     litellm.use_litellm_rust(True)  # enabled, but load_rust_ocr() returns None in CI
 
     captured = {}