diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs index 36d3d7da2e5..85f84cd2c2e 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs +++ b/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, + 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 { 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!({ diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 9225b6ff13c..e5c7c8db037 100644 --- a/litellm/ocr/main.py +++ b/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"] diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/test_litellm/ocr/test_ocr_file_input.py index e6216d7c580..35628d0aedc 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/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"