diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 56c2292d00c..df3f9d2096b 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -29,6 +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.bindings import native_exception_types from litellm.rust_bridge.configuration import rust_enabled from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -281,6 +282,33 @@ def _prepare_rust_ocr_call( ) +def _map_rust_ocr_error( + error: Exception, + prepared_request: _PreparedOCRRequest, + exception_types: tuple[type[BaseException], type[BaseException]] | None, +) -> Exception: + if exception_types is None: + return error + _, upstream_error = exception_types + if not isinstance(error, upstream_error): + return error + error_args: Final = cast( # cast-ok: BaseException.args is typed with Any in the standard library stubs + tuple[object, ...], error.args + ) + status_value: Final = error_args[0] if error_args else 0 + message_value: Final = error_args[1] if len(error_args) > 1 else str(error) + status: Final = status_value if isinstance(status_value, int) else 0 + message: Final = message_value if isinstance(message_value, str) else str(message_value) + error_factory: Final = cast( # cast-ok: the legacy provider interface leaves callable parameters untyped + Callable[..., Exception], prepared_request.provider_config.get_error_class + ) + return error_factory( + error_message=message, + status_code=status or 500, + headers={}, # mutable-ok: provider error factories require a concrete header dict + ) + + def _run_rust_ocr( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], @@ -291,16 +319,19 @@ def _run_rust_ocr( prepared_request=prepared_request, resolve_api_key=resolve_api_key, ) - rust_response: Final = rust_ocr_bridge.ocr( - 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, - ) + try: + rust_response: Final = rust_ocr_bridge.ocr( + 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, + ) + except Exception as error: + raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error if rust_response is None: return None return OCRResponse.model_validate(rust_response) @@ -316,16 +347,19 @@ async def _run_rust_aocr( prepared_request=prepared_request, resolve_api_key=resolve_api_key, ) - rust_response: Final = await rust_ocr_bridge.aocr( - 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, - ) + try: + rust_response: Final = await rust_ocr_bridge.aocr( + 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, + ) + except Exception as error: + raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error if rust_response is None: return None return OCRResponse.model_validate(rust_response) diff --git a/tests/rust-python-harness/shared/parity/compare.py b/tests/rust-python-harness/shared/parity/compare.py index adf85e5c8d7..88239ed042a 100644 --- a/tests/rust-python-harness/shared/parity/compare.py +++ b/tests/rust-python-harness/shared/parity/compare.py @@ -78,3 +78,4 @@ def assert_parity(baseline: Execution, candidate: Execution, baseline_user_agent validate_harness(baseline, candidate, baseline_user_agent) assert_request_parity(baseline.requests, candidate.requests) assert_value_parity(baseline.report, candidate.report) + assert_value_parity(baseline.callbacks, candidate.callbacks, path="$.callbacks") diff --git a/tests/rust-python-harness/shared/parity/models.py b/tests/rust-python-harness/shared/parity/models.py index 898b58d23ee..5612e218824 100644 --- a/tests/rust-python-harness/shared/parity/models.py +++ b/tests/rust-python-harness/shared/parity/models.py @@ -37,6 +37,25 @@ class SDKError(BaseModel): llm_provider: str | None +class CallbackObservation(BaseModel): + model_config = ConfigDict(frozen=True) + + hook: Literal[ + "log_success_event", + "async_log_success_event", + "log_failure_event", + "async_log_failure_event", + ] + phase: Literal["success", "failure"] + model: str | None + call_type: str | None + litellm_call_id: str | None + metadata: JsonValue + kwargs: JsonValue + payload: JsonValue + error: SDKError | None + + class SDKJsonChunk(BaseModel): model_config = ConfigDict(frozen=True) @@ -119,6 +138,7 @@ class Execution(BaseModel): requests: tuple[CapturedRequest, ...] report: SDKReport + callbacks: tuple[CallbackObservation, ...] | None = None class SDKCommand(BaseModel): @@ -133,6 +153,7 @@ class WorkerSuccess(BaseModel): status: Literal["ok"] = "ok" report: SDKReport + callbacks: tuple[CallbackObservation, ...] | None = None class WorkerFailure(BaseModel): diff --git a/tests/rust-python-harness/shared/parity/runner.py b/tests/rust-python-harness/shared/parity/runner.py index 43a583382cb..feae0201aee 100644 --- a/tests/rust-python-harness/shared/parity/runner.py +++ b/tests/rust-python-harness/shared/parity/runner.py @@ -39,9 +39,7 @@ class SubprocessRunner: return ( sys.executable, "-m", - ".".join( - self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts - ), + ".".join(self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts), "--parity-worker", provider_url, ) @@ -113,7 +111,11 @@ class SubprocessWorker: ) assert isinstance(result, WorkerSuccess) try: - return Execution(requests=self.provider.take_requests(len(responses)), report=result.report) + return Execution( + requests=self.provider.take_requests(len(responses)), + report=result.report, + callbacks=result.callbacks, + ) except AssertionError: self.provider.reset() raise diff --git a/tests/rust-python-harness/shared/parity/test_parity.py b/tests/rust-python-harness/shared/parity/test_parity.py index 83daccdf8ba..7c8355fbfd8 100644 --- a/tests/rust-python-harness/shared/parity/test_parity.py +++ b/tests/rust-python-harness/shared/parity/test_parity.py @@ -7,7 +7,7 @@ import pytest from pydantic import BaseModel, ConfigDict, JsonValue, PrivateAttr from .compare import assert_model_parity, assert_parity -from .models import CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report +from .models import CallbackObservation, CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report SENTINEL: Final = "python-parity-fallback" @@ -69,6 +69,56 @@ def test_parity_rejects_response_difference() -> None: assert_parity(python, rust, SENTINEL) +def test_parity_distinguishes_unobserved_callbacks_from_zero_events() -> None: + python: Final = _execution(user_agent=SENTINEL) + rust: Final = _execution(user_agent="litellm-rust").model_copy(update={"callbacks": ()}) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + +def test_parity_rejects_callback_payload_difference() -> None: + observation: Final = CallbackObservation( + hook="log_success_event", + phase="success", + model="test-model", + call_type="ocr", + litellm_call_id="test-call", + metadata={"profile": "success"}, + kwargs={"model": "test-model"}, + payload={"model": "test-model", "pages": []}, + error=None, + ) + python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)}) + rust: Final = _execution(user_agent="litellm-rust").model_copy( + update={"callbacks": (observation.model_copy(update={"payload": {"model": "changed", "pages": []}}),)} + ) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + +def test_parity_rejects_callback_kwargs_difference() -> None: + observation: Final = CallbackObservation( + hook="log_success_event", + phase="success", + model="test-model", + call_type="ocr", + litellm_call_id="test-call", + metadata={"profile": "success"}, + kwargs={"model": "test-model"}, + payload={"model": "test-model", "pages": []}, + error=None, + ) + python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)}) + rust: Final = _execution(user_agent="litellm-rust").model_copy( + update={"callbacks": (observation.model_copy(update={"kwargs": {"model": "changed"}}),)} + ) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + def test_parity_rejects_error_difference() -> None: python: Final = Execution( requests=(), diff --git a/tests/rust-python-harness/strategies/e2e_parity/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/__init__.py index f668e178eef..95346bbd496 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/__init__.py +++ b/tests/rust-python-harness/strategies/e2e_parity/__init__.py @@ -18,8 +18,8 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( coverage=Coverage.PARTIAL, module="tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.test_sdk_parity", note=( - "Recorded sync/async SDK parity; invalid-model provider errors differ, " - "and Reducto lacks a Rust contract." + "Recorded sync/async SDK parity with focused success/error callback profiles; " + "Reducto lacks a Rust contract, and known provider parity gaps remain." ), ), surface="sdk", diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py index e72980752f2..5c4abc78081 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py @@ -1,23 +1,31 @@ from __future__ import annotations import asyncio +import datetime +import queue import sys import tempfile +import time import traceback -from collections.abc import Callable, Coroutine, Generator +from collections.abc import Callable, Coroutine, Generator, Mapping from contextlib import contextmanager from enum import Enum from functools import partial from pathlib import Path from typing import Annotated, Final, Literal, cast +from urllib.parse import urlsplit, urlunsplit from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse from .....shared.parity.compare import assert_parity from .....shared.parity.fixtures.store import fixture_id, recorded_fixtures from .....shared.parity.models import ( + JSON_VALUE_ADAPTER, + CallbackObservation, + Execution, SDKCommand, SDKError, SDKReport, @@ -39,6 +47,9 @@ from .fixtures.config import configured_fixture_directory from .fixtures.models import OcrParityCase, OcrSdkInput API_KEY: Final = "test-key" +CALLBACK_DELAY_SECONDS: Final = 0.05 +CALLBACK_DRAIN_TIMEOUT_SECONDS: Final = 10.0 +CALLBACK_TERMINALS: Final[tuple[Literal["success", "failure"], ...]] = ("success", "failure") PYTHON_HTTP_SENTINEL: Final = "python-ocr-parity-fallback" PYTHON_VARIANT: Final = ExecutionVariant(name="Python", environment=(("LITELLM_RUST", "0"),)) RUST_VARIANT: Final = ExecutionVariant(name="Rust", environment=(("LITELLM_RUST", "1"),)) @@ -75,10 +86,148 @@ class InvalidOcrWorkerCase(BaseModel): case: InvalidOcrCase -OcrWorkerCase = Annotated[RecordedOcrWorkerCase | InvalidOcrWorkerCase, Field(discriminator="kind")] +class CallbackOcrWorkerCase(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["callback"] = "callback" + case: OcrParityCase + terminal: Literal["success", "failure"] + + +OcrWorkerCase = Annotated[ + RecordedOcrWorkerCase | InvalidOcrWorkerCase | CallbackOcrWorkerCase, + Field(discriminator="kind"), +] OCR_WORKER_CASE_ADAPTER: Final[TypeAdapter[OcrWorkerCase]] = TypeAdapter(OcrWorkerCase) +class RecordingCallback(CustomLogger): + def __init__(self) -> None: + self.message_logging: Final = True + self.turn_off_message_logging: Final = False + self._observations: Final[queue.SimpleQueue[CallbackObservation]] = queue.SimpleQueue() + + def _normalized_kwargs(self, value: object, key: str | None = None) -> JsonValue: + if value is None or isinstance(value, (bool, int, float)): + return value + if isinstance(value, str): + if key != "api_base": + return value + parsed: Final = urlsplit(value) + return urlunsplit(("", "", parsed.path, parsed.query, parsed.fragment)) + if isinstance(value, datetime.datetime): + return "datetime" + if isinstance(value, Exception): + return sdk_error_report(value).model_dump(mode="json") + if isinstance(value, BaseModel): + return self._normalized_kwargs(value.model_dump(mode="json"), key) + if isinstance(value, Mapping): + if any(not isinstance(map_key, str) for map_key in value): + raise TypeError("callback kwarg mappings must use string keys") + return { + map_key: self._normalized_kwargs(map_value, map_key) + for map_key, map_value in value.items() + } + if isinstance(value, (list, tuple)): + return [self._normalized_kwargs(item) for item in value] + raise TypeError(f"unsupported callback kwarg type: {type(value)}") + + def _record( + self, + hook: Literal[ + "log_success_event", + "async_log_success_event", + "log_failure_event", + "async_log_failure_event", + ], + phase: Literal["success", "failure"], + kwargs: dict[str, object], + response_obj: object, + ) -> None: + raw_litellm_params: Final = kwargs.get("litellm_params") + litellm_params: Final[Mapping[str, object]] = ( + cast(Mapping[str, object], raw_litellm_params) if isinstance(raw_litellm_params, Mapping) else {} + ) + raw_metadata: Final = litellm_params.get("metadata") + metadata_mapping: Final[Mapping[str, object]] = ( + cast(Mapping[str, object], raw_metadata) if isinstance(raw_metadata, Mapping) else {} + ) + metadata: Final = JSON_VALUE_ADAPTER.validate_python( + {key: metadata_mapping[key] for key in ("callback_profile", "sdk_route") if key in metadata_mapping} + ) + raw_error: Final = kwargs.get("exception") + error: Final = sdk_error_report(raw_error) if isinstance(raw_error, Exception) else None + normalized_kwargs: Final = self._normalized_kwargs(kwargs) + payload_source: Final = ( + response_obj.model_dump(mode="json") if isinstance(response_obj, BaseModel) else response_obj + ) + payload: Final = JSON_VALUE_ADAPTER.validate_python(payload_source) + raw_model: Final = kwargs.get("model") + raw_call_type: Final = kwargs.get("call_type") + raw_call_id: Final = kwargs.get("litellm_call_id") + self._observations.put( + CallbackObservation( + hook=hook, + phase=phase, + model=raw_model if isinstance(raw_model, str) else None, + call_type=str(raw_call_type) if raw_call_type is not None else None, + litellm_call_id=raw_call_id if isinstance(raw_call_id, str) else None, + metadata=metadata, + kwargs=normalized_kwargs, + payload=payload, + error=error, + ) + ) + + def log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + time.sleep(CALLBACK_DELAY_SECONDS) + self._record("log_success_event", "success", kwargs, response_obj) + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + await asyncio.sleep(CALLBACK_DELAY_SECONDS) + self._record("async_log_success_event", "success", kwargs, response_obj) + + def log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + time.sleep(CALLBACK_DELAY_SECONDS) + self._record("log_failure_event", "failure", kwargs, response_obj) + + async def async_log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + await asyncio.sleep(CALLBACK_DELAY_SECONDS) + self._record("async_log_failure_event", "failure", kwargs, response_obj) + + def observations(self) -> tuple[CallbackObservation, ...]: + observations: Final = tuple(self._observations.get_nowait() for _ in range(self._observations.qsize())) + return tuple(sorted(observations, key=lambda observation: observation.hook)) + + INVALID_OCR_CASES: Final = ( InvalidOcrCase( name="unsupported_provider", @@ -229,6 +378,106 @@ def _execute_invalid_sdk_case( return _execute_sdk_call(call_kwargs, route, event_loop) +def _callback_call_id(route: SDKRoute, terminal: Literal["success", "failure"]) -> str: + return f"ocr-callback-{route.value}-{terminal}" + + +def _callback_metadata(route: SDKRoute, terminal: Literal["success", "failure"]) -> dict[str, str]: + return {"callback_profile": terminal, "sdk_route": route.value} + + +def _drain_callback_delivery(route: SDKRoute, event_loop: asyncio.AbstractEventLoop) -> None: + if route is SDKRoute.AOCR: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + async def drain_async_callbacks() -> None: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=CALLBACK_DRAIN_TIMEOUT_SECONDS) + await GLOBAL_LOGGING_WORKER.stop() + + event_loop.run_until_complete(drain_async_callbacks()) + + from litellm.litellm_core_utils.thread_pool_executor import executor + + executor.shutdown(wait=True, cancel_futures=False) + + +def _execute_callback_sdk_case( + case: OcrParityCase, + route: SDKRoute, + terminal: Literal["success", "failure"], + mock_url: str, + event_loop: asyncio.AbstractEventLoop, +) -> WorkerSuccess: + callback: Final = RecordingCallback() + call_kwargs: Final = { + **_call_kwargs(case.litellm_input, mock_url, route), + "callbacks": [callback], + "litellm_call_id": _callback_call_id(route, terminal), + "litellm_trace_id": _callback_call_id(route, terminal), + "metadata": _callback_metadata(route, terminal), + } + report: Final = _execute_sdk_call(call_kwargs, route, event_loop) + _drain_callback_delivery(route, event_loop) + return WorkerSuccess(report=report, callbacks=callback.observations()) + + +def _assert_callback_lifecycle( + execution: Execution, + route: SDKRoute, + terminal: Literal["success", "failure"], +) -> None: + observations: Final = execution.callbacks + assert observations is not None, f"{route.value} {terminal} callbacks were not observed" + expected_hooks: Final = ( + (f"log_{terminal}_event",) + if route is SDKRoute.OCR + else ("async_log_success_event",) + if terminal == "success" + else ("async_log_failure_event", "log_failure_event") + ) + actual_hooks: Final = tuple(observation.hook for observation in observations) + assert actual_hooks == expected_hooks, ( + f"{route.value} {terminal} expected callback hooks {expected_hooks}, received {actual_hooks}" + ) + expected_call_id: Final = _callback_call_id(route, terminal) + expected_metadata: Final = _callback_metadata(route, terminal) + for observation in observations: + assert observation.phase == terminal + assert observation.call_type == route.value + assert observation.litellm_call_id == expected_call_id + assert observation.metadata == expected_metadata + assert observation.model + if terminal == "success": + assert isinstance(execution.report, SDKSuccess) + assert observation.payload == execution.report.response + assert observation.error is None + else: + assert isinstance(execution.report, SDKError) + assert observation.payload is None + assert observation.error is not None + assert observation.error.exception_type + assert observation.error.message + assert observation.error.status_code is not None + assert observation.error.status_code >= 400 + + +def _check_callback_ocr_sdk_parity( + case: OcrParityCase, + route: SDKRoute, + terminal: Literal["success", "failure"], + case_file: Path, + runner: SubprocessRunner, +) -> None: + with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers: + python_worker, rust_worker = workers + python: Final = python_worker.execute(case_file, route.value, case.provider_responses) + rust: Final = rust_worker.execute(case_file, route.value, case.provider_responses) + + _assert_callback_lifecycle(python, route, terminal) + _assert_callback_lifecycle(rust, route, terminal) + assert_parity(python, rust, PYTHON_HTTP_SENTINEL) + + def _check_recorded_ocr_sdk_parity( ocr_fixture: OcrParityCase, route: SDKRoute, @@ -276,6 +525,25 @@ def _write_worker_case(directory: Path, index: int, case: OcrWorkerCase) -> Path return case_file +def _callback_fixture( + fixtures: tuple[OcrParityCase, ...], + terminal: Literal["success", "failure"], +) -> OcrParityCase: + matching: Final = tuple( + fixture + for fixture in fixtures + if fixture.litellm_input.contract == "mistral" + and ( + all(response.status_code < 400 for response in fixture.provider_responses) + if terminal == "success" + else any(response.status_code >= 400 for response in fixture.provider_responses) + ) + ) + if not matching: + raise AssertionError(f"no recorded Mistral OCR {terminal} fixture is available for callback parity") + return min(matching, key=lambda fixture: fixture_id(fixture.litellm_input, fixture.litellm_input.model)) + + @contextmanager def parity_checks() -> Generator[tuple[E2ECheck, ...]]: fixtures: Final = tuple( @@ -298,6 +566,17 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]: _write_worker_case(directory, len(recorded_files) + index, InvalidOcrWorkerCase(case=case)) for index, case in enumerate(INVALID_OCR_CASES) ) + callback_cases: Final[tuple[tuple[Literal["success", "failure"], OcrParityCase], ...]] = tuple( + (terminal, _callback_fixture(fixtures, terminal)) for terminal in CALLBACK_TERMINALS + ) + callback_files: Final = tuple( + _write_worker_case( + directory, + len(recorded_files) + len(invalid_files) + index, + CallbackOcrWorkerCase(case=case, terminal=terminal), + ) + for index, (terminal, case) in enumerate(callback_cases) + ) with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers: recorded: Final = tuple( E2ECheck( @@ -315,7 +594,15 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]: for case, case_file in zip(INVALID_OCR_CASES, invalid_files, strict=True) for route in SDKRoute ) - yield (*recorded, *invalid) + callbacks: Final = tuple( + E2ECheck( + f"callback:{route.value}:{terminal}", + partial(_check_callback_ocr_sdk_parity, case, route, terminal, case_file, runner), + ) + for (terminal, case), case_file in zip(callback_cases, callback_files, strict=True) + for route in SDKRoute + ) + yield (*recorded, *invalid, *callbacks) def _execute_worker_command( @@ -333,6 +620,8 @@ def _execute_worker_command( return WorkerSuccess(report=_execute_sdk_case(recorded.litellm_input, route, mock_url, event_loop)) case InvalidOcrWorkerCase(case=invalid): return WorkerSuccess(report=_execute_invalid_sdk_case(invalid, route, mock_url, event_loop)) + case CallbackOcrWorkerCase(case=callback_case, terminal=terminal): + return _execute_callback_sdk_case(callback_case, route, terminal, mock_url, event_loop) except Exception: return WorkerFailure(error=traceback.format_exc()) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py index 3e6c4060134..0e771f0dc17 100644 --- a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py @@ -230,6 +230,10 @@ _HOST_ONLY_BRIDGE_EXCLUSIONS: Final = tuple( "test_ocr_exception_type_uses_resolved_provider_context", "Python wraps bridge exceptions into public errors.", ), + ( + "test_rust_upstream_error_uses_ocr_provider_error_mapping", + "Python maps native upstream errors through the selected OCR provider config.", + ), ("test_aocr_routes_to_async_rust_when_enabled", "Python selects and invokes the async native bridge."), ("test_aocr_exception_type_uses_resolved_provider_context", "Python wraps async bridge exceptions."), ("test_ocr_forwards_timeout_to_rust", "Python converts and forwards explicit timeouts."), diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index d28184b96a6..c34833221cc 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -9,6 +9,7 @@ import httpx import pytest import litellm +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge import configuration @@ -39,6 +40,10 @@ class CapturedException(Exception): pass +class RustUpstreamError(Exception): + pass + + class RecordingBridge: """A fake ``RustOcr`` callable that records the args it was handed.""" @@ -183,6 +188,9 @@ class FakeOCRConfig: ) -> str: return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" + def get_error_class(self, error_message: str, status_code: int, headers: dict[str, str]) -> BaseLLMException: + return BaseLLMException(status_code=status_code, message=error_message, headers=headers) + def build_prepared_request( *, @@ -486,6 +494,20 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): } +def test_rust_upstream_error_uses_ocr_provider_error_mapping(): + error = RustUpstreamError(400, '{"message":"invalid model"}') + + mapped = ocr_main._map_rust_ocr_error( + error, + build_prepared_request(), + (RuntimeError, RustUpstreamError), + ) + + assert isinstance(mapped, BaseLLMException) + assert mapped.status_code == 400 + assert mapped.message == '{"message":"invalid model"}' + + def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): bridge = RecordingBridge() litellm.rust(True)