fix(ocr): preserve provider error parity

litellm-rust/crates/python-bridge/src/lib.rs

@@ -72,6 +72,13 @@ fn core_error_to_pyerr(err: CoreError) -> PyErr {
     }
 }
 
+fn ocr_error_to_pyerr(err: CoreError) -> PyErr {
+    match err {
+        CoreError::Http { status, body } => RustUpstreamError::new_err((status, body)),
+        other => core_error_to_pyerr(other),
+    }
+}
+
 /// Map a core error for a route whose host keeps a Python implementation.
 ///
 /// The distinction the host needs is whether the provider was already called.
@@ -252,7 +259,7 @@ fn ocr(
 
     match result {
         Ok(value) => to_py(py, &value),
-        Err(err) => Err(core_error_to_pyerr(err)),
+        Err(err) => Err(ocr_error_to_pyerr(err)),
     }
 }
 
@@ -294,7 +301,7 @@ fn aocr(
             litellm_call_id: None,
         })
         .await
-        .map_err(core_error_to_pyerr)?;
+        .map_err(ocr_error_to_pyerr)?;
 
         Python::attach(|py| to_py(py, &value))
     })

litellm/rust_bridge/ocr.py

@@ -65,6 +65,28 @@ _rust_ocr_impl: RustOcr | None = None
 _rust_aocr_impl: RustAocr | None = None
 
 
+class _OcrProviderError(Exception):
+    def __init__(self, status_code: int, message: str, api_base: str | None) -> None:
+        super().__init__(message)
+        self.status_code: Final = status_code
+        self.response: Final = httpx.Response(
+            status_code=status_code,
+            request=httpx.Request("POST", api_base or "https://api.mistral.ai/v1/ocr"),
+        )
+
+
+def _raise_provider_error(error: BaseException, api_base: str | None) -> None:
+    from litellm.rust_bridge import get_native_bridge
+
+    native_bridge: Final = get_native_bridge()
+    upstream_error: Final = getattr(native_bridge, "RustUpstreamError", None) if native_bridge is not None else None
+    if upstream_error is None or not isinstance(error, upstream_error):
+        raise error
+    status: Final = error.args[0] if error.args else 0
+    message: Final = error.args[1] if len(error.args) > 1 else ""
+    raise _OcrProviderError(int(status) or 500, str(message), api_base) from error
+
+
 def use_litellm_rust(
     enabled: bool = True,
     *,
@@ -152,16 +174,19 @@ def ocr(
     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,
-        timeout_seconds=_timeout_to_seconds(timeout),
-    )
+    try:
+        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,
+            timeout_seconds=_timeout_to_seconds(timeout),
+        )
+    except Exception as error:
+        _raise_provider_error(error, api_base)
 
 
 async def aocr(
@@ -178,13 +203,16 @@ async def aocr(
     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,
-        timeout_seconds=_timeout_to_seconds(timeout),
-    )
+    try:
+        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,
+            timeout_seconds=_timeout_to_seconds(timeout),
+        )
+    except Exception as error:
+        _raise_provider_error(error, api_base)

tests/test_litellm/ocr/fixtures/mistral.py

@@ -90,7 +90,19 @@ class MistralOcrSdkInput(MistralCompatibleOcrSdkInput):
         return self
 
 
+class MistralProviderRejectedOcrSdkInput(MistralCompatibleOcrSdkInput):
+    boundary: str = Field(default="mistral", pattern=r"^mistral$")
+    model: Literal["mistral/invalid-ocr-model-for-parity"]
+    custom_llm_provider: Literal["mistral"] | None = None
+
+
 MISTRAL_MODEL: Final[MistralModel] = "mistral/mistral-ocr-latest"
+MISTRAL_PROVIDER_REJECTED_INPUTS: Final[tuple[MistralProviderRejectedOcrSdkInput, ...]] = (
+    MistralProviderRejectedOcrSdkInput(
+        model="mistral/invalid-ocr-model-for-parity",
+        document=pdf_document(),
+    ),
+)
 MistralFeatureLevel = Literal["2505", "2512", "4"]
 _MISTRAL_4_MODELS: Final = frozenset(
     {
@@ -256,5 +268,6 @@ def mistral_recording_targets(
                 _mistral_recording_strategy(inline_image_data_uri),
             ),
             invocation=invoke_with_api_key(client, api_key),
+            required_inputs=MISTRAL_PROVIDER_REJECTED_INPUTS,
         ),
     )

