diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 479f184dcca..8bd87ae21d3 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -53,7 +53,6 @@ class _PreparedOCRRequest: class _PreparedRustOCRCall: api_key: Optional[str] headers: dict[str, object] - complete_url: str def _timeout_to_seconds( @@ -211,7 +210,6 @@ def _prepare_rust_ocr_call( return _PreparedRustOCRCall( api_key=resolved_api_key, headers=cast(dict[str, object], resolved_headers), - complete_url=resolved_complete_url, ) @@ -239,7 +237,7 @@ def _run_rust_ocr( api_key=prepared.api_key, api_base=prepared_request.api_base, custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared_request.extra_headers, + extra_headers=prepared.headers, optional_params=prepared_request.optional_params, timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout), ) @@ -262,7 +260,7 @@ async def _run_rust_aocr( api_key=prepared.api_key, api_base=prepared_request.api_base, custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared_request.extra_headers, + extra_headers=prepared.headers, optional_params=prepared_request.optional_params, timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout), ) @@ -350,6 +348,9 @@ async def aocr( extra_headers=extra_headers, kwargs=kwargs, ) + model = prepared.model + custom_llm_provider = prepared.custom_llm_provider + local_vars.update({"model": model, "custom_llm_provider": custom_llm_provider}) if prepared.custom_llm_provider == "mistral" and rust_ocr_enabled(): rust_aocr = load_rust_aocr() @@ -487,6 +488,9 @@ def ocr( extra_headers=extra_headers, timeout=timeout, ) + model = prepared.model + custom_llm_provider = prepared.custom_llm_provider + local_vars.update({"model": model, "custom_llm_provider": custom_llm_provider}) # Optional Rust path: hand the whole Mistral OCR call to the Rust bridge. if prepared.custom_llm_provider == "mistral" and rust_ocr_enabled(): diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 31a346c09c3..aad5a4d6f42 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -32,6 +32,10 @@ FAKE_OCR_RESPONSE: dict[str, object] = { } +class CapturedException(Exception): + pass + + class RecordingBridge: """A fake ``RustOcr`` callable that records the args it was handed.""" @@ -96,6 +100,36 @@ class RecordingAsyncBridge: return dict(FAKE_OCR_RESPONSE) +class RaisingBridge: + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + raise RuntimeError("bridge failed") + + +class RaisingAsyncBridge: + async def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + raise RuntimeError("bridge failed") + + class RecordingLogging: """A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``.""" @@ -302,7 +336,10 @@ def test_run_rust_ocr_forwards_args_and_wraps_response(): "api_key": "sk-test", "api_base": "https://proxy.internal", "custom_llm_provider": "mistral", - "extra_headers": {"x-trace-id": "trace-1"}, + "extra_headers": { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + }, "optional_params": {"include_image_base64": True}, "timeout_seconds": 12.5, } @@ -413,11 +450,33 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" assert call["custom_llm_provider"] == "mistral" - assert call["extra_headers"] == {"x-trace-id": "trace-1"} + assert call["extra_headers"] == { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + } # Raw OCR params ride along in optional_params; Rust filters to supported keys. assert call["optional_params"].get("include_image_base64") is True +def test_ocr_exception_type_uses_resolved_provider_context( + monkeypatch: pytest.MonkeyPatch, +): + captured: dict[str, object] = {} + + def fake_exception_type(**kwargs: object) -> CapturedException: + captured.update(kwargs) + return CapturedException("wrapped") + + monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) + litellm.use_litellm_rust(True, ocr=RaisingBridge()) + + with pytest.raises(CapturedException): + litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert captured["model"] == "mistral-ocr-latest" + assert captured["custom_llm_provider"] == "mistral" + + @pytest.mark.asyncio async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge): response = await litellm.aocr( @@ -436,10 +495,33 @@ async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge): assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" assert call["custom_llm_provider"] == "mistral" - assert call["extra_headers"] == {"x-trace-id": "trace-1"} + assert call["extra_headers"] == { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + } assert call["optional_params"].get("include_image_base64") is True +@pytest.mark.asyncio +async def test_aocr_exception_type_uses_resolved_provider_context( + monkeypatch: pytest.MonkeyPatch, +): + captured: dict[str, object] = {} + + def fake_exception_type(**kwargs: object) -> CapturedException: + captured.update(kwargs) + return CapturedException("wrapped") + + monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) + litellm.use_litellm_rust(True, aocr=RaisingAsyncBridge()) + + with pytest.raises(CapturedException): + await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert captured["model"] == "mistral-ocr-latest" + assert captured["custom_llm_provider"] == "mistral" + + def test_ocr_forwards_timeout_to_rust(fake_bridge): """Caller-supplied timeout must flow into the Rust bridge so the fixed 600s client ceiling doesn't silently override shorter deadlines."""