fix(ocr): support latest staging toolchain

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

litellm-rust/Cargo.lock

@@ -28,6 +28,12 @@ version = "1.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
 
+[[package]]
+name = "autocfg"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
+
 [[package]]
 name = "axum"
 version = "0.7.9"
@@ -386,6 +392,20 @@ dependencies = [
  "percent-encoding",
 ]
 
+[[package]]
+name = "futures"
+version = "0.3.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "futures-io",
+ "futures-sink",
+ "futures-task",
+ "futures-util",
+]
+
 [[package]]
 name = "futures-channel"
 version = "0.3.32"

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

@@ -172,7 +172,7 @@ fn aocr(
             litellm_call_id: None,
         })
         .await
-        .map_err(|err| Python::with_gil(|py| core_error_to_pyerr(py, err)))?;
+        .map_err(|err| Python::attach(|py| core_error_to_pyerr(py, err)))?;
 
         Python::attach(|py| json_to_py(py, value))
     })

litellm/ocr/main.py

@@ -222,7 +222,7 @@ def _resolve_ocr_call_context(
 def _run_pre_call_logging(
     litellm_logging_obj: LiteLLMLoggingObj,
     model: str,
-    document: dict[str, Any],
+    document: dict[str, object],
     api_key: str | None,
     api_base: str | None,
     extra_headers: dict[str, object] | None,
@@ -246,7 +246,7 @@ def _run_pre_call_logging(
 def _run_rust_ocr(
     rust_ocr: RustOcr,
     model: str,
-    document: dict[str, Any],
+    document: dict[str, object],
     api_key: str | None,
     api_base: str | None,
     custom_llm_provider: str,
@@ -267,7 +267,7 @@ def _run_rust_ocr(
     return OCRResponse.model_validate(
         rust_ocr(
             model=model,
-            document=cast(dict[str, object], document),
+            document=document,
             api_key=api_key,
             api_base=api_base,
             custom_llm_provider=custom_llm_provider,
@@ -287,7 +287,7 @@ def _missing_rust_bridge_error() -> RuntimeError:
 async def _run_rust_aocr(
     rust_aocr: RustAocr,
     model: str,
-    document: dict[str, Any],
+    document: dict[str, object],
     api_key: str | None,
     api_base: str | None,
     custom_llm_provider: str,
@@ -308,7 +308,7 @@ async def _run_rust_aocr(
     return OCRResponse.model_validate(
         await rust_aocr(
             model=model,
-            document=cast(dict[str, object], document),
+            document=document,
             api_key=api_key,
             api_base=api_base,
             custom_llm_provider=custom_llm_provider,

tests/e2e/gateway/ocr_capture_proxy.py

@@ -154,19 +154,13 @@ def capture_proxy(tmp_path: Path) -> Iterator[CaptureProxy]:
     proxy_log_path = tmp_path / "capture-proxy.log"
     try:
         capture_port = _free_port()
-        capture_server = HTTPServer(
-            ("127.0.0.1", capture_port), _make_capture_handler(captures)
-        )
-        server_thread = threading.Thread(
-            target=capture_server.serve_forever, daemon=True
-        )
+        capture_server = HTTPServer(("127.0.0.1", capture_port), _make_capture_handler(captures))
+        server_thread = threading.Thread(target=capture_server.serve_forever, daemon=True)
         server_thread.start()
 
         proxy_port = _free_port()
         config_path = tmp_path / "capture-config.yml"
-        config_path.write_text(
-            yaml.safe_dump(_capture_config(capture_port, master_key).model_dump())
-        )
+        config_path.write_text(yaml.safe_dump(_capture_config(capture_port, master_key).model_dump()))
 
         proxy_log = proxy_log_path.open("w")
         proxy = subprocess.Popen(
@@ -187,18 +181,14 @@ def capture_proxy(tmp_path: Path) -> Iterator[CaptureProxy]:
             stderr=subprocess.STDOUT,
         )
         proxy_url = f"http://127.0.0.1:{proxy_port}"
-        if not _wait_for_liveness(
-            proxy_url, time.monotonic() + _LIVENESS_DEADLINE_SECONDS
-        ):
+        if not _wait_for_liveness(proxy_url, time.monotonic() + _LIVENESS_DEADLINE_SECONDS):
             proxy_log.flush()
             tail = _sanitize(proxy_log_path.read_text()[-4000:])
             pytest.fail(
                 f"capture proxy did not become live while the Rust bridge is available; "
                 f"sanitized proxy log at {proxy_log_path}\n{tail}"
             )
-        yield CaptureProxy(
-            proxy_url=proxy_url, master_key=master_key, captures=captures
-        )
+        yield CaptureProxy(proxy_url=proxy_url, master_key=master_key, captures=captures)
     finally:
         if proxy is not None:
             proxy.terminate()

tests/e2e/gateway/test_ocr_rust_e2e.py

@@ -158,6 +158,7 @@ class GatewayConfig(BaseModel):
 
     model_list: tuple[GatewayConfigEntry, ...]
 
+
 TEST_PDF_URL = (
     "https://cdn.jsdelivr.net/gh/BerriAI/litellm"
     "@d769e81c90d453240c61fc572cdb27fae06a89d0"
@@ -282,16 +283,8 @@ def _wire_payload(
         {
             "model": model,
             "document": document.model_dump(mode="json", exclude_none=True),
-            **(
-                params.model_dump(mode="json", by_alias=True)
-                if params is not None
-                else {}
-            ),
-            **(
-                canaries.model_dump(mode="json", by_alias=True)
-                if canaries is not None
-                else {}
-            ),
+            **(params.model_dump(mode="json", by_alias=True) if params is not None else {}),
+            **(canaries.model_dump(mode="json", by_alias=True) if canaries is not None else {}),
         }
     )
 
@@ -330,9 +323,7 @@ class OcrGateway:
                 content=_wire_payload(model, document, params),
             )
 
-    def create_model(
-        self, model_name: str, litellm_params: dict[str, str]
-    ) -> httpx.Response:
+    def create_model(self, model_name: str, litellm_params: dict[str, str]) -> httpx.Response:
         with self._client() as client:
             return client.post(
                 f"{self.base_url.rstrip('/')}/model/new",
@@ -366,18 +357,14 @@ class OcrGateway:
             if model_name in self.model_names():
                 return
             time.sleep(1)
-        raise AssertionError(
-            f"{model_name} did not appear on /model/info within {attempts}s"
-        )
+        raise AssertionError(f"{model_name} did not appear on /model/info within {attempts}s")
 
     def wait_for_model_absent(self, model_name: str, attempts: int = 20) -> None:
         for _ in range(attempts):
             if model_name not in self.model_names():
                 return
             time.sleep(1)
-        raise AssertionError(
-            f"{model_name} still present on /model/info after {attempts}s"
-        )
+        raise AssertionError(f"{model_name} still present on /model/info after {attempts}s")
 
 
 @dataclass(frozen=True)
@@ -389,9 +376,7 @@ class OcrResources:
 def resources() -> OcrResources:
     proxy_url = os.getenv("LITELLM_PROXY_URL")
     if not proxy_url:
-        pytest.skip(
-            "Start a Rust OCR proxy and set LITELLM_PROXY_URL, e.g. http://localhost:4000"
-        )
+        pytest.skip("Start a Rust OCR proxy and set LITELLM_PROXY_URL, e.g. http://localhost:4000")
     return OcrResources(
         gateway=OcrGateway(
             base_url=proxy_url,
@@ -402,39 +387,29 @@ def resources() -> OcrResources:
 
 class TestRustOcrGateway:
     def test_rust_ocr_models_are_on_gateway_config(self) -> None:
-        config = GatewayConfig.model_validate(
-            cast(object, yaml.safe_load(CONFIG_PATH.read_text()))
-        )
+        config = GatewayConfig.model_validate(cast(object, yaml.safe_load(CONFIG_PATH.read_text())))
         configured_models = frozenset(entry.model_name for entry in config.model_list)
 
         expected_models = {str(case.values[0]) for case in RUST_OCR_GATEWAY_CASES}
         assert expected_models.issubset(configured_models)
 
-    def test_running_gateway_loaded_rust_ocr_models(
-        self, resources: OcrResources
-    ) -> None:
+    def test_running_gateway_loaded_rust_ocr_models(self, resources: OcrResources) -> None:
         expected_models = {str(case.values[0]) for case in RUST_OCR_GATEWAY_CASES}
         assert expected_models.issubset(resources.gateway.model_names())
 
     @pytest.mark.parametrize(("model", "document"), RUST_OCR_GATEWAY_CASES)
-    def test_rust_ocr_model_gateway_response(
-        self, resources: OcrResources, model: str, document: OcrDocument
-    ) -> None:
+    def test_rust_ocr_model_gateway_response(self, resources: OcrResources, model: str, document: OcrDocument) -> None:
         response = resources.gateway.ocr(model, document)
 
         assert response.status_code == 200, response.text
         OcrResponseEnvelope.model_validate_json(response.content)
 
     @pytest.mark.e2e
-    def test_rust_ocr_mistral_live_forwards_supported_params(
-        self, resources: OcrResources
-    ) -> None:
+    def test_rust_ocr_mistral_live_forwards_supported_params(self, resources: OcrResources) -> None:
         if not os.getenv("MISTRAL_API_KEY"):
             pytest.skip("MISTRAL_API_KEY not set for live Mistral OCR call")
 
-        response = resources.gateway.ocr(
-            "rust-ocr-mistral", CAPTURE_DOCUMENT, SUPPORTED_PARAMS
-        )
+        response = resources.gateway.ocr("rust-ocr-mistral", CAPTURE_DOCUMENT, SUPPORTED_PARAMS)
 
         assert response.status_code == 200, response.text
         parsed = OcrResponseEnvelope.model_validate_json(response.content)
@@ -463,16 +438,12 @@ def test_rust_ocr_proxy_forwards_full_contract_to_capture_endpoint(
     assert response.status_code == 200, response.text
     OcrResponseEnvelope.model_validate_json(response.content)
 
-    captured = MistralOcrUpstreamRequest.model_validate_json(
-        capture_proxy.captures.get(timeout=10)
-    )
+    captured = MistralOcrUpstreamRequest.model_validate_json(capture_proxy.captures.get(timeout=10))
     assert captured == EXPECTED_UPSTREAM
 
 
 class TestRustOcrDynamicDeployment:
-    def test_os_environ_api_key_deployment_lifecycle(
-        self, resources: OcrResources
-    ) -> None:
+    def test_os_environ_api_key_deployment_lifecycle(self, resources: OcrResources) -> None:
         if not os.getenv("MISTRAL_API_KEY"):
             pytest.skip("Set MISTRAL_API_KEY on the proxy for the live OCR lifecycle")
 

tests/ocr_tests/test_ocr_azure_document_intelligence.py

@@ -42,4 +42,3 @@ class TestAzureDocumentIntelligenceOCR(BaseOCRTest):
             "api_key": api_key,
             "api_base": endpoint,
         }
-

tests/ocr_tests/test_ocr_vertex_ai.py

@@ -64,9 +64,7 @@ class TestVertexAIMistralOCR(BaseOCRTest):
         if os.environ.get("LITELLM_RUN_LIVE_VERTEX_MISTRAL_OCR_TESTS") != "1":
             pytest.skip("Live Vertex AI Mistral OCR E2E tests are opt-in")
         if os.environ.get("CASSETTE_REDIS_URL"):
-            pytest.skip(
-                "Live Vertex AI Mistral OCR E2E tests cannot run under VCR replay"
-            )
+            pytest.skip("Live Vertex AI Mistral OCR E2E tests cannot run under VCR replay")
 
     def get_base_ocr_call_args(self) -> dict:
         """

tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py

@@ -29,25 +29,16 @@ class TestDocIntelligenceApiBaseResolution:
     def test_generic_azure_ai_base_does_not_hijack_doc_intelligence(self, monkeypatch):
         monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE)
 
-        assert (
-            _resolved_api_base("azure_ai/doc-intelligence/prebuilt-layout", None)
-            is None
-        )
+        assert _resolved_api_base("azure_ai/doc-intelligence/prebuilt-layout", None) is None
 
     def test_explicit_api_base_is_honoured_for_doc_intelligence(self, monkeypatch):
         monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE)
 
         custom = "https://my-di.cognitiveservices.azure.com"
 
-        assert (
-            _resolved_api_base("azure_ai/doc-intelligence/prebuilt-layout", custom)
-            == custom
-        )
+        assert _resolved_api_base("azure_ai/doc-intelligence/prebuilt-layout", custom) == custom
 
     def test_generic_azure_ai_base_still_applies_to_mistral_ocr(self, monkeypatch):
         monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE)
 
-        assert (
-            _resolved_api_base("azure_ai/mistral-document-ai-2505", None)
-            == _AZURE_AI_API_BASE
-        )
+        assert _resolved_api_base("azure_ai/mistral-document-ai-2505", None) == _AZURE_AI_API_BASE

tests/test_litellm/ocr/test_ocr_file_input.py

@@ -71,9 +71,7 @@ class TestConvertFileDocumentToUrlDocument:
             tmp_path = Path(f.name)
 
         try:
-            result = convert_file_document_to_url_document(
-                {"type": "file", "file": tmp_path}
-            )
+            result = convert_file_document_to_url_document({"type": "file", "file": tmp_path})
 
             assert result["type"] == "document_url"
             assert result["document_url"].startswith("data:application/pdf;base64,")
@@ -93,9 +91,7 @@ class TestConvertFileDocumentToUrlDocument:
             tmp_path = Path(f.name)
 
         try:
-            result = convert_file_document_to_url_document(
-                {"type": "file", "file": tmp_path}
-            )
+            result = convert_file_document_to_url_document({"type": "file", "file": tmp_path})
 
             assert result["type"] == "image_url"
             assert result["image_url"].startswith("data:image/png;base64,")
@@ -110,9 +106,7 @@ class TestConvertFileDocumentToUrlDocument:
         request handler the value is attacker-controlled, and opening it as
         a path is an arbitrary local file read on the proxy host."""
         with pytest.raises(ValueError, match="does not accept bare str values"):
-            convert_file_document_to_url_document(
-                {"type": "file", "file": "/etc/passwd"}
-            )
+            convert_file_document_to_url_document({"type": "file", "file": "/etc/passwd"})
 
     def test_should_convert_pathlib_path(self):
         """pathlib.Path objects should work the same as string paths."""
@@ -124,9 +118,7 @@ class TestConvertFileDocumentToUrlDocument:
             tmp_path = Path(f.name)
 
         try:
-            result = convert_file_document_to_url_document(
-                {"type": "file", "file": tmp_path}
-            )
+            result = convert_file_document_to_url_document({"type": "file", "file": tmp_path})
 
             assert result["type"] == "document_url"
             assert result["document_url"].startswith("data:application/pdf;base64,")
@@ -137,9 +129,7 @@ class TestConvertFileDocumentToUrlDocument:
         """Raw bytes should be converted using a fallback MIME type."""
         content = b"raw bytes content"
 
-        result = convert_file_document_to_url_document(
-            {"type": "file", "file": content}
-        )
+        result = convert_file_document_to_url_document({"type": "file", "file": content})
 
         assert result["type"] == "document_url"
         assert "base64," in result["document_url"]
@@ -162,9 +152,7 @@ class TestConvertFileDocumentToUrlDocument:
         """Raw bytes with an image MIME type should produce type=image_url."""
         content = b"raw image content"
 
-        result = convert_file_document_to_url_document(
-            {"type": "file", "file": content, "mime_type": "image/jpeg"}
-        )
+        result = convert_file_document_to_url_document({"type": "file", "file": content, "mime_type": "image/jpeg"})
 
         assert result["type"] == "image_url"
         assert result["image_url"].startswith("data:image/jpeg;base64,")
@@ -174,9 +162,7 @@ class TestConvertFileDocumentToUrlDocument:
         content = b"file-like content"
         file_obj = BytesIO(content)
 
-        result = convert_file_document_to_url_document(
-            {"type": "file", "file": file_obj}
-        )
+        result = convert_file_document_to_url_document({"type": "file", "file": file_obj})
 
         assert result["type"] == "document_url"
         assert "base64," in result["document_url"]
@@ -187,9 +173,7 @@ class TestConvertFileDocumentToUrlDocument:
         file_obj = BytesIO(content)
         file_obj.name = "test_image.png"
 
-        result = convert_file_document_to_url_document(
-            {"type": "file", "file": file_obj}
-        )
+        result = convert_file_document_to_url_document({"type": "file", "file": file_obj})
 
         assert result["type"] == "image_url"
         assert result["image_url"].startswith("data:image/png;base64,")
@@ -202,9 +186,7 @@ class TestConvertFileDocumentToUrlDocument:
     def test_should_raise_error_for_nonexistent_pathlib_path(self):
         """Non-existent pathlib.Path should raise a path-free input error."""
         with pytest.raises(ValueError, match="does not exist") as exc_info:
-            convert_file_document_to_url_document(
-                {"type": "file", "file": Path("/nonexistent/path/to/file.pdf")}
-            )
+            convert_file_document_to_url_document({"type": "file", "file": Path("/nonexistent/path/to/file.pdf")})
 
         assert "/nonexistent/path/to/file.pdf" not in str(exc_info.value)
 
@@ -215,9 +197,7 @@ class TestConvertFileDocumentToUrlDocument:
 
         try:
             with pytest.raises(ValueError, match="File is empty"):
-                convert_file_document_to_url_document(
-                    {"type": "file", "file": tmp_path}
-                )
+                convert_file_document_to_url_document({"type": "file", "file": tmp_path})
         finally:
             os.unlink(str(tmp_path))
 
@@ -248,9 +228,7 @@ class TestConvertFileDocumentToUrlDocument:
             tmp_path = Path(f.name)
 
         try:
-            result = convert_file_document_to_url_document(
-                {"type": "file", "file": tmp_path, "mime_type": "image/png"}
-            )
+            result = convert_file_document_to_url_document({"type": "file", "file": tmp_path, "mime_type": "image/png"})
 
             assert result["type"] == "image_url"
             assert result["image_url"].startswith("data:image/png;base64,")
@@ -477,7 +455,5 @@ class TestProxySecurityGuard:
         result = await self._parse_multipart(mock_request)
 
         assert result["document"]["type"] == "document_url"
-        assert result["document"]["document_url"].startswith(
-            "data:application/pdf;base64,"
-        )
+        assert result["document"]["document_url"].startswith("data:application/pdf;base64,")
         assert result["model"] == "mistral/mistral-ocr-latest"

tests/test_litellm/ocr/test_rust_bridge.py

@@ -633,9 +633,7 @@ RUST_OCR_ERROR_CASES = [
     pytest.param(403, litellm.PermissionDeniedError, 403, id="403_permission_denied"),
     pytest.param(404, litellm.NotFoundError, 404, id="404_not_found"),
     pytest.param(408, litellm.Timeout, 408, id="408_timeout"),
-    pytest.param(
-        422, litellm.UnprocessableEntityError, 422, id="422_unprocessable_entity"
-    ),
+    pytest.param(422, litellm.UnprocessableEntityError, 422, id="422_unprocessable_entity"),
     pytest.param(429, litellm.RateLimitError, 429, id="429_rate_limit"),
     pytest.param(500, litellm.InternalServerError, 500, id="500_internal"),
     pytest.param(502, litellm.BadGatewayError, 502, id="502_bad_gateway"),
@@ -644,9 +642,7 @@ RUST_OCR_ERROR_CASES = [
 ]
 
 
-@pytest.mark.parametrize(
-    ("status_code", "expected_exception", "expected_status"), RUST_OCR_ERROR_CASES
-)
+@pytest.mark.parametrize(("status_code", "expected_exception", "expected_status"), RUST_OCR_ERROR_CASES)
 def test_rust_ocr_error_maps_to_public_exception(
     status_code: int | None,
     expected_exception: type[Exception],
@@ -666,9 +662,7 @@ def test_rust_ocr_error_maps_to_public_exception(
     assert "upstream boom" in str(exc)
 
 
-@pytest.mark.parametrize(
-    ("status_code", "expected_exception", "expected_status"), RUST_OCR_ERROR_CASES
-)
+@pytest.mark.parametrize(("status_code", "expected_exception", "expected_status"), RUST_OCR_ERROR_CASES)
 def test_ocr_raises_typed_exception_from_rust_error(
     status_code: int | None,
     expected_exception: type[Exception],
@@ -683,18 +677,14 @@ def test_ocr_raises_typed_exception_from_rust_error(
     assert exc_info.value.llm_provider == "mistral"
 
 
-@pytest.mark.parametrize(
-    ("status_code", "expected_exception", "expected_status"), RUST_OCR_ERROR_CASES
-)
+@pytest.mark.parametrize(("status_code", "expected_exception", "expected_status"), RUST_OCR_ERROR_CASES)
 @pytest.mark.asyncio
 async def test_aocr_raises_typed_exception_from_rust_error(
     status_code: int | None,
     expected_exception: type[Exception],
     expected_status: int,
 ) -> None:
-    rust_bridge._set_rust_ocr_bridge(
-        aocr=RustErrorAsyncBridge("upstream boom", status_code)
-    )
+    rust_bridge._set_rust_ocr_bridge(aocr=RustErrorAsyncBridge("upstream boom", status_code))
 
     with pytest.raises(expected_exception) as exc_info:
         await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
@@ -885,9 +875,7 @@ def test_raise_ocr_exception_keeps_validation_error_off_bad_request(
 def test_ocr_forwards_os_environ_api_key_reference_to_rust(
     fake_bridge: RecordingBridge,
 ) -> None:
-    litellm.ocr(
-        model=MODEL, document=DOCUMENT, api_key="os.environ/MISTRAL_OCR_TEST_KEY"
-    )
+    litellm.ocr(model=MODEL, document=DOCUMENT, api_key="os.environ/MISTRAL_OCR_TEST_KEY")
 
     assert fake_bridge.calls[0]["api_key"] == "os.environ/MISTRAL_OCR_TEST_KEY"
 
@@ -922,8 +910,6 @@ def test_ocr_forwards_provider_derived_os_environ_references_to_rust(
 async def test_aocr_forwards_os_environ_api_key_reference_to_rust(
     fake_async_bridge: RecordingAsyncBridge,
 ) -> None:
-    await litellm.aocr(
-        model=MODEL, document=DOCUMENT, api_key="os.environ/MISTRAL_OCR_TEST_KEY"
-    )
+    await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="os.environ/MISTRAL_OCR_TEST_KEY")
 
     assert fake_async_bridge.calls[0]["api_key"] == "os.environ/MISTRAL_OCR_TEST_KEY"