refactor(ocr-inputs): share magic-byte detector via litellm-core ocr module

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

litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs

@@ -4,6 +4,7 @@ use std::time::{Duration, Instant};
 use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
 use base64::Engine;
 use litellm_core::error::CoreError;
+use litellm_core::ocr::mime::{mime_from_file_name, sniff_document_mime};
 use litellm_core::ocr::transformation::OcrProviderConfig;
 use litellm_core::CoreResult;
 use reqwest::Url;
@@ -280,47 +281,6 @@ async fn read_response_with_limit(
     Ok(bytes)
 }
 
-fn sniff_mime_from_magic_bytes(bytes: &[u8]) -> Option<&'static str> {
-    if bytes.starts_with(b"%PDF-") {
-        return Some("application/pdf");
-    }
-    if bytes.starts_with(&[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a]) {
-        return Some("image/png");
-    }
-    if bytes.starts_with(&[0xff, 0xd8, 0xff]) {
-        return Some("image/jpeg");
-    }
-    if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
-        return Some("image/gif");
-    }
-    if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
-        return Some("image/webp");
-    }
-    if bytes.starts_with(&[0x49, 0x49, 0x2a, 0x00]) || bytes.starts_with(&[0x4d, 0x4d, 0x00, 0x2a])
-    {
-        return Some("image/tiff");
-    }
-    if bytes.starts_with(b"BM") {
-        return Some("image/bmp");
-    }
-    None
-}
-
-fn mime_from_url_extension(url_path: &str) -> Option<&'static str> {
-    let file_name = url_path.rsplit('/').next()?;
-    let extension = file_name.rsplit_once('.')?.1.to_ascii_lowercase();
-    match extension.as_str() {
-        "pdf" => Some("application/pdf"),
-        "png" => Some("image/png"),
-        "jpg" | "jpeg" => Some("image/jpeg"),
-        "gif" => Some("image/gif"),
-        "webp" => Some("image/webp"),
-        "tiff" | "tif" => Some("image/tiff"),
-        "bmp" => Some("image/bmp"),
-        _ => None,
-    }
-}
-
 fn resolve_document_mime(
     header_content_type: Option<String>,
     bytes: &[u8],
@@ -333,8 +293,8 @@ fn resolve_document_mime(
             return header.clone();
         }
     }
-    sniff_mime_from_magic_bytes(bytes)
-        .or_else(|| mime_from_url_extension(url_path))
+    sniff_document_mime(bytes)
+        .or_else(|| mime_from_file_name(url_path))
         .map(str::to_string)
         .or(header_content_type)
         .unwrap_or_else(|| "application/octet-stream".to_string())
@@ -655,52 +615,6 @@ mod tests {
         ));
     }
 
-    #[test]
-    fn sniff_mime_detects_supported_signatures() {
-        assert_eq!(
-            sniff_mime_from_magic_bytes(b"%PDF-1.7\nrest"),
-            Some("application/pdf")
-        );
-        assert_eq!(
-            sniff_mime_from_magic_bytes(b"\x89PNG\r\n\x1a\nrest"),
-            Some("image/png")
-        );
-        assert_eq!(
-            sniff_mime_from_magic_bytes(b"\xff\xd8\xff\xe0rest"),
-            Some("image/jpeg")
-        );
-        assert_eq!(
-            sniff_mime_from_magic_bytes(b"GIF89arest"),
-            Some("image/gif")
-        );
-        assert_eq!(
-            sniff_mime_from_magic_bytes(b"RIFF\x00\x00\x00\x00WEBPrest"),
-            Some("image/webp")
-        );
-        assert_eq!(
-            sniff_mime_from_magic_bytes(b"II*\x00rest"),
-            Some("image/tiff")
-        );
-        assert_eq!(sniff_mime_from_magic_bytes(b"BMrest"), Some("image/bmp"));
-        assert_eq!(sniff_mime_from_magic_bytes(b"plain text"), None);
-        assert_eq!(
-            sniff_mime_from_magic_bytes(b"RIFF\x00\x00\x00\x00WAVErest"),
-            None
-        );
-    }
-
-    #[test]
-    fn mime_from_url_extension_maps_known_suffixes() {
-        assert_eq!(
-            mime_from_url_extension("/files/report.pdf"),
-            Some("application/pdf")
-        );
-        assert_eq!(mime_from_url_extension("/a/b/c.PNG"), Some("image/png"));
-        assert_eq!(mime_from_url_extension("/scan.jpeg"), Some("image/jpeg"));
-        assert_eq!(mime_from_url_extension("/no-extension"), None);
-        assert_eq!(mime_from_url_extension("/dotless/path"), None);
-    }
-
     #[test]
     fn resolve_document_mime_prefers_specific_header() {
         assert_eq!(

litellm-rust/crates/core/src/ocr/mime.rs

@@ -0,0 +1,105 @@
+pub fn sniff_document_mime(bytes: &[u8]) -> Option<&'static str> {
+    if bytes.starts_with(b"%PDF-") {
+        return Some("application/pdf");
+    }
+    if bytes.starts_with(&[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a]) {
+        return Some("image/png");
+    }
+    if bytes.starts_with(&[0xff, 0xd8, 0xff]) {
+        return Some("image/jpeg");
+    }
+    if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
+        return Some("image/gif");
+    }
+    if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
+        return Some("image/webp");
+    }
+    if bytes.starts_with(&[0x49, 0x49, 0x2a, 0x00]) || bytes.starts_with(&[0x4d, 0x4d, 0x00, 0x2a])
+    {
+        return Some("image/tiff");
+    }
+    if bytes.starts_with(b"BM") {
+        return Some("image/bmp");
+    }
+    None
+}
+
+pub fn mime_from_file_name(file_name: &str) -> Option<&'static str> {
+    let base_name = file_name.rsplit(['/', '\\']).next()?;
+    let extension = base_name.rsplit_once('.')?.1.to_ascii_lowercase();
+    match extension.as_str() {
+        "pdf" => Some("application/pdf"),
+        "png" => Some("image/png"),
+        "jpg" | "jpeg" => Some("image/jpeg"),
+        "gif" => Some("image/gif"),
+        "webp" => Some("image/webp"),
+        "tiff" | "tif" => Some("image/tiff"),
+        "bmp" => Some("image/bmp"),
+        _ => None,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn sniff_document_mime_detects_supported_signatures() {
+        assert_eq!(
+            sniff_document_mime(b"%PDF-1.7\ncontent"),
+            Some("application/pdf")
+        );
+        assert_eq!(
+            sniff_document_mime(&[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a, 0x00]),
+            Some("image/png")
+        );
+        assert_eq!(
+            sniff_document_mime(&[0xff, 0xd8, 0xff, 0xe0]),
+            Some("image/jpeg")
+        );
+        assert_eq!(sniff_document_mime(b"GIF87a....."), Some("image/gif"));
+        assert_eq!(sniff_document_mime(b"GIF89a....."), Some("image/gif"));
+        assert_eq!(
+            sniff_document_mime(b"RIFF\x00\x00\x00\x00WEBPVP8 "),
+            Some("image/webp")
+        );
+        assert_eq!(
+            sniff_document_mime(&[0x49, 0x49, 0x2a, 0x00]),
+            Some("image/tiff")
+        );
+        assert_eq!(
+            sniff_document_mime(&[0x4d, 0x4d, 0x00, 0x2a]),
+            Some("image/tiff")
+        );
+        assert_eq!(sniff_document_mime(b"BMxxxx"), Some("image/bmp"));
+    }
+
+    #[test]
+    fn sniff_document_mime_returns_none_for_unknown_or_riff_non_webp() {
+        assert_eq!(sniff_document_mime(b"plain text payload"), None);
+        assert_eq!(sniff_document_mime(b"RIFF\x00\x00\x00\x00WAVEfmt "), None);
+        assert_eq!(sniff_document_mime(b""), None);
+        assert_eq!(sniff_document_mime(b"RIFF"), None);
+    }
+
+    #[test]
+    fn mime_from_file_name_maps_known_suffixes() {
+        assert_eq!(mime_from_file_name("report.pdf"), Some("application/pdf"));
+        assert_eq!(mime_from_file_name("/docs/scan.PNG"), Some("image/png"));
+        assert_eq!(mime_from_file_name("a/b/photo.jpeg"), Some("image/jpeg"));
+        assert_eq!(mime_from_file_name("photo.JPG"), Some("image/jpeg"));
+        assert_eq!(mime_from_file_name("anim.gif"), Some("image/gif"));
+        assert_eq!(mime_from_file_name("sticker.webp"), Some("image/webp"));
+        assert_eq!(mime_from_file_name("scan.tif"), Some("image/tiff"));
+        assert_eq!(mime_from_file_name("scan.tiff"), Some("image/tiff"));
+        assert_eq!(mime_from_file_name("bitmap.bmp"), Some("image/bmp"));
+    }
+
+    #[test]
+    fn mime_from_file_name_returns_none_when_unmapped_or_missing_extension() {
+        assert_eq!(mime_from_file_name("archive.zip"), None);
+        assert_eq!(mime_from_file_name("noextension"), None);
+        assert_eq!(mime_from_file_name("/trailing/"), None);
+        assert_eq!(mime_from_file_name(""), None);
+    }
+}

litellm-rust/crates/core/src/ocr/mod.rs

@@ -1,2 +1,3 @@
+pub mod mime;
 pub mod transformation;
 pub mod types;

litellm/ocr/main.py

@@ -488,8 +488,8 @@ def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str,
     if not file_bytes:
         raise ValueError("File is empty or could not be read")
 
-    explicit_mime = document.get("mime_type")
-    explicit_mime = explicit_mime if isinstance(explicit_mime, str) else None
+    raw_mime = document.get("mime_type")
+    explicit_mime = raw_mime if isinstance(raw_mime, str) else None
     resolved_mime = next(
         (
             candidate

tests/e2e/gateway/test_ocr_rust_e2e.py

@@ -17,6 +17,8 @@ import httpx
 import pytest
 import yaml
 
+from litellm.proxy.ocr_endpoints.endpoints import _build_document_from_upload
+
 TEST_PDF_URL = (
     "https://cdn.jsdelivr.net/gh/BerriAI/litellm"
     "@d769e81c90d453240c61fc572cdb27fae06a89d0"
@@ -28,6 +30,15 @@ TEST_IMAGE_URL = (
     "/tests/image_gen_tests/test_image.png"
 )
 
+REPO_ROOT = Path(__file__).resolve().parents[3]
+FIXTURE_PDF = REPO_ROOT / "tests" / "llm_translation" / "fixtures" / "dummy.pdf"
+FIXTURE_IMAGE = REPO_ROOT / "tests" / "image_gen_tests" / "test_image.png"
+
+RUST_OCR_UPLOAD_CASES = [
+    pytest.param(FIXTURE_PDF, "application/pdf", "document_url", id="pdf_octet_stream"),
+    pytest.param(FIXTURE_IMAGE, "image/png", "image_url", id="image_octet_stream"),
+]
+
 RUST_OCR_GATEWAY_CASES = [
     pytest.param(
         "rust-ocr-mistral",
@@ -92,6 +103,17 @@ class OcrGateway:
                 json={"model": model, "document": document},
             )
 
+    def ocr_upload(self, model: str, content: bytes, upload_name: str) -> httpx.Response:
+        with httpx.Client(
+            timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120"))
+        ) as client:
+            return client.post(
+                f"{self.base_url.rstrip('/')}/v1/ocr",
+                headers={"Authorization": f"Bearer {self.master_key}"},
+                data={"model": model},
+                files={"file": (upload_name, content, "application/octet-stream")},
+            )
+
 
 @dataclass(frozen=True)
 class OcrResources:
@@ -146,3 +168,27 @@ class TestRustOcrGateway:
 
         assert response.status_code == 200, response.text
         _assert_ocr_response_shape(response.json())
+
+    @pytest.mark.parametrize(
+        ("fixture_path", "expected_mime", "expected_document_type"),
+        RUST_OCR_UPLOAD_CASES,
+    )
+    def test_rust_ocr_octet_stream_upload_response(
+        self,
+        resources: OcrResources,
+        fixture_path: Path,
+        expected_mime: str,
+        expected_document_type: str,
+    ) -> None:
+        if not fixture_path.is_file():
+            pytest.skip(f"Missing OCR fixture: {fixture_path}")
+
+        content = fixture_path.read_bytes()
+        document = _build_document_from_upload(content, "document", "application/octet-stream")
+        assert document["type"] == expected_document_type
+        assert document[expected_document_type].startswith(f"data:{expected_mime};base64,")
+
+        response = resources.gateway.ocr_upload("rust-ocr-mistral", content, "document")
+
+        assert response.status_code == 200, response.text
+        _assert_ocr_response_shape(response.json())