feat: package Rust OCR bridge in LiteLLM wheel by ishaan-berri · Pull Request #31265 · BerriAI/litellm

ishaan-berri · 2026-06-25T02:47:33Z

Summary

Synced directly on top of litellm_internal_staging

This makes the default litellm pip package build and ship the Rust OCR bridge inside the existing Python package instead of publishing a separate Rust wheel package

switches the root package build backend to maturin and builds the PyO3 crate as litellm.rust_bridge._native
adds a small litellm.rust_bridge loader package so Python code does not import a top-level native module
updates the OCR Rust opt-in path to load the packaged native bridge and keep Python fallback behavior when unavailable
adds Rust only to Docker builder stages; no runtime Rust dependency and no inline Dockerfile smoke-test bloat
ports the old uv_build source excludes into [tool.maturin], including litellm/proxy/enterprise
ignores local Rust/native build artifacts so Docker context does not pick up generated outputs
enforces typed Rust core translations by failing cargo test if litellm-rust/crates/core/src uses raw serde_json::Value, serde_json::Map, or json!
replaces the remaining raw JSON OCR and realtime core payloads with typed Rust structs/enums, leaving JSON conversion at the gateway/bridge boundary

Related: #31263

Validation

uv lock
uv build --sdist --wheel
artifact inspection: generated wheel/sdist include the native bridge and litellm/rust_bridge/loader.py, and exclude litellm/proxy/enterprise
uv run black litellm/rust_bridge litellm/ocr/rust_bridge.py tests/test_litellm/ocr/test_rust_bridge.py
uv run ruff check litellm/rust_bridge litellm/ocr/rust_bridge.py tests/test_litellm/ocr/test_rust_bridge.py
uv run pytest tests/test_litellm/ocr/test_rust_bridge.py -q (18 passed)
(cd litellm-rust && cargo fmt --check && cargo clippy --workspace --all-targets --locked -- -D warnings && cargo test --workspace --locked)
cargo test --workspace --locked includes typed_core_boundary::core_translation_code_does_not_use_raw_json_values
docker build -t litellm-rust-pip-binary-test .
docker run --rm --entrypoint python litellm-rust-pip-binary-test -c "from litellm.rust_bridge import native_bridge_available; from litellm.ocr.rust_bridge import load_rust_ocr; import litellm.rust_bridge._native as native; assert native_bridge_available(); assert load_rust_ocr() is native.ocr; print('main docker rust bridge ok')"
docker build -f docker/Dockerfile.non_root -t litellm-rust-pip-binary-non-root-test .
docker run --rm --entrypoint python litellm-rust-pip-binary-non-root-test -c "from litellm.rust_bridge import native_bridge_available; from litellm.ocr.rust_bridge import load_rust_ocr; import litellm.rust_bridge._native as native; assert native_bridge_available(); assert load_rust_ocr() is native.ocr; print('non-root docker rust bridge ok')"
git merge-tree $(git merge-base HEAD upstream/litellm_internal_staging) HEAD upstream/litellm_internal_staging showed no conflicts

Type

New Feature
Bug Fix
Test

CLAassistant · 2026-06-25T02:47:40Z

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.}

greptile-apps · 2026-06-25T02:51:07Z

Greptile Summary

This PR packages the Rust OCR bridge directly into the litellm wheel by switching the build backend from uv_build to maturin and exposing the compiled extension as litellm.rust_bridge._native, eliminating the need for a separate Rust wheel package.

Adds a litellm/rust_bridge loader package that wraps the compiled _native extension with graceful None fallback when the native module is absent, and updates litellm/ocr/rust_bridge.py to route through this loader instead of the old top-level litellm_python_bridge import.
Adds rust to Docker builder stages and embeds a build-time smoke test that asserts the native module imports and the OCR bridge rejects a missing API key.
Drops [tool.uv.build-backend] including its source-exclude list; the enterprise directory exclusion is not carried over to the new [tool.maturin] configuration, which risks bundling litellm.proxy.enterprise into the OSS wheel.

