diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py new file mode 100644 index 00000000000..cae349be998 --- /dev/null +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -0,0 +1,120 @@ +"""Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``).""" + +import importlib +import sys +import types + +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse + +# `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` +# function onto `litellm.ocr` and shadows the submodule — so import the modules +# explicitly via importlib rather than attribute traversal. +ocr_main = importlib.import_module("litellm.ocr.main") +rust_bridge = importlib.import_module("litellm.ocr.rust_bridge") + +MODEL = "mistral/mistral-ocr-latest" +DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} + +FAKE_OCR_RESPONSE = { + "pages": [{"index": 0, "markdown": "hello world"}], + "model": "mistral-ocr-2505-completion", + "document_annotation": None, + "usage_info": {"pages_processed": 1}, + "object": "ocr", +} + + +@pytest.fixture(autouse=True) +def _reset_rust_flag(): + """Keep the global toggle isolated between tests.""" + rust_bridge.use_litellm_rust(False) + yield + rust_bridge.use_litellm_rust(False) + + +@pytest.fixture +def fake_bridge(monkeypatch): + """Install a fake compiled ``litellm_python_bridge`` module and record calls.""" + calls = [] + + def _ocr(model, document, api_key, api_base, optional_params): + calls.append( + { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "optional_params": optional_params, + } + ) + return dict(FAKE_OCR_RESPONSE) + + module = types.ModuleType("litellm_python_bridge") + module.ocr = _ocr # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "litellm_python_bridge", module) + return calls + + +def test_use_litellm_rust_toggles_flag(): + assert rust_bridge.rust_ocr_enabled() is False + litellm.use_litellm_rust() + assert rust_bridge.rust_ocr_enabled() is True + litellm.use_litellm_rust(False) + assert rust_bridge.rust_ocr_enabled() is False + + +def test_rust_ocr_wraps_bridge_response(fake_bridge): + response = rust_bridge.rust_ocr( + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-test", + api_base=None, + optional_params={"include_image_base64": True}, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "hello world" + assert response.model == "mistral-ocr-2505-completion" + assert fake_bridge[0]["model"] == "mistral-ocr-latest" + assert fake_bridge[0]["optional_params"] == {"include_image_base64": True} + + +def test_ocr_routes_to_rust_when_enabled(fake_bridge): + litellm.use_litellm_rust() + + response = litellm.ocr( + model=MODEL, + document=DOCUMENT, + api_key="sk-test", + include_image_base64=True, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "hello world" + assert len(fake_bridge) == 1 + call = fake_bridge[0] + # Provider prefix is stripped before reaching the bridge. + assert call["model"] == "mistral-ocr-latest" + assert call["document"] == DOCUMENT + assert call["api_key"] == "sk-test" + # Raw OCR params ride along in optional_params; Rust filters to supported keys. + assert call["optional_params"].get("include_image_base64") is True + + +def test_ocr_skips_rust_when_disabled(monkeypatch, fake_bridge): + """With the flag off, ocr() must take the normal Python provider path.""" + called = {} + + def _fake_handler(*_args, **_kwargs): + called["hit"] = True + return OCRResponse(pages=[], model="mistral-ocr-latest") + + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", _fake_handler) + + litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert called.get("hit") is True + assert fake_bridge == [] # Rust bridge never invoked