fix(ocr-inputs): restore MIME and public-URL conversion parity

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

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

@@ -280,6 +280,66 @@ 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],
+    url_path: &str,
+) -> String {
+    if let Some(header) = &header_content_type {
+        if !header.eq_ignore_ascii_case("application/octet-stream")
+            && !header.eq_ignore_ascii_case("binary/octet-stream")
+        {
+            return header.clone();
+        }
+    }
+    sniff_mime_from_magic_bytes(bytes)
+        .or_else(|| mime_from_url_extension(url_path))
+        .map(str::to_string)
+        .or(header_content_type)
+        .unwrap_or_else(|| "application/octet-stream".to_string())
+}
+
 pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult<Value> {
     let Some((field, url)) = document_url_field(&document)? else {
         return Ok(document);
@@ -297,16 +357,16 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes
             body: truncate_error_body(&body),
         });
     }
-    let content_type = response
+    let header_content_type = response
         .headers()
         .get(reqwest::header::CONTENT_TYPE)
         .and_then(|value| value.to_str().ok())
         .and_then(|value| value.split(';').next())
         .map(str::trim)
         .filter(|value| !value.is_empty())
-        .unwrap_or("application/octet-stream")
-        .to_string();
+        .map(str::to_string);
     let bytes = read_response_with_limit(response, &final_url).await?;
