scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

refactor(ocr): dispatch complete operations through shared runtime

Browse source
This commit is contained in:
Yujong Lee 2026-09-05 12:59:50 -07:00
parent 0f886a6c30
commit 227b561aa3
5 changed files with 201 additions and 231 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

192
litellm/ocr/main.py
View file

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

110
litellm/rust_bridge/ocr.py
View file

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

2
ruff-strict-budget.json
View file

@ -246,7 +246,7 @@
"limit": 109
},
"TRY300": {
"limit": 852
"limit": 850
},
"UP028": {
"limit": 2

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

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

4
type-discipline-budget.json
View file

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