tests/test_litellm/ocr/fixtures/models.py

@@ -11,7 +11,7 @@ from tests.test_litellm.ocr.fixtures.azure import (
     AzureMistralOcrSdkInput,
 )
 from tests.test_litellm.ocr.fixtures.base import OcrSdkInputBase
-from tests.test_litellm.ocr.fixtures.mistral import MistralOcrSdkInput
+from tests.test_litellm.ocr.fixtures.mistral import MistralOcrSdkInput, MistralProviderRejectedOcrSdkInput
 from tests.test_litellm.ocr.fixtures.reducto import ReductoParseLegacySdkInput, ReductoParseV3SdkInput
 from tests.test_litellm.ocr.fixtures.vertex import VertexDeepSeekOcrSdkInput, VertexMistralOcrSdkInput
 
@@ -22,14 +22,20 @@ def _ocr_boundary(value: object) -> str | None:
     if isinstance(value, Mapping):
         mapping: Final = cast(Mapping[object, object], value)
         boundary: Final = mapping.get("boundary")
+        model: Final = mapping.get("model")
+        if boundary == "mistral" and model == "mistral/invalid-ocr-model-for-parity":
+            return "mistral_provider_rejected"
         return boundary if isinstance(boundary, str) else None
     if isinstance(value, OcrSdkInputBase):