+    let content_type = resolve_document_mime(header_content_type, &bytes, final_url.path());
     let data_uri = format!(
         "data:{content_type};base64,{}",
         BASE64_STANDARD.encode(bytes)
@@ -595,6 +655,96 @@ 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!(
+            resolve_document_mime(Some("image/png".to_string()), b"%PDF-1.4", "/x.pdf"),
+            "image/png"
+        );
+    }
+
+    #[test]
+    fn resolve_document_mime_sniffs_when_header_is_octet_stream() {
+        assert_eq!(
+            resolve_document_mime(
+                Some("application/octet-stream".to_string()),
+                b"%PDF-1.4 payload",
+                "/x"
+            ),
+            "application/pdf"
+        );
+    }
+
+    #[test]
+    fn resolve_document_mime_uses_url_extension_when_header_and_bytes_ambiguous() {
+        assert_eq!(
+            resolve_document_mime(None, b"unrecognized bytes", "/docs/file.pdf"),
+            "application/pdf"
+        );
+    }
+
+    #[test]
+    fn resolve_document_mime_falls_back_to_octet_stream() {
+        assert_eq!(
+            resolve_document_mime(None, b"unrecognized bytes", "/docs/file"),
+            "application/octet-stream"
+        );
+        assert_eq!(
+            resolve_document_mime(
+                Some("application/octet-stream".to_string()),
+                b"unrecognized bytes",
+                "/docs/file"
+            ),
+            "application/octet-stream"
+        );
+    }
+
     #[tokio::test]
     async fn convert_document_url_leaves_data_uri_untouched() {
         let document = json!({

litellm/ocr/main.py

@@ -401,6 +401,30 @@ def get_mime_type(file_path: str) -> str:
     return guessed or "application/octet-stream"
 
 
+def sniff_mime_type_from_bytes(data: bytes) -> str | None:
+    """
+    Detect a document/image MIME type from a byte payload's magic-number prefix.
+
+    Returns None when the prefix matches none of the OCR-supported signatures,
+    leaving the caller to fall back to an explicit MIME type or octet-stream.
+    """
+    if data.startswith(b"%PDF-"):
+        return "application/pdf"
+    if data.startswith(b"\x89PNG\r\n\x1a\n"):
+        return "image/png"
+    if data.startswith(b"\xff\xd8\xff"):
+        return "image/jpeg"
+    if data.startswith((b"GIF87a", b"GIF89a")):
+        return "image/gif"
+    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
+        return "image/webp"
+    if data.startswith((b"II*\x00", b"MM\x00*")):
+        return "image/tiff"
+    if data.startswith(b"BM"):
+        return "image/bmp"
+    return None
+
+
 def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, str]:
     """
     Convert a file-type document dict to a document_url-type document dict
@@ -470,6 +494,9 @@ 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")
 
+    if mime_type == "application/octet-stream":
+        mime_type = sniff_mime_type_from_bytes(file_bytes) or mime_type
+
     if "mime_type" in document:
         mime_type = document["mime_type"]
 

tests/test_litellm/ocr/test_ocr_file_input.py

@@ -19,7 +19,11 @@ from unittest.mock import AsyncMock, MagicMock
 import orjson
 import pytest
 
-from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type
+from litellm.ocr.main import (
+    convert_file_document_to_url_document,
+    get_mime_type,
+    sniff_mime_type_from_bytes,
+)
 
 
 class TestGetMimeType:
@@ -59,6 +63,36 @@ class TestGetMimeType:
         assert isinstance(result, str)
 
 
+class TestSniffMimeTypeFromBytes:
+    def test_should_detect_pdf_from_magic_bytes(self):
+        assert sniff_mime_type_from_bytes(b"%PDF-1.7\n%rest") == "application/pdf"
+
+    def test_should_detect_png_from_magic_bytes(self):
+        assert sniff_mime_type_from_bytes(b"\x89PNG\r\n\x1a\n rest") == "image/png"
+
+    def test_should_detect_jpeg_from_magic_bytes(self):
+        assert sniff_mime_type_from_bytes(b"\xff\xd8\xff\xe0 rest") == "image/jpeg"
+
+    def test_should_detect_gif_from_magic_bytes(self):
+        assert sniff_mime_type_from_bytes(b"GIF89a rest") == "image/gif"
+
+    def test_should_detect_webp_from_magic_bytes(self):
+        assert sniff_mime_type_from_bytes(b"RIFF\x00\x00\x00\x00WEBPrest") == "image/webp"
+
+    def test_should_detect_tiff_from_magic_bytes(self):
+        assert sniff_mime_type_from_bytes(b"II*\x00 rest") == "image/tiff"
+        assert sniff_mime_type_from_bytes(b"MM\x00* rest") == "image/tiff"
+
+    def test_should_detect_bmp_from_magic_bytes(self):
+        assert sniff_mime_type_from_bytes(b"BM rest") == "image/bmp"
+
+    def test_should_return_none_for_unknown_prefix(self):
+        assert sniff_mime_type_from_bytes(b"not a known file") is None
+
+    def test_should_not_confuse_riff_without_webp_tag(self):
+        assert sniff_mime_type_from_bytes(b"RIFF\x00\x00\x00\x00WAVErest") is None
+
+
 class TestConvertFileDocumentToUrlDocument:
     def test_should_convert_pdf_pathlib_path_to_document_url(self):
         """pathlib.Path to a PDF should produce type=document_url with base64 data URI.
@@ -147,6 +181,62 @@ class TestConvertFileDocumentToUrlDocument:
         b64_data = result["document_url"].split(";base64,")[1]
         assert base64.b64decode(b64_data) == content
 
+    def test_should_infer_pdf_mime_from_raw_bytes_magic_number(self):
+        """Raw PDF bytes with no name or explicit MIME must be detected as application/pdf,
+        not application/octet-stream (which Mistral/Azure AI reject)."""
+        content = b"%PDF-1.4\n1 0 obj\n<< >>\nendobj\n"
+
+        result = convert_file_document_to_url_document(
+            {"type": "file", "file": content}
+        )
+
+        assert result["type"] == "document_url"
+        assert result["document_url"].startswith("data:application/pdf;base64,")
+
+    def test_should_infer_png_mime_from_raw_bytes_magic_number(self):
+        """Raw PNG bytes must be detected as image/png and produce an image_url."""
+        content = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR"
+
+        result = convert_file_document_to_url_document(
+            {"type": "file", "file": content}
+        )
+
+        assert result["type"] == "image_url"
+        assert result["image_url"].startswith("data:image/png;base64,")
+
+    def test_should_infer_pdf_mime_from_unnamed_bytesio_magic_number(self):
+        """An unnamed BytesIO carrying PDF bytes must be detected as application/pdf."""
+        file_obj = BytesIO(b"%PDF-1.5\n%\xe2\xe3\xcf\xd3\n")
+
+        result = convert_file_document_to_url_document(
+            {"type": "file", "file": file_obj}
+        )
+
+        assert result["type"] == "document_url"
+        assert result["document_url"].startswith("data:application/pdf;base64,")
+
+    def test_should_prefer_explicit_mime_over_sniffed_magic_number(self):
+        """An explicit mime_type must win over content sniffing."""
+        content = b"%PDF-1.4 pretends to be a pdf"
+
+        result = convert_file_document_to_url_document(
+            {"type": "file", "file": content, "mime_type": "image/png"}
+        )
+
+        assert result["type"] == "image_url"
+        assert result["image_url"].startswith("data:image/png;base64,")
+
+    def test_should_fallback_to_octet_stream_for_unrecognized_raw_bytes(self):
+        """Bytes matching no known signature stay octet-stream so behavior is unchanged."""
+        content = b"totally unrecognized payload"
+
+        result = convert_file_document_to_url_document(
+            {"type": "file", "file": content}
+        )
+
+        assert result["type"] == "document_url"
+        assert result["document_url"].startswith("data:application/octet-stream;base64,")
+
     def test_should_convert_raw_bytes_with_explicit_mime_type(self):
         """Raw bytes with explicit mime_type should use the specified MIME type."""
         content = b"raw pdf content"