scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

cleanup

Browse source
This commit is contained in:
Yujong Lee 2026-09-16 14:16:15 -07:00
parent 64f2a3d098
commit 64cd6538a6
16 changed files with 163 additions and 388 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs
View file

@ -159,8 +159,8 @@ fn redact(
}
pub(super) fn response(py: Python<'_>, response: &LiteLLMOcrResponse) -> PyResult<Py<PyAny>> {
py.import("litellm.rust_bridge.ocr.native")?
.getattr("_response")?
py.import("litellm.rust_bridge.ocr.callbacks")?
.getattr("response")?
.call1((to_py(py, response)?,))
.map(Bound::unbind)
}
@ -172,7 +172,7 @@ pub(super) fn map_failure(
provider: &str,
) -> PyResult<Py<PyBaseException>> {
Ok(py
.import("litellm.rust_bridge.ocr.lifecycle")?
.import("litellm.rust_bridge.ocr.callbacks")?
.getattr("map_failure")?
.call1((error, request, provider))?
.extract()?)

28
litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs
View file

@ -279,8 +279,7 @@ impl litellm_core::ocr::hooks::OcrHooks for BridgeOcrHooks {
}
}
#[pyfunction]
fn _ocr_lifecycle(
fn run_ocr(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
@ -310,6 +309,27 @@ fn _ocr_lifecycle(
run_call(py, call, host)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_function(wrap_pyfunction!(_ocr_lifecycle, module)?)
#[pyfunction]
fn ocr(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_ocr(py, request, args, kwargs, false)
}
#[pyfunction]
fn aocr(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_ocr(py, request, args, kwargs, true)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_function(wrap_pyfunction!(ocr, module)?)?;
module.add_function(wrap_pyfunction!(aocr, module)?)
}

2
litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs
View file

@ -3,12 +3,10 @@ mod document;
mod errors;
mod lifecycle;
mod project;
mod value;
use pyo3::prelude::*;
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
value::register(module)?;
document::register(module)?;
lifecycle::register(module)
}

80
litellm-rust/crates/python-bridge/src/routes/ocr/value.rs
View file

@ -1,80 +0,0 @@
use litellm_core::ocr::Error;
use std::future::Future;
use litellm_core::ocr::wire::{OcrWireRequest, decode_request};
use pyo3::prelude::*;
use serde_json::Value;
use super::errors::to_pyerr as ocr_error_to_pyerr;
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
fn prepare_ocr(
inputs: OcrInputs,
) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
let document = inputs.document;
let options = RouteOptions::from_python(RouteOptionsInputs {
model: inputs.model,
api_key: inputs.api_key,
api_base: inputs.api_base,
custom_llm_provider: inputs.custom_llm_provider,
extra_headers: inputs.extra_headers,
timeout_seconds: inputs.timeout_seconds,
})?;
let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
let input_sources = inputs
.input_sources
.map(serde_json::from_value)
.transpose()
.map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?
.unwrap_or_default();
Ok(async move {
let RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout,
} = options;
let request = decode_request(OcrWireRequest {
model,
document,
api_key,
api_base,
custom_llm_provider,
extra_headers,
optional_params,
input_sources,
timeout_seconds: timeout.map(|value| value.as_secs_f64()),
})?;
litellm_core::ocr::ocr(request)
.await
.map(|response| response.into_json())
})
}
bridge_route! {
sync = ocr,
asynchronous = aocr,
inputs = OcrInputs,
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
document: serde_json::Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
input_sources: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_ocr,
errors = ocr_error_to_pyerr,
}

2
litellm/__init__.py
View file

@ -1434,7 +1434,7 @@ from .skills.main import (
adelete_skill,
)
from .containers.main import *
from .ocr.rust import *
from .ocr.dispatch import *
from .rust_bridge import rust
from .rag.main import *
from .sandbox.main import *

2
litellm/ocr/__init__.py
View file

@ -1,5 +1,5 @@
"""OCR module for LiteLLM."""
from .rust import aocr, ocr
from .dispatch import aocr, ocr
__all__ = ["aocr", "ocr"]

27
litellm/ocr/rust.py → litellm/ocr/dispatch.py
View file

