mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
refactor(ocr): dispatch complete operations through shared runtime
This commit is contained in:
Yujong Lee
yujonglee
parent
20de0ce755
commit
f9a1fd42af
6 changed files with 264 additions and 231 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue