scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

test(ocr): cover sync ocr + async aocr rust bridge routing

Browse source
This commit is contained in:
Ishaan Jaffer 2026-06-23 14:16:20 -07:00
parent 27aa3c10b2
commit 27d6fa9303
No known key found for this signature in database
1 changed files with 178 additions and 178 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

356
tests/test_litellm/ocr/test_rust_bridge.py
View file

@ -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)