@ -7,8 +7,7 @@ from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.ocr import main
from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type
from litellm.rust_bridge.catalog import Context, Route
from litellm.rust_bridge.ocr.lifecycle import NATIVE_OCR_LIFECYCLE, NativeOcrLifecycle
from litellm.rust_bridge.ocr.native import LiteLLMOcrRequest
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR, LiteLLMOcrRequest, NativeAocr
from litellm.rust_bridge.runtime import arun, run
__all__ = ("aocr", "convert_file_document_to_url_document", "get_mime_type", "ocr")
@ -48,18 +47,16 @@ def ocr(
**kwargs: object, # kwargs-ok: preserve the public OCR call shape
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
request: Final = _public_request("ocr", args, kwargs)
fallback: Final = cast( # cast-ok: forward the original call shape through the Python @client decorator
python_ocr: Final = cast( # cast-ok: forward the original call shape through the Python @client decorator
Callable[..., OCRResponse | Coroutine[object, object, OCRResponse]], main.ocr
)
if request.kwargs.get("aocr"):
return fallback(*args, **kwargs)
if request.kwargs.get("aocr") is True:
return python_ocr(*args, **kwargs)
return run(
_context(request),
binding=NATIVE_OCR_LIFECYCLE,
native=lambda hook: cast( # cast-ok: False selects the synchronous result
OCRResponse, hook(request, args, kwargs, False)
),
python=lambda: fallback(*args, **kwargs),
binding=NATIVE_OCR,
native=lambda hook: hook(request, args, kwargs),
python=lambda: python_ocr(*args, **kwargs),
)
@ -69,14 +66,10 @@ async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: pr
Callable[..., Awaitable[OCRResponse]], main.aocr
)
async def native(hook: NativeOcrLifecycle) -> OCRResponse:
return await cast( # cast-ok: True selects the asynchronous result
Awaitable[OCRResponse], hook(request, args, kwargs, True)
)
async def native(hook: NativeAocr) -> OCRResponse:
return await hook(request, args, kwargs)
return await arun(
_context(request), binding=NATIVE_OCR_LIFECYCLE, native=native, python=lambda: fallback(*args, **kwargs)
)
return await arun(_context(request), binding=NATIVE_AOCR, native=native, python=lambda: fallback(*args, **kwargs))
def _context(request: LiteLLMOcrRequest) -> Context:

41
litellm/rust_bridge/_native.pyi
View file

@ -1,37 +1,23 @@
from asyncio import Future
from collections.abc import Coroutine, Mapping, Sequence
from typing import Literal, Never, TypeAlias, final
from typing import Never, final
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.ocr.native import LiteLLMOcrRequest
_InputSource: TypeAlias = Literal["request", "deployment", "environment"]
from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest
class RustBridgeDeclined(Exception): ...
class RustUpstreamError(Exception): ...
def ocr(
model: str,
document: object,
api_key: str | None = None,
api_base: str | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, object] | None = None,
optional_params: Mapping[str, object] | None = None,
input_sources: Mapping[str, _InputSource] | None = None,
timeout_seconds: float | None = None,
) -> dict[str, object]: ...
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: dict[str, object],
) -> OCRResponse: ...
def aocr(
model: str,
document: object,
api_key: str | None = None,
api_base: str | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, object] | None = None,
optional_params: Mapping[str, object] | None = None,
input_sources: Mapping[str, _InputSource] | None = None,
timeout_seconds: float | None = None,
) -> Future[dict[str, object]]: ...
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: dict[str, object],
) -> Coroutine[object, object, OCRResponse]: ...
_OCR_MAX_FILE_BYTES: int
@ -42,12 +28,6 @@ def _ocr_upload_document(
) -> dict[str, str]: ...
def _ocr_file_document(document: Mapping[str, object]) -> dict[str, str]: ...
def _ocr_mime_type(file_name: str) -> str: ...
def _ocr_lifecycle(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: dict[str, object],
asynchronous: bool,
) -> OCRResponse | Coroutine[object, object, OCRResponse]: ...
def transcription(
model: str,
audio: object,
@ -145,7 +125,6 @@ __all__ = [
"RustUpstreamError",
"TokenCounter",
"_ocr_file_document",
"_ocr_lifecycle",
"_ocr_mime_type",
"_ocr_upload_document",
"achat_completions",

37
litellm/rust_bridge/ocr/lifecycle.py → litellm/rust_bridge/ocr/callbacks.py
View file

@ -1,22 +1,16 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final, Protocol, cast # noqa: TID251 # adapts the public exception mapper
from pydantic import TypeAdapter
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.ocr.native import LiteLLMOcrRequest
from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse
from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest
class NativeOcrLifecycle(Protocol):
def __call__(
self,
request: LiteLLMOcrRequest,
args: Sequence[object],
kwargs: Mapping[str, object],
asynchronous: bool,
) -> OCRResponse | Awaitable[OCRResponse]: ...
_RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object])
class ExceptionMapper(Protocol):
@ -31,13 +25,14 @@ class ExceptionMapper(Protocol):
) -> Exception: ...
def _binding(value: object) -> NativeOcrLifecycle | None:
if not callable(value):
return None
return cast("NativeOcrLifecycle", value) # cast-ok: callable validated at the native binding boundary
NATIVE_OCR_LIFECYCLE: Final = NativeBinding("_ocr_lifecycle", validate=_binding)
def response(value: Mapping[str, object]) -> OCRResponse:
provider_native_response: Final = value.get(PROVIDER_NATIVE_RESPONSE_KEY)
normalized: Final = OCRResponse.model_validate(
MappingProxyType({key: item for key, item in value.items() if key != PROVIDER_NATIVE_RESPONSE_KEY})
)
if isinstance(provider_native_response, Mapping):
normalized.set_provider_native_response(_RESPONSE_ADAPTER.validate_python(provider_native_response))
return normalized
def arguments(request: LiteLLMOcrRequest) -> Mapping[str, object]:

