diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index b918f013700..4a1f9f57fc0 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -29,6 +29,13 @@ 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.request import ( + NativeRequestOptions, + PreparedNativeCall, + provider_connection_params, + provider_request_params, +) +from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -52,14 +59,6 @@ class _PreparedOCRRequest: litellm_logging_obj: LiteLLMLoggingObj -@dataclass -class _PreparedRustOCRCall: - api_key: str | None - api_base: str | None - headers: dict[str, object] - optional_params: dict[str, object] - - _RUST_OCR_PROVIDERS: Final = { "mistral", "azure_ai", @@ -236,7 +235,7 @@ def _rust_bridge_api_base( def _prepare_rust_ocr_call( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], -) -> _PreparedRustOCRCall: +) -> PreparedNativeCall[rust_ocr_bridge.NativeOCRRequest]: provider_config: Final = prepared_request.provider_config api_key_env_var: Final = provider_config.get_api_key_env_var() resolved_api_key: Final = prepared_request.api_key or ( @@ -270,62 +269,59 @@ def _prepare_rust_ocr_call( "headers": resolved_headers, }, ) - return _PreparedRustOCRCall( - api_key=resolved_api_key, - api_base=rust_api_base, - headers=cast(dict[str, object], resolved_headers), - optional_params=rust_optional_params, + return PreparedNativeCall( + request=rust_ocr_bridge.NativeOCRRequest( + model=prepared_request.model, + document=prepared_request.document, + optional_params=provider_request_params(rust_optional_params), + options=NativeRequestOptions( + provider_connection=provider_connection_params(rust_optional_params), + api_key=resolved_api_key, + api_base=rust_api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=cast( # cast-ok: provider header normalization returns string-object pairs + dict[str, object], resolved_headers + ), + timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), + ), + ) ) 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, + ), + 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, + ), + 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 @@ -422,40 +418,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_ocr_bridge.rust_ocr_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, @@ -694,34 +683,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_ocr_bridge.rust_ocr_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 7bd7cdc717c..320f4b6718f 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,40 +2,28 @@ 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.bindings import UNCHANGED, Unchanged -from litellm.rust_bridge.protocols import RustAocr, RustOcr -from litellm.rust_bridge.request import ( - NativeOCRRequest, - NativeRequestContext, - NativeRequestOptions, - PreparedNativeCall, - call_native, - provider_connection_params, - provider_request_params, -) -from litellm.rust_bridge.runtime import ( +from . import configuration as _configuration +from .bindings import UNCHANGED, Unchanged +from .protocols import RustAocr, RustOcr +from .request import NativeOCRRequest, PreparedNativeCall, call_native +from .runtime import ( BridgeErrorContext, EndpointDispatch, - always_enabled, - 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") _OCR: Final[EndpointDispatch[RustOcr, RustAocr]] = EndpointDispatch.native( route="ocr", sync=lambda native: native.ocr, asynchronous=lambda native: native.aocr, - enabled=always_enabled, + enabled=_configuration.rust_ocr_enabled, ) @@ -64,71 +52,39 @@ def load_rust_aocr() -> RustAocr | None: return _OCR.asynchronous.load() -def ocr( +def dispatch_ocr( *, + prepare: Callable[[], PreparedNativeCall[NativeOCRRequest]], + 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: PreparedNativeCall( - NativeOCRRequest( - model=model, - document=document, - optional_params=provider_request_params(optional_params), - options=NativeRequestOptions( - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=_timeout_to_seconds(timeout), - provider_connection=provider_connection_params(optional_params), - ), - ), - context=NativeRequestContext(), - ), + prepare=prepare, call=call_native, - fallback=lambda: None, - adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + fallback=fallback, + adapt=adapt, + error_context=BridgeErrorContext(provider=provider, model=model), + eligible=eligible, ) -async def aocr( +async def adispatch_ocr( *, + prepare: Callable[[], PreparedNativeCall[NativeOCRRequest]], + 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: PreparedNativeCall( - NativeOCRRequest( - model=model, - document=document, - optional_params=provider_request_params(optional_params), - options=NativeRequestOptions( - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=_timeout_to_seconds(timeout), - provider_connection=provider_connection_params(optional_params), - ), - ), - context=NativeRequestContext(), - ), + prepare=prepare, call=call_native, - fallback=async_none, - adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + 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..dc4e8de0260 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -246,7 +246,7 @@ "limit": 109 }, "TRY300": { - "limit": 852 + "limit": 850 }, "UP028": { "limit": 2 diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 9fdb98baa75..dc9b8c0eb17 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -3,7 +3,8 @@ 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 @@ -11,7 +12,8 @@ import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge import configuration -from litellm.rust_bridge.request import NativeOCRRequest, NativeRequestContext +from litellm.rust_bridge.request import NativeOCRRequest, NativeRequestContext, NativeRequestOptions, PreparedNativeCall +from litellm.rust_bridge.timeouts import timeout_to_seconds # `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` # function onto `litellm.ocr` and shadows the submodule, so import the modules @@ -370,9 +372,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(): @@ -381,15 +383,26 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - response = rust_bridge.ocr( + response = rust_bridge.dispatch_ocr( + prepare=lambda: PreparedNativeCall( + request=NativeOCRRequest( + model="mistral-ocr-latest", + document=DOCUMENT, + optional_params={"include_image_base64": True, "pages": [0]}, + options=NativeRequestOptions( + api_key="sk-test", + api_base="https://proxy.internal", + custom_llm_provider="mistral", + extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, + timeout_seconds=12.5, + ), + ), + ), + 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 @@ -416,15 +429,28 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): litellm.rust(True) rust_bridge.set_rust_ocr(aocr=bridge) - response = await rust_bridge.aocr( + + async def unexpected_fallback(): + pytest.fail("unexpected Python fallback") + + response = await rust_bridge.adispatch_ocr( + prepare=lambda: PreparedNativeCall( + request=NativeOCRRequest( + model="mistral-ocr-maas", + document=DOCUMENT, + optional_params={}, + options=NativeRequestOptions( + custom_llm_provider="vertex_ai", + provider_connection={"vertex_project": "project-1"}, + timeout_seconds=42.0, + ), + ), + ), + 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 @@ -447,6 +473,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): 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", @@ -480,6 +507,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): 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, ) @@ -496,6 +524,7 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver(): 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, @@ -517,6 +546,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): 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", @@ -536,6 +566,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): 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", @@ -569,6 +600,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", @@ -587,6 +619,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): 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", @@ -605,6 +638,7 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): 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", @@ -626,6 +660,7 @@ def test_run_rust_ocr_runs_pre_call_logging(): 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", @@ -784,36 +819,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.set_rust_ocr(ocr=bridge) +@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) - 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 == [] + def unexpected_preparation(*_args: object, **_kwargs: object) -> None: + pytest.fail("Python fallback must not resolve native credentials or emit native pre_call") + monkeypatch.setattr(ocr_main, "_prepare_rust_ocr_call", unexpected_preparation) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fallback) -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 + 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") + ) - captured = {} - - 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 57fbd6875d6..a1c7b31407d 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22165 + "limit": 22155 }, "LIT002": { "limit": 26745 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1031 + "limit": 1030 }, "LIT007": { "limit": 0