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 85f84cd2c2e..0df7cab4acf 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 @@ -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, 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!( diff --git a/litellm-rust/crates/core/src/ocr/mime.rs b/litellm-rust/crates/core/src/ocr/mime.rs new file mode 100644 index 00000000000..34b5d7d3df1 --- /dev/null +++ b/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); + } +} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index ec2fbb969a6..c0a0d9bb565 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,2 +1,3 @@ +pub mod mime; pub mod transformation; pub mod types; diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 548a5984f43..cab7e7d7164 100644 --- a/litellm/ocr/main.py +++ b/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 diff --git a/tests/e2e/gateway/test_ocr_rust_e2e.py b/tests/e2e/gateway/test_ocr_rust_e2e.py index dcf85898365..7559c18f9d4 100644 --- a/tests/e2e/gateway/test_ocr_rust_e2e.py +++ b/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())