57
litellm/rust_bridge/ocr/entrypoints.py Normal file
View file

@ -0,0 +1,57 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping
from dataclasses import dataclass
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
import httpx
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.bindings import NativeBinding
@dataclass(frozen=True, slots=True)
class LiteLLMOcrRequest:
model: str
document: Mapping[str, object]
api_key: str | None
api_base: str | None
timeout: float | httpx.Timeout | None
custom_llm_provider: str | None
extra_headers: dict[str, object] | None
kwargs: Mapping[str, object]
input_sources: Mapping[str, str] | None = None
class NativeOcr(Protocol):
def __call__(
self,
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse: ...
class NativeAocr(Protocol):
def __call__(
self,
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Awaitable[OCRResponse]: ...
def _ocr_binding(value: object) -> NativeOcr | None:
if not callable(value):
return None
return cast("NativeOcr", value) # cast-ok: callable validated at the native binding boundary
def _aocr_binding(value: object) -> NativeAocr | None:
if not callable(value):
return None
return cast("NativeAocr", value) # cast-ok: callable validated at the native binding boundary
NATIVE_OCR: Final = NativeBinding("ocr", validate=_ocr_binding)
NATIVE_AOCR: Final = NativeBinding("aocr", validate=_aocr_binding)

145
litellm/rust_bridge/ocr/native.py
View file

@ -1,145 +0,0 @@
"""Thin Python wrapper for the native Rust OCR bridge."""
from __future__ import annotations
from collections.abc import Awaitable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables
import httpx
from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds
@dataclass(frozen=True, slots=True)
class LiteLLMOcrRequest:
model: str
document: Mapping[str, object]
api_key: str | None
api_base: str | None
timeout: float | httpx.Timeout | None
custom_llm_provider: str | None
extra_headers: dict[str, object] | None
kwargs: Mapping[str, object]
input_sources: Mapping[str, str] | None = None
class RustOcr(Protocol):
def __call__(
self,
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],
input_sources: dict[str, str],
timeout_seconds: float | None,
) -> dict[str, object]:
raise NotImplementedError
class RustAocr(Protocol):
def __call__(
self,
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],
input_sources: dict[str, str],
timeout_seconds: float | None,
) -> Awaitable[dict[str, object]]:
raise NotImplementedError
def _as_ocr(value: object) -> RustOcr | None:
return cast(RustOcr, value) if callable(value) else None
def _as_aocr(value: object) -> RustAocr | None:
return cast(RustAocr, value) if callable(value) else None
_OCR: Final = NativeBinding("ocr", validate=_as_ocr)
_AOCR: Final = NativeBinding("aocr", validate=_as_aocr)
def load_rust_ocr() -> RustOcr | None:
return _OCR.load()
def load_rust_aocr() -> RustAocr | None:
return _AOCR.load()
def _response(response: Mapping[str, object]) -> OCRResponse:
provider_native_response: Final = response.get(PROVIDER_NATIVE_RESPONSE_KEY)
normalized: Final = OCRResponse.model_validate(
MappingProxyType({key: value for key, value in response.items() if key != PROVIDER_NATIVE_RESPONSE_KEY})
)
if isinstance(provider_native_response, Mapping):
normalized.set_provider_native_response(provider_native_response)
return normalized
def ocr(
*,
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,
input_sources: Mapping[str, str] | None = None,
) -> dict[str, object] | None:
rust_ocr: Final = load_rust_ocr()
if rust_ocr is None:
return None
return 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,
input_sources=dict(input_sources or {}), # mutable-ok: native boundary requires a concrete dict
timeout_seconds=_timeout_to_seconds(timeout),
)
async def aocr(
*,
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,
input_sources: Mapping[str, str] | None = None,
) -> dict[str, object] | None:
rust_aocr: Final = load_rust_aocr()
if rust_aocr is None:
return None
return await 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,
input_sources=dict(input_sources or {}), # mutable-ok: native boundary requires a concrete dict
timeout_seconds=_timeout_to_seconds(timeout),
)