Confidence Score: 3/5

The Python and Rust logic changes are correct, but the pyproject.toml change inadvertently removes an explicit exclusion that kept the enterprise directory out of the OSS wheel; this should be resolved before merging.

The new loader, protocol types, and test patches are all well-constructed. The riskiest part of the change is pyproject.toml: the old [tool.uv.build-backend] block explicitly excluded litellm/proxy/enterprise from the built package, and that exclusion was not ported to [tool.maturin]. Because litellm/proxy/enterprise/__init__.py exists and is tracked by git, maturin (with python-source = ".") will include it in the OSS wheel, which is the opposite of the previous behaviour.

pyproject.toml — the missing exclude for litellm/proxy/enterprise in [tool.maturin] is the most important thing to verify before merging.

Important Files Changed

Filename	Overview
pyproject.toml	Switches build backend from uv_build to maturin; removes source-exclude rules including the enterprise directory exclusion, risking bundling enterprise code into the OSS wheel.
litellm/rust_bridge/init.py	New thin package that re-exports `get_native_bridge` and `native_bridge_available` from the loader; clean and minimal.
litellm/rust_bridge/loader.py	New lazy loader for the compiled `_native` extension; correctly returns None on ImportError to preserve Python fallback.
litellm/ocr/rust_bridge.py	Re-routes OCR bridge loading through `litellm.rust_bridge.get_native_bridge()` instead of the former top-level `litellm_python_bridge` import; fallback behaviour unchanged.
Dockerfile	Adds `rust` to Alpine builder stage and embeds a Python smoke test for the native bridge; test is duplicated verbatim in the non-root Dockerfile.
docker/Dockerfile.non_root	Same Rust toolchain addition and identical smoke-test block as main Dockerfile; duplication is the only concern.
tests/test_litellm/ocr/test_rust_bridge.py	Test improvements: replaces brittle `sys.modules` injection with `monkeypatch.setattr` on `get_native_bridge`; adds missing `monkeypatch` for `load_rust_ocr` in the fallback test; no coverage regressions detected.
litellm-rust/crates/python-bridge/Cargo.toml	Renames the cdylib from `litellm_python_bridge` to `_native` to match the new `litellm.rust_bridge._native` module path.
litellm-rust/crates/python-bridge/src/lib.rs	Updates the `#[pymodule]` initializer name to `_native`; trivial rename, no logic change.
.dockerignore	Adds `litellm-rust/target/` to prevent Rust build artifacts from entering the Docker build context; correct and necessary.

_{Reviews (1): Last reviewed commit: "feat: package rust ocr bridge in litellm..." | Re-trigger Greptile}

greptile-apps · 2026-06-25T02:51:10Z

+[tool.maturin]
+manifest-path = "litellm-rust/crates/python-bridge/Cargo.toml"
+module-name = "litellm.rust_bridge._native"
+python-source = "."
+bindings = "pyo3"


Enterprise directory will be bundled into the OSS wheel

The old [tool.uv.build-backend] explicitly excluded litellm/proxy/enterprise from the built package. That block was removed when switching to maturin, but no equivalent exclude was added to [tool.maturin]. litellm/proxy/enterprise/ has an __init__.py and is tracked by git, so maturin (with python-source = ".") will include litellm.proxy.enterprise as a subpackage of the OSS wheel — the inverse of the old behaviour. Add an explicit exclude to [tool.maturin]: exclude = ["litellm/proxy/enterprise/**"]

greptile-apps · 2026-06-25T02:51:11Z

