diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs index 0febedc01c3..302a31a759d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs @@ -159,8 +159,8 @@ fn redact( } pub(super) fn response(py: Python<'_>, response: &LiteLLMOcrResponse) -> PyResult> { - 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> { Ok(py - .import("litellm.rust_bridge.ocr.lifecycle")? + .import("litellm.rust_bridge.ocr.callbacks")? .getattr("map_failure")? .call1((error, request, provider))? .extract()?) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs index 32794936899..096ceb47897 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -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> { + run_ocr(py, request, args, kwargs, false) +} + +#[pyfunction] +fn aocr( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, +) -> PyResult> { + 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)?) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index f17bf249b7f..f3683501a62 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -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) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs deleted file mode 100644 index b7d53a97fd6..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs +++ /dev/null @@ -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> + 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, - api_base: Option, - custom_llm_provider: Option, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - extra_headers: Option, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - optional_params: Option, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - input_sources: Option, - timeout_seconds: Option, - }, - prepare = prepare_ocr, - errors = ocr_error_to_pyerr, -} diff --git a/litellm/__init__.py b/litellm/__init__.py index a56d988e801..c6f03172d8e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 * diff --git a/litellm/ocr/__init__.py b/litellm/ocr/__init__.py index a171009564f..4c48f91f76e 100644 --- a/litellm/ocr/__init__.py +++ b/litellm/ocr/__init__.py @@ -1,5 +1,5 @@ """OCR module for LiteLLM.""" -from .rust import aocr, ocr +from .dispatch import aocr, ocr __all__ = ["aocr", "ocr"] diff --git a/litellm/ocr/rust.py b/litellm/ocr/dispatch.py similarity index 72% rename from litellm/ocr/rust.py rename to litellm/ocr/dispatch.py index 5f290e58d14..41f9cc93f2c 100644 --- a/litellm/ocr/rust.py +++ b/litellm/ocr/dispatch.py @@ -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: diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index a20bc1c0811..05bb417f079 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -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", diff --git a/litellm/rust_bridge/ocr/lifecycle.py b/litellm/rust_bridge/ocr/callbacks.py similarity index 57% rename from litellm/rust_bridge/ocr/lifecycle.py rename to litellm/rust_bridge/ocr/callbacks.py index b3a022e46b3..6c2c0573779 100644 --- a/litellm/rust_bridge/ocr/lifecycle.py +++ b/litellm/rust_bridge/ocr/callbacks.py @@ -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]: diff --git a/litellm/rust_bridge/ocr/entrypoints.py b/litellm/rust_bridge/ocr/entrypoints.py new file mode 100644 index 00000000000..5b87634ec16 --- /dev/null +++ b/litellm/rust_bridge/ocr/entrypoints.py @@ -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) diff --git a/litellm/rust_bridge/ocr/native.py b/litellm/rust_bridge/ocr/native.py deleted file mode 100644 index de8a93dd8b1..00000000000 --- a/litellm/rust_bridge/ocr/native.py +++ /dev/null @@ -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), - ) diff --git a/tests/test_litellm/rust_bridge/ocr/test_lifecycle.py b/tests/test_litellm/ocr/test_dispatch.py similarity index 87% rename from tests/test_litellm/rust_bridge/ocr/test_lifecycle.py rename to tests/test_litellm/ocr/test_dispatch.py index fd0a1591305..0dad3cbb466 100644 --- a/tests/test_litellm/rust_bridge/ocr/test_lifecycle.py +++ b/tests/test_litellm/ocr/test_dispatch.py @@ -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) diff --git a/tests/test_litellm/ocr/test_main.py b/tests/test_litellm/ocr/test_main.py index e0d2b5cfeb0..8ff796e388e 100644 --- a/tests/test_litellm/ocr/test_main.py +++ b/tests/test_litellm/ocr/test_main.py @@ -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]) diff --git a/tests/test_litellm/ocr/test_ocr_native_format.py b/tests/test_litellm/rust_bridge/ocr/test_callbacks.py similarity index 76% rename from tests/test_litellm/ocr/test_ocr_native_format.py rename to tests/test_litellm/rust_bridge/ocr/test_callbacks.py index 87d3faaf0fc..c5e9d60ff86 100644 --- a/tests/test_litellm/ocr/test_ocr_native_format.py +++ b/tests/test_litellm/rust_bridge/ocr/test_callbacks.py @@ -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", diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 8eeee1941e9..aa9794a73a6 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -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) diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py index 7657eee2872..8eccbea1a73 100644 --- a/tests/test_litellm_rust/test_ocr.py +++ b/tests/test_litellm_rust/test_ocr.py @@ -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):