diff --git a/.circleci/config.yml b/.circleci/config.yml index abcdbf451879..d532d3106e5f 100644 --- a/.circleci/config.yml +++ b/.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 diff --git a/.dockerignore b/.dockerignore index a487d2a859ab..6b80caeaf9f2 100644 --- a/.dockerignore +++ b/.dockerignore @@ -49,6 +49,8 @@ build/ *.egg-info/ .DS_Store **/node_modules +litellm-rust/target/ +litellm/rust_bridge/_native*.so *.log .env .env.local diff --git a/Dockerfile b/Dockerfile index af49dc8d8cfe..681681f28cc2 100644 --- a/Dockerfile +++ b/Dockerfile @@ -21,6 +21,7 @@ RUN apk add --no-cache \ gcc \ python3 \ python3-dev \ + rust \ openssl \ openssl-dev \ nodejs \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index ab02b43d0f95..6bb925aa7238 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -19,6 +19,7 @@ RUN for i in 1 2 3; do \ python3 \ python3-dev \ gcc \ + rust \ bash \ coreutils \ curl \ diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index f5b29f49cfd5..4db328186043 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/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] diff --git a/litellm-rust/crates/python-bridge/build.rs b/litellm-rust/crates/python-bridge/build.rs new file mode 100644 index 000000000000..0f7293007b23 --- /dev/null +++ b/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"); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 50ec7fceeaca..8c8416b6bd3b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -93,7 +93,7 @@ fn gil_stats(py: Python<'_>) -> PyResult> { } #[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(()) diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 61f9e8ca69ae..4e688c42e1f4 100644 --- a/litellm/ocr/rust_bridge.py +++ b/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) diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py new file mode 100644 index 000000000000..ec89e3b65b4c --- /dev/null +++ b/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"] diff --git a/litellm/rust_bridge/loader.py b/litellm/rust_bridge/loader.py new file mode 100644 index 000000000000..3ae0bf9307c6 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index ae612ae2184c..30b9a84e09a1 100644 --- a/pyproject.toml +++ b/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 - diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 7e028064e4cc..11448aed828b 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/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,7 +163,12 @@ 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) @@ -130,20 +176,29 @@ def test_explicit_ocr_none_clears_injected_impl(): 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 = {} diff --git a/uv.lock b/uv.lock index 8193bf14f63b..d81d6e5d8a0a 100644 --- a/uv.lock +++ b/uv.lock @@ -3121,15 +3121,15 @@ wheels = [ [[package]] name = "langgraph-checkpoint" -version = "4.1.0" +version = "4.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/02/b4/6005c5dd88ad484fe6235d4c43a0d2cee7e91b08ad85a180985c2662df87/langgraph_checkpoint-4.1.0.tar.gz", hash = "sha256:e5bb304e30fc1363ac8fcb5f7dee5ca2185d77fe475b0d01de2c5f91324c2c21", size = 181942, upload-time = "2026-05-12T03:33:49.888Z" } +sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/93/74/d3be2b41955e20ccd624dba5f6fe9d38dcee385ba470a6e13ed86732fc86/langgraph_checkpoint-4.1.0-py3-none-any.whl", hash = "sha256:8bc2a0466a20c38b865ce6671b42093fd5c041133f32351cae4222e0eeaf7fb5", size = 56047, upload-time = "2026-05-12T03:33:48.548Z" }, + { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, ] [[package]]