diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0ed951b789f..2f7f22032be 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 13428 + "limit": 13427 }, "reportArgumentType": { - "limit": 2196 + "limit": 2194 }, "reportAssignmentType": { "limit": 319 @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44247 + "limit": 44136 }, "reportUnknownLambdaType": { "limit": 109 diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 56c2292d00c..36e3b8838e2 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -29,7 +29,7 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge -from litellm.rust_bridge.configuration import rust_enabled +from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -284,51 +284,57 @@ def _prepare_rust_ocr_call( def _run_rust_ocr( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], -) -> OCRResponse | None: - if rust_ocr_bridge.load_rust_ocr() is None: - return None - prepared: Final = _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ) - rust_response: Final = rust_ocr_bridge.ocr( + fallback: Callable[[], OCRResponse | Coroutine[object, object, OCRResponse]], +) -> OCRResponse | Coroutine[object, object, OCRResponse]: + return rust_ocr_bridge.dispatch_ocr( + prepare=lambda: _prepare_rust_ocr_call( + prepared_request=prepared_request, + resolve_api_key=resolve_api_key, + ), + call=lambda native, prepared: native( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), + ), + fallback=fallback, + adapt=OCRResponse.model_validate, model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, + provider=prepared_request.custom_llm_provider, + eligible=_rust_ocr_supported(prepared_request), ) - if rust_response is None: - return None - return OCRResponse.model_validate(rust_response) async def _run_rust_aocr( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], -) -> OCRResponse | None: - if rust_ocr_bridge.load_rust_aocr() is None: - return None - prepared: Final = _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ) - rust_response: Final = await rust_ocr_bridge.aocr( + fallback: Callable[[], Coroutine[object, object, OCRResponse]], +) -> OCRResponse: + return await rust_ocr_bridge.adispatch_ocr( + prepare=lambda: _prepare_rust_ocr_call( + prepared_request=prepared_request, + resolve_api_key=resolve_api_key, + ), + call=lambda native, prepared: native( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), + ), + fallback=fallback, + adapt=OCRResponse.model_validate, model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, + provider=prepared_request.custom_llm_provider, + eligible=_rust_ocr_supported(prepared_request), ) - if rust_response is None: - return None - return OCRResponse.model_validate(rust_response) @client @@ -425,40 +431,33 @@ async def aocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and rust_enabled(): - from litellm.secret_managers.main import get_secret_str + from litellm.secret_managers.main import get_secret_str - rust_response: Final = await _run_rust_aocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, + async def python_fallback() -> OCRResponse: + pending: Final = base_llm_http_handler.ocr( + model=prepared.model, + document=prepared.document, + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + aocr=True, + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, ) - if rust_response is None: - verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response + response: Final = await pending if asyncio.iscoroutine(pending) else pending + if response is None: + raise ValueError(f"Got an unexpected None response from the OCR API: {response}") + return response - response = base_llm_http_handler.ocr( - model=prepared.model, - document=prepared.document, - optional_params=prepared.optional_params, - timeout=prepared.effective_timeout, - logging_obj=prepared.litellm_logging_obj, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared.custom_llm_provider, - aocr=True, - headers=prepared.extra_headers, - provider_config=prepared.provider_config, - litellm_params=prepared.litellm_params, + return await _run_rust_aocr( + prepared_request=prepared, + resolve_api_key=get_secret_str, + fallback=python_fallback, ) - - if asyncio.iscoroutine(response): - response = await response - - if response is None: - raise ValueError(f"Got an unexpected None response from the OCR API: {response}") - - return response except Exception as e: raise litellm.exception_type( model=model, @@ -697,34 +696,29 @@ def ocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and rust_enabled(): - from litellm.secret_managers.main import get_secret_str + from litellm.secret_managers.main import get_secret_str - rust_response: Final = _run_rust_ocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, + def python_fallback() -> OCRResponse | Coroutine[object, object, OCRResponse]: + return base_llm_http_handler.ocr( + model=prepared.model, + document=prepared.document, + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + aocr=_is_async, + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, ) - if rust_response is None: - verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response - response: Final = base_llm_http_handler.ocr( - model=prepared.model, - document=prepared.document, - optional_params=prepared.optional_params, - timeout=prepared.effective_timeout, - logging_obj=prepared.litellm_logging_obj, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared.custom_llm_provider, - aocr=_is_async, - headers=prepared.extra_headers, - provider_config=prepared.provider_config, - litellm_params=prepared.litellm_params, + return _run_rust_ocr( + prepared_request=prepared, + resolve_api_key=get_secret_str, + fallback=python_fallback, ) - - return response except Exception as e: raise litellm.exception_type( model=model, diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 734e48fd3ed..169966da8be 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,30 +2,50 @@ from __future__ import annotations -from typing import Final +from collections.abc import Awaitable, Callable, Mapping +from typing import Final, TypeVar -import httpx - -from litellm.rust_bridge import configuration as _configuration -from litellm.rust_bridge.protocols import RustAocr, RustOcr -from litellm.rust_bridge.runtime import ( +from . import configuration as _configuration +from .bindings import UNCHANGED, Unchanged +from .protocols import RustAocr, RustOcr +from .runtime import ( BridgeErrorContext, EndpointDispatch, NativeErrorPolicy, - async_none, - identity, ) -from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds + +rust_ocr_enabled = _configuration.rust_ocr_enabled +rust = _configuration.rust +ResultT = TypeVar("ResultT") +RequestT = TypeVar("RequestT") + _OCR: Final[EndpointDispatch[RustOcr, RustAocr]] = EndpointDispatch.native( route="ocr", sync=lambda native: native.ocr, asynchronous=lambda native: native.aocr, - enabled=_configuration.rust_enabled, + enabled=_configuration.rust_ocr_enabled, error_policy=NativeErrorPolicy.PROPAGATE, ) +def set_rust_ocr( + *, + ocr: RustOcr | None | Unchanged = UNCHANGED, + aocr: RustAocr | None | Unchanged = UNCHANGED, +) -> None: + if not isinstance(ocr, Unchanged): + if ocr is None: + _OCR.sync.reset() + else: + _OCR.sync.override(ocr) + if not isinstance(aocr, Unchanged): + if aocr is None: + _OCR.asynchronous.reset() + else: + _OCR.asynchronous.override(aocr) + + def load_rust_ocr() -> RustOcr | None: return _OCR.sync.load() @@ -34,59 +54,41 @@ def load_rust_aocr() -> RustAocr | None: return _OCR.asynchronous.load() -def ocr( +def dispatch_ocr( *, + prepare: Callable[[], RequestT], + call: Callable[[RustOcr, RequestT], Mapping[str, object]], + fallback: Callable[[], ResultT], + adapt: Callable[[Mapping[str, object]], ResultT], model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: + provider: str, + eligible: bool, +) -> ResultT: return _OCR.invoke( - prepare=lambda: _timeout_to_seconds(timeout), - call=lambda rust_ocr, timeout_seconds: rust_ocr( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=timeout_seconds, - ), - fallback=lambda: None, - adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + prepare=prepare, + call=call, + fallback=fallback, + adapt=adapt, + error_context=BridgeErrorContext(provider=provider, model=model), + eligible=eligible, ) -async def aocr( +async def adispatch_ocr( *, + prepare: Callable[[], RequestT], + call: Callable[[RustAocr, RequestT], Awaitable[Mapping[str, object]]], + fallback: Callable[[], Awaitable[ResultT]], + adapt: Callable[[Mapping[str, object]], ResultT], model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: + provider: str, + eligible: bool, +) -> ResultT: return await _OCR.ainvoke( - prepare=lambda: _timeout_to_seconds(timeout), - call=lambda rust_aocr, timeout_seconds: rust_aocr( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=timeout_seconds, - ), - fallback=async_none, - adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + prepare=prepare, + call=call, + fallback=fallback, + adapt=adapt, + error_context=BridgeErrorContext(provider=provider, model=model), + eligible=eligible, ) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 55699214431..7a9abf1db97 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -156,7 +156,7 @@ "limit": 215 }, "PLW0603": { - "limit": 188 + "limit": 186 }, "PLW1508": { "limit": 190 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1033 + "limit": 1031 }, "TRY002": { "limit": 524 @@ -246,7 +246,7 @@ "limit": 109 }, "TRY300": { - "limit": 852 + "limit": 848 }, "UP028": { "limit": 2 diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 28e7acb8a65..6ad08170e35 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -3,13 +3,15 @@ import builtins import importlib import types -from typing import Any +from typing import Any, Final +from unittest.mock import AsyncMock, Mock import httpx import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.rust_bridge import configuration # `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` @@ -17,7 +19,6 @@ from litellm.rust_bridge import configuration # explicitly via importlib rather than attribute traversal. ocr_main = importlib.import_module("litellm.ocr.main") rust_bridge = importlib.import_module("litellm.rust_bridge.ocr") -rust_bridge_bindings = importlib.import_module("litellm.rust_bridge.bindings") rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") MODEL = "mistral/mistral-ocr-latest" @@ -216,13 +217,11 @@ def build_prepared_request( @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge._OCR.sync.reset() - rust_bridge._OCR.asynchronous.reset() + rust_bridge.set_rust_ocr(ocr=None, aocr=None) configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - rust_bridge._OCR.sync.reset() - rust_bridge._OCR.asynchronous.reset() + rust_bridge.set_rust_ocr(ocr=None, aocr=None) configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -232,7 +231,7 @@ def fake_bridge(): """Enable the Rust path with an injected recording bridge (no native wheel).""" bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) return bridge @@ -241,14 +240,27 @@ def fake_async_bridge(): """Enable the async Rust path with an injected recording bridge.""" bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge._OCR.asynchronous.override(bridge) + rust_bridge.set_rust_ocr(aocr=bridge) return bridge +def test_rust_toggles_flag(): + assert rust_bridge.rust_ocr_enabled() is False + litellm.rust(True) + assert rust_bridge.rust_ocr_enabled() is True + litellm.rust(False) + assert rust_bridge.rust_ocr_enabled() is False + + +def test_env_var_enables_rust_ocr(monkeypatch): + monkeypatch.setenv("LITELLM_RUST", "1") + assert rust_bridge.rust_ocr_enabled() is True + + def test_load_rust_ocr_returns_injected_impl(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) assert rust_bridge.load_rust_ocr() is bridge @@ -312,7 +324,7 @@ def test_native_bridge_available_reflects_loader(monkeypatch): def test_load_rust_aocr_returns_injected_impl(): bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge._OCR.asynchronous.override(bridge) + rust_bridge.set_rust_ocr(aocr=bridge) assert rust_bridge.load_rust_aocr() is bridge @@ -321,8 +333,7 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) - rust_bridge._OCR.asynchronous.override(async_bridge) + rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) litellm.rust(False) assert rust_bridge.load_rust_ocr() is bridge @@ -334,18 +345,16 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): def test_explicit_ocr_none_clears_injected_impl(monkeypatch): monkeypatch.setattr( - rust_bridge_bindings, + importlib.import_module("litellm.rust_bridge.bindings"), "get_native_bridge", lambda: None, ) bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) - rust_bridge._OCR.asynchronous.override(async_bridge) + rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) - rust_bridge._OCR.sync.override(None) - rust_bridge._OCR.asynchronous.override(None) + rust_bridge.set_rust_ocr(ocr=None, aocr=None) assert rust_bridge.load_rust_ocr() is None assert rust_bridge.load_rust_aocr() is None @@ -354,7 +363,7 @@ 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( - rust_bridge_bindings, + importlib.import_module("litellm.rust_bridge.bindings"), "get_native_bridge", lambda: None, ) @@ -371,7 +380,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] monkeypatch.setattr( - rust_bridge_bindings, + importlib.import_module("litellm.rust_bridge.bindings"), "get_native_bridge", lambda: fake_module, ) @@ -382,9 +391,9 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): def test_timeout_to_seconds_handles_float_timeout_and_none(): - assert rust_bridge._timeout_to_seconds(12.5) == 12.5 - assert rust_bridge._timeout_to_seconds(None) is None - assert rust_bridge._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 + assert timeout_to_seconds(12.5) == 12.5 + assert timeout_to_seconds(None) is None + assert timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): @@ -392,16 +401,24 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) - response = rust_bridge.ocr( + rust_bridge.set_rust_ocr(ocr=bridge) + response = rust_bridge.dispatch_ocr( + prepare=lambda: 12.5, + call=lambda native, timeout: native( + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-test", + api_base="https://proxy.internal", + custom_llm_provider="mistral", + extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, + optional_params={"include_image_base64": True, "pages": [0]}, + timeout_seconds=timeout, + ), + fallback=lambda: pytest.fail("unexpected Python fallback"), + adapt=dict, model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-test", - api_base="https://proxy.internal", - custom_llm_provider="mistral", - extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True, "pages": [0]}, - timeout=12.5, + provider="mistral", + eligible=True, ) assert response == FAKE_OCR_RESPONSE @@ -427,16 +444,28 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): litellm.rust(True) - rust_bridge._OCR.asynchronous.override(bridge) - response = await rust_bridge.aocr( + rust_bridge.set_rust_ocr(aocr=bridge) + + async def unexpected_fallback(): + pytest.fail("unexpected Python fallback") + + response = await rust_bridge.adispatch_ocr( + prepare=lambda: 42.0, + call=lambda native, timeout: native( + model="mistral-ocr-maas", + document=DOCUMENT, + api_key=None, + api_base=None, + custom_llm_provider="vertex_ai", + extra_headers=None, + optional_params={"vertex_project": "project-1"}, + timeout_seconds=timeout, + ), + fallback=unexpected_fallback, + adapt=dict, model="mistral-ocr-maas", - document=DOCUMENT, - api_key=None, - api_base=None, - custom_llm_provider="vertex_ai", - extra_headers=None, - optional_params={"vertex_project": "project-1"}, - timeout=httpx.Timeout(30.0, read=42.0), + provider="vertex_ai", + eligible=True, ) assert response == FAKE_OCR_RESPONSE @@ -456,9 +485,10 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): bridge = RecordingBridge() logging_obj = RecordingLogging() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) response = ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( logging_obj=logging_obj, api_base="https://proxy.internal", @@ -489,9 +519,10 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request(api_key=None, timeout=None), resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None, ) @@ -502,12 +533,13 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): def test_run_rust_ocr_prefers_explicit_key_over_resolver(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) def _resolver(name: str) -> str | None: raise AssertionError(f"resolver should not be called for {name}") ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( api_key="sk-explicit", timeout=None, @@ -522,13 +554,14 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): bridge = RecordingBridge() resolver_calls = [] litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) def _resolver(name): resolver_calls.append(name) return "sk-provider-env" ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"), model="provider-ocr-model", @@ -545,9 +578,10 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", @@ -572,7 +606,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) def _resolver(name: str) -> str | None: return { @@ -581,6 +615,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana }.get(name) ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", @@ -596,9 +631,10 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( custom_llm_provider="azure_ai", model="pixtral-12b-2409", @@ -614,9 +650,10 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( custom_llm_provider="azure_ai", model="doc-intelligence/prebuilt-layout", @@ -635,9 +672,10 @@ def test_run_rust_ocr_runs_pre_call_logging(): logging_obj = RecordingLogging() bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.sync.override(bridge) + rust_bridge.set_rust_ocr(ocr=bridge) ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( logging_obj=logging_obj, api_base="https://api.mistral.ai/v1", @@ -661,13 +699,15 @@ def test_run_rust_ocr_runs_pre_call_logging(): } -def test_ocr_routes_to_rust_when_enabled(fake_bridge): +@pytest.mark.parametrize("request_flag", (False, True)) +def test_ocr_routes_to_rust_when_enabled(fake_bridge, request_flag): response = litellm.ocr( model=MODEL, document=DOCUMENT, api_key="sk-test", extra_headers={"x-trace-id": "trace-1"}, include_image_base64=True, + rust=request_flag, ) assert isinstance(response, OCRResponse) @@ -723,7 +763,7 @@ def test_ocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge._OCR.sync.override(RaisingBridge()) + rust_bridge.set_rust_ocr(ocr=RaisingBridge()) with pytest.raises(CapturedException): litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -769,7 +809,7 @@ async def test_aocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge._OCR.asynchronous.override(RaisingAsyncBridge()) + rust_bridge.set_rust_ocr(aocr=RaisingAsyncBridge()) with pytest.raises(CapturedException): await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -794,34 +834,31 @@ def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): assert fake_bridge.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.rust(False) - rust_bridge._OCR.sync.override(bridge) - # 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 == [] +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", (False, True)) +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_ocr_fallback_skips_native_preparation( + monkeypatch: pytest.MonkeyPatch, enabled: bool, asynchronous: bool +) -> None: + monkeypatch.setattr(importlib.import_module("litellm.rust_bridge.bindings"), "get_native_bridge", lambda: None) + litellm.rust(enabled) + expected: Final = OCRResponse(pages=[], model="mistral-ocr-latest", object="ocr") + fallback: Final = AsyncMock(return_value=expected) if asynchronous else Mock(return_value=expected) + def unexpected_preparation(*_args: object, **_kwargs: object) -> None: + pytest.fail("Python fallback must not resolve native credentials or emit native pre_call") -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(rust_bridge, "load_rust_ocr", lambda: None) - litellm.rust(True) # enabled, but load_rust_ocr() returns None in CI + monkeypatch.setattr(ocr_main, "_prepare_rust_ocr_call", unexpected_preparation) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fallback) - captured = {} + response: Final = ( + await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + if asynchronous + else litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + ) - def fake_handler_ocr(**kwargs): - captured["called"] = True - 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 isinstance(response, OCRResponse) + assert response is expected + fallback.assert_called_once() def test_ocr_provider_configs_expose_api_key_env_vars(): diff --git a/type-discipline-budget.json b/type-discipline-budget.json index fb968f6ffcd..1c5671c0b57 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22180 + "limit": 22173 }, "LIT002": { "limit": 26729 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1031 + "limit": 1027 }, "LIT007": { "limit": 0