+        if isinstance(value, MistralProviderRejectedOcrSdkInput):
+            return "mistral_provider_rejected"
         return value.boundary
     return None
 
 
 OcrSdkInput = Annotated[
     Annotated[MistralOcrSdkInput, Tag("mistral")]
+    | Annotated[MistralProviderRejectedOcrSdkInput, Tag("mistral_provider_rejected")]
     | Annotated[AzureMistralOcrSdkInput, Tag("azure_mistral")]
     | Annotated[VertexMistralOcrSdkInput, Tag("vertex_mistral")]
     | Annotated[AzureDocumentIntelligenceOcrSdkInput, Tag("azure_document_intelligence")]

tests/test_litellm/ocr/test_record_fixtures.py

@@ -19,7 +19,7 @@ from tests.test_litellm.ocr.fixtures.azure import (
 )
 from tests.test_litellm.ocr.fixtures.base import OcrSdkInputBase
 from tests.test_litellm.ocr.fixtures.common import OcrFixtureClient, OcrRecordingTarget
-from tests.test_litellm.ocr.fixtures.mistral import MISTRAL_MODELS
+from tests.test_litellm.ocr.fixtures.mistral import MISTRAL_MODELS, MISTRAL_PROVIDER_REJECTED_INPUTS
 from tests.test_litellm.ocr.fixtures.record import (
     discover_targets as discover_targets_with_media,
 )
@@ -301,7 +301,7 @@ def test_vertex_deepseek_recording_reaches_documented_image_branch() -> None:
     assert _document_transport(case_input) == ("image_url", "data")
 
 
-def test_ocr_targets_have_no_hardcoded_required_inputs() -> None:
+def test_only_intentional_provider_failures_are_fixed_inputs() -> None:
     targets: Final = discover_targets(
         {
             "MISTRAL_API_KEY": "mistral-secret",
@@ -316,7 +316,11 @@ def test_ocr_targets_have_no_hardcoded_required_inputs() -> None:
         _UNUSED_OCR_CLIENT,
     )
 
-    assert all(target.required_inputs == () for target in targets)
+    mistral: Final = next(target for target in targets if target.name == "mistral-ocr")
+    assert mistral.required_inputs == MISTRAL_PROVIDER_REJECTED_INPUTS
+    generated: Final = generate_case_inputs(mistral.strategy, examples=20)
+    assert all(case_input not in mistral.required_inputs for case_input in generated)
+    assert all(target.required_inputs == () for target in targets if target is not mistral)
 
 
 def test_mistral_adapters_preserve_omitted_optional_params() -> None:

tests/test_litellm/ocr/test_sdk_parity.py

@@ -197,6 +197,25 @@ def _execute_sdk_case(
     return _execute_sdk_call(call_kwargs, route, event_loop)
 
 
+def _execute_recorded_sdk_case(
+    sdk_input: OcrSdkInput,
+    route: SDKRoute,
+    mock_url: str,
+    event_loop: asyncio.AbstractEventLoop,
+) -> OCRResponse | SDKError:
+    import litellm
+
+    call_kwargs: Final = _call_kwargs(sdk_input, mock_url, route)
+    try:
+        if route is SDKRoute.OCR:
+            sync_route: Final = cast(Callable[..., OCRResponse], litellm.ocr)
+            return sync_route(**call_kwargs)
+        async_route: Final = cast(Callable[..., Coroutine[object, object, OCRResponse]], litellm.aocr)
+        return event_loop.run_until_complete(async_route(**call_kwargs))
+    except Exception as error:
+        return sdk_error_report(error)
+
+
 def _execute_invalid_sdk_case(
     case: InvalidOcrCase,
     route: SDKRoute,
@@ -214,17 +233,6 @@ def _execute_invalid_sdk_case(
     return _execute_sdk_call(call_kwargs, route, event_loop)
 
 
-def _call_sdk_case(sdk_input: OcrSdkInput, route: SDKRoute, mock_url: str) -> OCRResponse:
-    import litellm
-
-    call_kwargs: Final = _call_kwargs(sdk_input, mock_url, route)
-    if route is SDKRoute.OCR:
-        sync_route: Final = cast(Callable[..., OCRResponse], litellm.ocr)
-        return sync_route(**call_kwargs)
-    async_route: Final = cast(Callable[..., Coroutine[object, object, OCRResponse]], litellm.aocr)
-    return asyncio.run(async_route(**call_kwargs))
-
-
 class _RustOcrSpy:
     def __init__(self, delegate: RustOcr) -> None:
         self.delegate: Final = delegate
@@ -333,27 +341,37 @@ def test_recorded_ocr_sdk_parity(
     route: SDKRoute,
 ) -> None:
     sync_spy, async_spy = _native_spies()
-    with _restore_rust_ocr_state(), replay_server() as provider:
-        rust_ocr_bridge.use_litellm_rust(False, ocr=sync_spy, aocr=async_spy)
-        python: Final = run_in_process(
-            provider,
-            ocr_fixture.provider_responses,
-            lambda mock_url: _call_sdk_case(ocr_fixture.litellm_input, route, mock_url),
-        )
-        assert sync_spy.calls == 0
-        assert async_spy.calls == 0
+    event_loop: Final = asyncio.new_event_loop()
+    try:
+        with _restore_rust_ocr_state(), replay_server() as provider:
+            rust_ocr_bridge.use_litellm_rust(False, ocr=sync_spy, aocr=async_spy)
+            python: Final = run_in_process(
+                provider,
+                ocr_fixture.provider_responses,
+                lambda mock_url: _execute_recorded_sdk_case(ocr_fixture.litellm_input, route, mock_url, event_loop),
+            )
+            assert sync_spy.calls == 0
+            assert async_spy.calls == 0
 
-        rust_ocr_bridge.use_litellm_rust(True, ocr=sync_spy, aocr=async_spy)
-        rust: Final = run_in_process(
-            provider,
-            ocr_fixture.provider_responses,
-            lambda mock_url: _call_sdk_case(ocr_fixture.litellm_input, route, mock_url),
-        )
+            rust_ocr_bridge.use_litellm_rust(True, ocr=sync_spy, aocr=async_spy)
+            rust: Final = run_in_process(
+                provider,
+                ocr_fixture.provider_responses,
+                lambda mock_url: _execute_recorded_sdk_case(ocr_fixture.litellm_input, route, mock_url, event_loop),
+            )
+    finally:
+        event_loop.close()
 
     assert sync_spy.calls == (1 if route is SDKRoute.OCR else 0)
     assert async_spy.calls == (1 if route is SDKRoute.AOCR else 0)
     assert_request_parity(python.requests, rust.requests)
-    assert_model_parity(python.response, rust.response)
+    if any(response.status_code >= 400 for response in ocr_fixture.provider_responses):
+        assert isinstance(python.response, SDKError)
+    if isinstance(python.response, SDKError):
+        assert python.response == rust.response
+    else:
+        assert isinstance(rust.response, OCRResponse)
+        assert_model_parity(python.response, rust.response)
 
 
 @pytest.mark.parametrize("case", INVALID_OCR_CASES, ids=tuple(case.name for case in INVALID_OCR_CASES))