+RUN python - <<'PY'
+import importlib
+import os
+
+from litellm.ocr.rust_bridge import load_rust_ocr
+from litellm.rust_bridge import native_bridge_available
+
+os.environ.pop("MISTRAL_API_KEY", None)
+
+native = importlib.import_module("litellm.rust_bridge._native")
+assert native_bridge_available()
+assert load_rust_ocr() is native.ocr
+
+try:
+    native.ocr(
+        model="mistral-ocr-latest",
+        document={"type": "document_url", "document_url": "https://example.com/doc.pdf"},
+        api_key=None,
+        api_base=None,
+        custom_llm_provider="mistral",
+        extra_headers={},
+        optional_params={},
+        timeout_seconds=1.0,
+    )
+except ValueError as exc:
+    assert "Missing Mistral API Key" in str(exc)
+else:
+    raise AssertionError("expected Rust OCR bridge to reject missing Mistral API key")
+PY
+
 RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
    find /app/.venv -type d -path "*/tornado/test" -delete


Identical smoke-test script copy-pasted into both Dockerfiles

The same 30-line Python block appears verbatim in Dockerfile and docker/Dockerfile.non_root. The two copies will drift if the Rust error message, function signature, or import path ever changes — a single edit site would be missed. Consider extracting the script to docker/smoke_test_rust_bridge.py and referencing it with COPY + RUN python docker/smoke_test_rust_bridge.py (or equivalent) in both files.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

codecov · 2026-06-25T02:51:44Z

Codecov Report

❌ Patch coverage is 93.75000% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/rust_bridge/loader.py	90.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

ishaan-berri · 2026-06-25T03:04:13Z

Addressed Greptile feedback and reduced PR scope: retargeted/rebased onto litellm_internal_staging, ported the old enterprise/source excludes into [tool.maturin], removed the inline Docker smoke-test blocks, and kept Dockerfile changes to builder Rust only. Rebuilt both Docker images and verified native bridge imports from both containers.

ishaan-berri · 2026-06-25T03:05:05Z

Superseded by #31267. Recreated from BerriAI/litellm:litellm_hotfix_rust_pip_binary so the source-branch guard passes.

ishaan-berri mentioned this pull request Jun 25, 2026

LiteLLM Rust Migration - the fastest and litest AI Gateway (sub 1ms overheads) #31263

Open

greptile-apps Bot reviewed Jun 25, 2026

View reviewed changes

feat: package rust ocr bridge in litellm wheel

1ffe1d8

ishaan-berri force-pushed the codex/litellm-rust-pip-binary branch from 3aea24e to 1ffe1d8 Compare June 25, 2026 03:03

ishaan-berri changed the base branch from main to litellm_internal_staging June 25, 2026 03:03

ishaan-berri requested a review from a team June 25, 2026 03:03

ishaan-berri mentioned this pull request Jun 25, 2026

feat: package Rust OCR bridge in LiteLLM wheel #31267

Merged

ishaan-berri closed this Jun 25, 2026

Navigation Menu

Uh oh!

Uh oh!

feat: package Rust OCR bridge in LiteLLM wheel - #31265

feat: package Rust OCR bridge in LiteLLM wheel#31265
ishaan-berri wants to merge 1 commit into
BerriAI:litellm_internal_stagingfrom
ishaan-berri:codex/litellm-rust-pip-binary

ishaan-berri commented Jun 25, 2026 •

edited

Loading

Uh oh!

CLAassistant commented Jun 25, 2026

Uh oh!

greptile-apps Bot commented Jun 25, 2026

Important Files Changed

Uh oh!

greptile-apps Bot Jun 25, 2026

Uh oh!

greptile-apps Bot Jun 25, 2026

Uh oh!

codecov Bot commented Jun 25, 2026 •

edited

Loading

Uh oh!

ishaan-berri commented Jun 25, 2026

Uh oh!

ishaan-berri commented Jun 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Uh oh!

Conversation

ishaan-berri commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Type

Uh oh!

CLAassistant commented Jun 25, 2026

Uh oh!

greptile-apps Bot commented Jun 25, 2026

Greptile Summary

Confidence Score: 3/5

Important Files Changed

Uh oh!

greptile-apps Bot Jun 25, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jun 25, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ishaan-berri commented Jun 25, 2026

Uh oh!

ishaan-berri commented Jun 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ishaan-berri commented Jun 25, 2026 •

edited

Loading

codecov Bot commented Jun 25, 2026 •

edited

Loading