58
tests/test_litellm/rust_bridge/ocr/test_lifecycle.py → tests/test_litellm/ocr/test_dispatch.py
View file

@ -8,8 +8,7 @@ import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.ocr import main as python_ocr
from litellm.rust_bridge import bindings, configuration, runtime
from litellm.rust_bridge.ocr.lifecycle import NATIVE_OCR_LIFECYCLE
from litellm.rust_bridge.ocr.native import LiteLLMOcrRequest
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR, LiteLLMOcrRequest
@pytest.fixture(autouse=True)
@ -17,17 +16,21 @@ def isolated_ocr_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[Non
monkeypatch.delenv("LITELLM_RUST", raising=False)
configuration.reset_rust_configuration()
yield
NATIVE_OCR_LIFECYCLE.reset()
NATIVE_OCR.reset()
NATIVE_AOCR.reset()
configuration.reset_rust_configuration()
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_unavailable_native_uses_legacy(monkeypatch: pytest.MonkeyPatch, asynchronous: bool) -> None:
async def test_unavailable_native_uses_python(monkeypatch: pytest.MonkeyPatch, asynchronous: bool) -> None:
response: Final = OCRResponse(pages=[], model="mistral-ocr-latest")
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(python_ocr, "aocr" if asynchronous else "ocr", fallback)
NATIVE_OCR_LIFECYCLE.override(None)
if asynchronous:
NATIVE_AOCR.override(None)
else:
NATIVE_OCR.override(None)
document: Final = {"type": "document_url", "document_url": "https://example.com"}
result: Final = (
@ -44,67 +47,64 @@ def test_admitted_failure_is_returned_without_replay() -> None:
failure: Final = RuntimeError("admitted")
native: Final = Mock(side_effect=failure)
litellm.rust(True)
NATIVE_OCR_LIFECYCLE.override(native)
NATIVE_OCR.override(native)
try:
with pytest.raises(RuntimeError) as caught:
litellm.ocr("mistral/mistral-ocr-latest", {"type": "document_url", "document_url": "https://example.com"})
assert caught.value is failure
finally:
NATIVE_OCR_LIFECYCLE.reset()
NATIVE_OCR.reset()
litellm.rust(None)
assert native.call_count == 1
def test_public_binding_keeps_positional_fields_and_defaults_out_of_native_hook_kwargs() -> None:
document: Final = {"type": "document_url", "document_url": "https://example.com"}
captured: Final = []
captured: Final[list[tuple[LiteLLMOcrRequest, tuple[object, ...], Mapping[str, object]]]] = []
def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
asynchronous: bool,
) -> OCRResponse:
captured.append((request, args, kwargs, asynchronous))
captured.append((request, args, kwargs))
return OCRResponse(pages=[], model=request.model)
litellm.rust(True)
NATIVE_OCR_LIFECYCLE.override(native)
NATIVE_OCR.override(native)
try:
response: Final = litellm.ocr("mistral/mistral-ocr-latest", document)
finally:
NATIVE_OCR_LIFECYCLE.reset()
NATIVE_OCR.reset()
litellm.rust(None)
request, call_args, hook_kwargs, asynchronous = captured[0]
request, call_args, hook_kwargs = captured[0]
assert response.model == "mistral/mistral-ocr-latest"
assert request.model == "mistral/mistral-ocr-latest"
assert request.document is document
assert call_args == ("mistral/mistral-ocr-latest", document)
assert hook_kwargs == {}
assert asynchronous is False
def test_public_binding_keeps_keyword_model_and_document_in_native_hook_kwargs() -> None:
document: Final = {"type": "document_url", "document_url": "https://example.com"}
captured: Final = []
captured: Final[list[Mapping[str, object]]] = []
def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
asynchronous: bool,
) -> OCRResponse:
assert args == ()
captured.append(kwargs)
return OCRResponse(pages=[], model=request.model)
litellm.rust(True)
NATIVE_OCR_LIFECYCLE.override(native)
NATIVE_OCR.override(native)
try:
litellm.ocr(model="mistral/mistral-ocr-latest", document=document)
finally:
NATIVE_OCR_LIFECYCLE.reset()
NATIVE_OCR.reset()
litellm.rust(None)
assert captured[0]["model"] == "mistral/mistral-ocr-latest"
@ -117,12 +117,12 @@ def test_public_duplicate_argument_error_does_not_depend_on_native_selection(ena
native: Final = Mock(side_effect=AssertionError("binding errors precede admission"))
document: Final = {"type": "document_url", "document_url": "https://example.com"}
litellm.rust(enabled)
NATIVE_OCR_LIFECYCLE.override(native)
NATIVE_OCR.override(native)
try:
with pytest.raises(TypeError, match=r"ocr\(\) got multiple values for argument 'model'"):
litellm.ocr("mistral/mistral-ocr-latest", document, model="duplicate")
finally:
NATIVE_OCR_LIFECYCLE.reset()
NATIVE_OCR.reset()
litellm.rust(None)
assert native.call_count == 0
@ -131,12 +131,12 @@ def test_public_duplicate_argument_error_does_not_depend_on_native_selection(ena
def test_public_missing_required_argument_error_does_not_depend_on_native_selection(enabled: bool) -> None:
native: Final = Mock(side_effect=AssertionError("binding errors precede admission"))
litellm.rust(enabled)
NATIVE_OCR_LIFECYCLE.override(native)
NATIVE_OCR.override(native)
try:
with pytest.raises(TypeError, match=r"ocr\(\) missing 1 required positional argument: 'document'"):
litellm.ocr("mistral/mistral-ocr-latest")
finally:
NATIVE_OCR_LIFECYCLE.reset()
NATIVE_OCR.reset()
litellm.rust(None)
assert native.call_count == 0
@ -177,8 +177,11 @@ async def test_native_is_enabled_by_default(
monkeypatch.setenv("LITELLM_RUST", environment)
response: Final = OCRResponse(pages=[], model="mistral-ocr-latest")
native: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
NATIVE_OCR_LIFECYCLE.override(native)
fallback: Final = Mock(side_effect=AssertionError("legacy must not run"))
if asynchronous:
NATIVE_AOCR.override(native)
else:
NATIVE_OCR.override(native)
fallback: Final = Mock(side_effect=AssertionError("Python must not run"))
monkeypatch.setattr(python_ocr, "aocr" if asynchronous else "ocr", fallback)
result: Final = (
@ -203,12 +206,15 @@ class Upstream(Exception):
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("declined", [False, True])
async def test_only_native_declines_replay_on_legacy(
async def test_only_native_declines_replay_on_python(
monkeypatch: pytest.MonkeyPatch, asynchronous: bool, declined: bool
) -> None:
failure: Final = Declined("unsupported") if declined else RuntimeError("provider already called")
native: Final = AsyncMock(side_effect=failure) if asynchronous else Mock(side_effect=failure)
NATIVE_OCR_LIFECYCLE.override(native)
if asynchronous:
NATIVE_AOCR.override(native)
else:
NATIVE_OCR.override(native)
monkeypatch.setattr(runtime, "native_exception_types", lambda: (Declined, Upstream))
response: Final = OCRResponse(pages=[], model="mistral-ocr-latest")
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)

8
tests/test_litellm/ocr/test_main.py
View file

@ -16,7 +16,7 @@ from litellm.llms.custom_httpx import llm_http_handler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.ocr.main import _prepare_ocr_request
from litellm.rust_bridge import bindings, configuration, runtime
from litellm.rust_bridge.ocr.lifecycle import NATIVE_OCR_LIFECYCLE
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
@pytest.fixture
@ -44,7 +44,8 @@ async def provider(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[Mock]:
monkeypatch.setattr(llm_http_handler, "_get_httpx_client", lambda: sync_handler)
monkeypatch.setattr(llm_http_handler, "get_async_httpx_client", lambda llm_provider: async_handler)
yield handler
NATIVE_OCR_LIFECYCLE.reset()
NATIVE_OCR.reset()
NATIVE_AOCR.reset()
configuration.reset_rust_configuration()
@ -59,7 +60,8 @@ async def test_python_request_response_and_callbacks(
if dispatch != "disabled":
monkeypatch.setenv("LITELLM_RUST", "1")
NATIVE_OCR_LIFECYCLE.override(Mock(side_effect=Declined()) if dispatch == "declined" else None)
binding: Final = NATIVE_AOCR if mode == "async" else NATIVE_OCR
binding.override(Mock(side_effect=Declined()) if dispatch == "declined" else None)
monkeypatch.setattr(runtime, "native_exception_types", lambda: (Declined, RuntimeError))
logger: Final = Mock(spec=CustomLogger)
monkeypatch.setattr(litellm, "input_callback", [logger])

8
tests/test_litellm/ocr/test_ocr_native_format.py → tests/test_litellm/rust_bridge/ocr/test_callbacks.py
View file

@ -1,13 +1,9 @@
"""
Tests for the OCR `req_format` option in the SDK request path.
"""
from litellm.rust_bridge.ocr import native as rust_ocr_bridge
from litellm.rust_bridge.ocr.callbacks import response as build_ocr_response
def test_rust_ocr_response_retains_provider_native_response():
provider_response = {"status": "succeeded", "analyzeResult": {"content": "native"}}
response = rust_ocr_bridge._response(
response = build_ocr_response(
{
"pages": [],
"model": "prebuilt-layout",

4
tests/test_litellm_rust/ocr/test_lifecycle.py
View file

@ -580,7 +580,7 @@ async def test_retained_argument_aliases_and_body_roots_survive_envelope_replace
def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_server: RecordingServer) -> None:
from litellm.ocr.rust import _public_request
from litellm.ocr.dispatch import _public_request
from litellm.rust_bridge import _native
ocr_server.expected_requests = 0
@ -594,7 +594,7 @@ def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_serv
def create():
file: Final = File()
kwargs: Final = {"model": "mistral/mistral-ocr-latest", "document": {"type": "file", "file": file}}
coroutine: Final = _native._ocr_lifecycle(_public_request("aocr", (), kwargs), (), kwargs, True)
coroutine: Final = _native.aocr(_public_request("aocr", (), kwargs), (), kwargs)
file.owner = coroutine
coroutine.close()
return weakref.ref(file)

46
tests/test_litellm_rust/test_ocr.py
View file

@ -8,7 +8,6 @@ from typing import Final
import pytest
import litellm
from litellm.rust_bridge.ocr import native as rust_ocr_bridge
pytestmark = pytest.mark.requires_rust_extension
@ -71,35 +70,6 @@ def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[dict[str, object]]
thread.join()
def test_native_ocr_with_compiled_rust_extension(
ocr_server: tuple[ThreadingHTTPServer, list[dict[str, object]]],
) -> None:
server, requests = ocr_server
address: Final = server.server_address
host: Final = str(address[0])
port: Final = int(address[1])
response: Final = rust_ocr_bridge.ocr(
model="mistral-ocr-latest",
document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
api_key="test-key",
api_base=f"http://{host}:{port}",
custom_llm_provider="mistral",
extra_headers=None,
optional_params={},
timeout=None,
)
assert response is not None
assert response["pages"][0]["markdown"] == "native OCR response"
assert len(requests) == 1
assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx")
assert requests[0]["body"] == {
"model": "mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
}
@pytest.mark.parametrize(
"file_input,mime_type,expected_type,expected_field,expected_uri",
[
@ -219,22 +189,6 @@ async def test_native_ocr_failures_do_not_retry_on_python(ocr_server, asynchrono
assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx")
@pytest.mark.parametrize("custom_provider", ["mistral", "not-a-provider"])
def test_native_ocr_rejects_invalid_input_before_network(ocr_server, custom_provider):
from litellm.rust_bridge import _native
server, requests = ocr_server
with pytest.raises(ValueError, match="Document URL is required"):
_native.ocr(
model="mistral-ocr-latest",
custom_llm_provider=custom_provider,
document={"type": "document_url"},
api_key="test-key",
api_base=f"http://127.0.0.1:{server.server_port}",
)
assert requests == []
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_native_ocr_enforces_request_deadline_without_fallback(ocr_server, asynchronous):