diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 7e028064e4c..0db00366af4 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,10 +1,12 @@ -"""Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``).""" +"""Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``). +Covers both the sync ``ocr`` and async ``aocr`` bridge entry points with an +injected fake bridge — the compiled extension is not built in CI. +""" + +import asyncio import importlib -import sys -import types -import httpx import pytest import litellm @@ -29,25 +31,65 @@ FAKE_OCR_RESPONSE = { class RecordingBridge: - """A fake ``RustOcr`` callable that records the args it was handed.""" + """A fake bridge exposing sync ``ocr`` and async ``aocr``, recording calls.""" def __init__(self): - self.calls = [] + self.ocr_calls = [] + self.aocr_calls = [] - def __call__( - self, model, document, api_key, api_base, optional_params, timeout_seconds + def _record(self, store, kwargs): + store.append(kwargs) + return dict(FAKE_OCR_RESPONSE) + + def ocr( + self, + provider, + model, + document, + api_key=None, + api_base=None, + extra_headers=None, + timeout_seconds=None, + params=None, ): - self.calls.append( + return self._record( + self.ocr_calls, { + "provider": provider, "model": model, "document": document, "api_key": api_key, "api_base": api_base, - "optional_params": optional_params, + "extra_headers": extra_headers, "timeout_seconds": timeout_seconds, - } + "params": params, + }, + ) + + async def aocr( + self, + provider, + model, + document, + api_key=None, + api_base=None, + extra_headers=None, + timeout_seconds=None, + params=None, + ): + return self._record( + self.aocr_calls, + { + "provider": provider, + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "extra_headers": extra_headers, + "timeout_seconds": timeout_seconds, + "params": params, + }, ) - return dict(FAKE_OCR_RESPONSE) class RecordingLogging: @@ -70,7 +112,7 @@ class FakeOCRConfig: def validate_environment( self, *, headers, model, api_key, api_base, litellm_params ): - return {"authorization": f"Bearer {api_key}"} + return {**headers, "authorization": f"Bearer {api_key}"} def get_complete_url(self, *, api_base, model, optional_params, litellm_params): return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" @@ -79,19 +121,22 @@ class FakeOCRConfig: @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge.use_litellm_rust(False, ocr=None) + rust_bridge.use_litellm_rust(False, bridge=None) yield - rust_bridge.use_litellm_rust(False, ocr=None) + rust_bridge.use_litellm_rust(False, bridge=None) @pytest.fixture def fake_bridge(): """Enable the Rust path with an injected recording bridge (no native wheel).""" bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.use_litellm_rust(True, bridge=bridge) return bridge +# --------------------------------------------------------------------------- # +# Toggle + bridge loading +# --------------------------------------------------------------------------- # def test_use_litellm_rust_toggles_flag(): assert rust_bridge.rust_ocr_enabled() is False litellm.use_litellm_rust() @@ -100,224 +145,180 @@ def test_use_litellm_rust_toggles_flag(): assert rust_bridge.rust_ocr_enabled() is False -def test_load_rust_ocr_returns_injected_impl(): +def test_load_rust_bridge_returns_injected_bridge(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) - assert rust_bridge.load_rust_ocr() is bridge + litellm.use_litellm_rust(True, bridge=bridge) + assert rust_bridge.load_rust_bridge() is bridge -def test_toggle_without_ocr_arg_preserves_injected_impl(): - """Regression: routine enable/disable calls must not clobber a prior injection. - - Earlier, ``use_litellm_rust()`` unconditionally assigned the keyword default - of ``None`` to ``_rust_ocr_impl``, silently dropping a custom bridge whenever - a caller toggled the flag without re-passing ``ocr=``. - """ +def test_toggle_without_bridge_arg_preserves_injection(): + """Routine enable/disable must not clobber a prior injection.""" bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.use_litellm_rust(True, bridge=bridge) litellm.use_litellm_rust(False) - assert rust_bridge.load_rust_ocr() is bridge + assert rust_bridge.load_rust_bridge() is bridge litellm.use_litellm_rust(True) - assert rust_bridge.load_rust_ocr() is bridge + assert rust_bridge.load_rust_bridge() is bridge -def test_explicit_ocr_none_clears_injected_impl(): +def test_explicit_bridge_none_clears_injection(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.use_litellm_rust(True, bridge=bridge) - litellm.use_litellm_rust(True, ocr=None) - assert rust_bridge.load_rust_ocr() is None + litellm.use_litellm_rust(True, bridge=None) + assert rust_bridge.load_rust_bridge() is None -def test_load_rust_ocr_none_when_extension_absent(): - """With no injected impl and no compiled wheel, the loader returns None so the - caller degrades to the Python path instead of raising ImportError.""" - 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_bridge_none_when_extension_absent(): + litellm.use_litellm_rust(True) # no injection; extension isn't built in CI + assert rust_bridge.load_rust_bridge() is None -def test_load_rust_ocr_uses_compiled_extension(monkeypatch): - """With no injected impl but a compiled ``litellm_python_bridge`` 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") - fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] - monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module) - - litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension - assert rust_bridge.load_rust_ocr() is fake_module.ocr +def test_rust_supports_only_known_providers(): + assert rust_bridge.rust_supports("mistral") is True + assert rust_bridge.rust_supports("azure_ai") is False + assert rust_bridge.rust_supports("openai") is False +# --------------------------------------------------------------------------- # +# Helper: timeout + sync runner +# --------------------------------------------------------------------------- # def test_timeout_to_seconds_handles_float_timeout_and_none(): + import httpx + assert ocr_main._timeout_to_seconds(12.5) == 12.5 assert ocr_main._timeout_to_seconds(None) is None assert ocr_main._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 -def test_run_rust_ocr_forwards_args_and_wraps_response(): - bridge = RecordingBridge() - logging_obj = RecordingLogging() - - response = ocr_main._run_rust_ocr( - rust_ocr=bridge, +def _rust_call(logging_obj, **overrides): + base = dict( logging_obj=logging_obj, provider_config=FakeOCRConfig(), resolve_api_key=lambda _name: None, + provider="mistral", model="mistral-ocr-latest", document=DOCUMENT, api_key="sk-test", api_base="https://proxy.internal", + extra_headers=None, optional_params={"include_image_base64": True}, litellm_params={}, timeout_seconds=12.5, ) + base.update(overrides) + return ocr_main._RustOcrCall(**base) + + +def test_run_rust_ocr_forwards_args_and_wraps_response(): + bridge = RecordingBridge() + response = ocr_main._run_rust_ocr(bridge, _rust_call(RecordingLogging())) assert isinstance(response, OCRResponse) assert response.pages[0].markdown == "hello world" - call = bridge.calls[0] - assert call == { - "model": "mistral-ocr-latest", - "document": DOCUMENT, - "api_key": "sk-test", - "api_base": "https://proxy.internal", - "optional_params": {"include_image_base64": True}, - "timeout_seconds": 12.5, - } - - -def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): - """No explicit api_key: the resolver (get_secret_str in production) supplies it, - so secret-manager backends (AWS/Azure/GCP/Vault) work like the Python path.""" - bridge = RecordingBridge() - - ocr_main._run_rust_ocr( - rust_ocr=bridge, - logging_obj=RecordingLogging(), - provider_config=FakeOCRConfig(), - resolve_api_key=lambda name: ( - "sk-from-vault" if name == "MISTRAL_API_KEY" else None - ), - model="mistral-ocr-latest", - document=DOCUMENT, - api_key=None, - api_base=None, - optional_params={}, - litellm_params={}, - timeout_seconds=None, - ) - - assert bridge.calls[0]["api_key"] == "sk-from-vault" - - -def test_run_rust_ocr_prefers_explicit_key_over_resolver(): - bridge = RecordingBridge() - resolver_calls = [] - - def _resolver(name): - resolver_calls.append(name) - return "sk-from-vault" - - ocr_main._run_rust_ocr( - rust_ocr=bridge, - logging_obj=RecordingLogging(), - provider_config=FakeOCRConfig(), - resolve_api_key=_resolver, - model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-explicit", - api_base=None, - optional_params={}, - litellm_params={}, - timeout_seconds=None, - ) - - assert bridge.calls[0]["api_key"] == "sk-explicit" - assert resolver_calls == [] # resolver never consulted when a key is supplied - - -def test_run_rust_ocr_runs_pre_call_logging(): - """The Rust shortcut must run pre_call so callbacks and spend tracking fire.""" - logging_obj = RecordingLogging() - - ocr_main._run_rust_ocr( - rust_ocr=RecordingBridge(), - logging_obj=logging_obj, - provider_config=FakeOCRConfig(), - resolve_api_key=lambda _name: None, - model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-test", - api_base="https://api.mistral.ai/v1", - optional_params={"include_image_base64": True}, - litellm_params={}, - timeout_seconds=None, - ) - - assert logging_obj.pre_call_kwargs is not None - assert logging_obj.pre_call_kwargs["input"] == "OCR document processing" - additional_args = logging_obj.pre_call_kwargs["additional_args"] - complete_input = additional_args["complete_input_dict"] - assert complete_input["document"] == DOCUMENT - assert complete_input["include_image_base64"] is True - # The logged request mirrors what Rust sends: resolved URL + headers. - assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" - assert additional_args["headers"] == {"authorization": "Bearer sk-test"} - - -def test_ocr_routes_to_rust_when_enabled(fake_bridge): - 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.calls) == 1 - call = fake_bridge.calls[0] - # Provider prefix is stripped before reaching the bridge. + call = bridge.ocr_calls[0] + assert call["provider"] == "mistral" 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 + assert call["timeout_seconds"] == 12.5 + assert call["params"] == {"include_image_base64": True} + + +def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): + """No explicit api_key: the resolver (get_secret_str in prod) supplies it.""" + bridge = RecordingBridge() + ocr_main._run_rust_ocr( + bridge, + _rust_call( + RecordingLogging(), + api_key=None, + resolve_api_key=lambda name: ( + "sk-from-vault" if name == "MISTRAL_API_KEY" else None + ), + ), + ) + assert bridge.ocr_calls[0]["api_key"] == "sk-from-vault" + + +def test_run_rust_ocr_runs_pre_call_logging(): + logging_obj = RecordingLogging() + ocr_main._run_rust_ocr( + RecordingBridge(), + _rust_call(logging_obj, api_base="https://api.mistral.ai/v1"), + ) + + assert logging_obj.pre_call_kwargs is not None + additional = logging_obj.pre_call_kwargs["additional_args"] + assert additional["complete_input_dict"]["document"] == DOCUMENT + assert additional["api_base"] == "https://api.mistral.ai/v1/ocr" + assert additional["headers"] == {"authorization": "Bearer sk-test"} + + +def test_arun_rust_ocr_awaits_bridge_aocr(): + bridge = RecordingBridge() + response = asyncio.run( + ocr_main._arun_rust_ocr(bridge, _rust_call(RecordingLogging())) + ) + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "hello world" + # The async entry point was used, not the sync one. + assert len(bridge.aocr_calls) == 1 + assert len(bridge.ocr_calls) == 0 + assert bridge.aocr_calls[0]["params"] == {"include_image_base64": True} + + +# --------------------------------------------------------------------------- # +# Public ocr() / aocr() routing +# --------------------------------------------------------------------------- # +def test_ocr_routes_to_rust_sync(fake_bridge): + 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.ocr_calls) == 1 + call = fake_bridge.ocr_calls[0] + assert call["provider"] == "mistral" + assert call["model"] == "mistral-ocr-latest" # provider prefix stripped + assert call["params"].get("include_image_base64") is True + + +def test_aocr_routes_to_rust_async(fake_bridge): + response = asyncio.run( + litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + ) + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "hello world" + # Went through the async bridge (no executor thread held during HTTP). + assert len(fake_bridge.aocr_calls) == 1 + assert len(fake_bridge.ocr_calls) == 0 + assert fake_bridge.aocr_calls[0]["model"] == "mistral-ocr-latest" def test_ocr_forwards_timeout_to_rust(fake_bridge): - """Caller-supplied timeout must flow into the Rust bridge so the fixed 600s - client ceiling doesn't silently override shorter deadlines.""" litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test", timeout=12.5) - - assert fake_bridge.calls[0]["timeout_seconds"] == 12.5 + assert fake_bridge.ocr_calls[0]["timeout_seconds"] == 12.5 def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): - """When no explicit timeout is given, the library default (request_timeout) - must still be forwarded so the Rust path matches the Python path's deadline.""" from litellm.constants import request_timeout litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) + assert fake_bridge.ocr_calls[0]["timeout_seconds"] == float(request_timeout) def test_ocr_does_not_route_to_rust_when_disabled(): - """With the flag off, the bridge must not be consulted even if an impl exists.""" bridge = RecordingBridge() - litellm.use_litellm_rust(False, ocr=bridge) - + litellm.use_litellm_rust(False, bridge=bridge) assert rust_bridge.rust_ocr_enabled() is False - # The impl stays available for injection, but the disabled flag gates usage, - # so ocr() never reaches the Rust path (asserted via the enabled-path test). - assert bridge.calls == [] + assert bridge.ocr_calls == [] 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.""" - litellm.use_litellm_rust(True) # enabled, but load_rust_ocr() returns None in CI + """Rust enabled but no bridge: ocr() must degrade to the Python handler.""" + litellm.use_litellm_rust(True) # enabled, but load_rust_bridge() returns None captured = {} @@ -326,8 +327,7 @@ def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): return OCRResponse(pages=[], model="mistral-ocr-latest", object="ocr") monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fake_handler_ocr) - response = litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - assert captured.get("called") is True # Python path was used + assert captured.get("called") is True assert isinstance(response, OCRResponse)