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 42fabe1bd1e..68488963345 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,7 +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::mime::resolve_document_mime; use litellm_core::ocr::transformation::OcrProviderConfig; use litellm_core::CoreResult; use reqwest::Url; @@ -281,25 +281,6 @@ async fn read_response_with_limit( Ok(bytes) } -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.to_ascii_lowercase(); - } - } - sniff_document_mime(bytes) - .or_else(|| mime_from_file_name(url_path)) - .map(str::to_string) - .or_else(|| header_content_type.map(|header| header.to_ascii_lowercase())) - .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); @@ -321,12 +302,13 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes .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()) .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 content_type = resolve_document_mime( + header_content_type.as_deref(), + &bytes, + Some(final_url.path()), + ); let data_uri = format!( "data:{content_type};base64,{}", BASE64_STANDARD.encode(bytes) @@ -615,70 +597,6 @@ mod tests { )); } - #[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_normalizes_specific_header_casing() { - assert_eq!( - resolve_document_mime(Some("Application/PDF".to_string()), b"%PDF-1.4", "/x"), - "application/pdf" - ); - } - - #[test] - fn resolve_document_mime_sniffs_when_header_is_binary_octet_stream() { - assert_eq!( - resolve_document_mime( - Some("Binary/Octet-Stream".to_string()), - b"%PDF-1.4 payload", - "/x" - ), - "application/pdf" - ); - } - - #[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-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs new file mode 100644 index 00000000000..c5895f03230 --- /dev/null +++ b/litellm-rust/crates/core/src/constants.rs @@ -0,0 +1,9 @@ +pub const MIME_APPLICATION_PDF: &str = "application/pdf"; +pub const MIME_IMAGE_PNG: &str = "image/png"; +pub const MIME_IMAGE_JPEG: &str = "image/jpeg"; +pub const MIME_IMAGE_GIF: &str = "image/gif"; +pub const MIME_IMAGE_WEBP: &str = "image/webp"; +pub const MIME_IMAGE_TIFF: &str = "image/tiff"; +pub const MIME_IMAGE_BMP: &str = "image/bmp"; +pub const MIME_APPLICATION_OCTET_STREAM: &str = "application/octet-stream"; +pub const MIME_BINARY_OCTET_STREAM: &str = "binary/octet-stream"; diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 2ac479cc725..e0e15d8071c 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,3 +1,4 @@ +pub mod constants; pub mod error; pub mod ocr; pub mod providers; diff --git a/litellm-rust/crates/core/src/ocr/mime.rs b/litellm-rust/crates/core/src/ocr/mime.rs index abff428b39a..d1c6ddcc0a6 100644 --- a/litellm-rust/crates/core/src/ocr/mime.rs +++ b/litellm-rust/crates/core/src/ocr/mime.rs @@ -1,22 +1,27 @@ +use crate::constants::{ + MIME_APPLICATION_OCTET_STREAM, MIME_APPLICATION_PDF, MIME_BINARY_OCTET_STREAM, MIME_IMAGE_BMP, + MIME_IMAGE_GIF, MIME_IMAGE_JPEG, MIME_IMAGE_PNG, MIME_IMAGE_TIFF, MIME_IMAGE_WEBP, +}; + pub fn sniff_document_mime(bytes: &[u8]) -> Option<&'static str> { if bytes.starts_with(b"%PDF-") { - return Some("application/pdf"); + return Some(MIME_APPLICATION_PDF); } if bytes.starts_with(&[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a]) { - return Some("image/png"); + return Some(MIME_IMAGE_PNG); } if bytes.starts_with(&[0xff, 0xd8, 0xff]) { - return Some("image/jpeg"); + return Some(MIME_IMAGE_JPEG); } if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") { - return Some("image/gif"); + return Some(MIME_IMAGE_GIF); } if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" { - return Some("image/webp"); + return Some(MIME_IMAGE_WEBP); } if bytes.starts_with(&[0x49, 0x49, 0x2a, 0x00]) || bytes.starts_with(&[0x4d, 0x4d, 0x00, 0x2a]) { - return Some("image/tiff"); + return Some(MIME_IMAGE_TIFF); } None } @@ -25,17 +30,45 @@ 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"), + "pdf" => Some(MIME_APPLICATION_PDF), + "png" => Some(MIME_IMAGE_PNG), + "jpg" | "jpeg" => Some(MIME_IMAGE_JPEG), + "gif" => Some(MIME_IMAGE_GIF), + "webp" => Some(MIME_IMAGE_WEBP), + "tiff" | "tif" => Some(MIME_IMAGE_TIFF), + "bmp" => Some(MIME_IMAGE_BMP), _ => None, } } +pub fn normalize_declared_mime(declared: &str) -> Option { + let base = declared + .split(';') + .next() + .unwrap_or("") + .trim() + .to_ascii_lowercase(); + if base.is_empty() || base == MIME_APPLICATION_OCTET_STREAM || base == MIME_BINARY_OCTET_STREAM + { + return None; + } + Some(base) +} + +pub fn resolve_document_mime( + declared: Option<&str>, + bytes: &[u8], + file_name: Option<&str>, +) -> String { + if let Some(specific) = declared.and_then(normalize_declared_mime) { + return specific; + } + sniff_document_mime(bytes) + .or_else(|| file_name.and_then(mime_from_file_name)) + .map(str::to_string) + .unwrap_or_else(|| MIME_APPLICATION_OCTET_STREAM.to_string()) +} + #[cfg(test)] mod tests { use super::*; @@ -44,29 +77,29 @@ mod tests { fn sniff_document_mime_detects_supported_signatures() { assert_eq!( sniff_document_mime(b"%PDF-1.7\ncontent"), - Some("application/pdf") + Some(MIME_APPLICATION_PDF) ); assert_eq!( sniff_document_mime(&[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a, 0x00]), - Some("image/png") + Some(MIME_IMAGE_PNG) ); assert_eq!( sniff_document_mime(&[0xff, 0xd8, 0xff, 0xe0]), - Some("image/jpeg") + Some(MIME_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"GIF87a....."), Some(MIME_IMAGE_GIF)); + assert_eq!(sniff_document_mime(b"GIF89a....."), Some(MIME_IMAGE_GIF)); assert_eq!( sniff_document_mime(b"RIFF\x00\x00\x00\x00WEBPVP8 "), - Some("image/webp") + Some(MIME_IMAGE_WEBP) ); assert_eq!( sniff_document_mime(&[0x49, 0x49, 0x2a, 0x00]), - Some("image/tiff") + Some(MIME_IMAGE_TIFF) ); assert_eq!( sniff_document_mime(&[0x4d, 0x4d, 0x00, 0x2a]), - Some("image/tiff") + Some(MIME_IMAGE_TIFF) ); } @@ -86,15 +119,18 @@ mod tests { #[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")); + assert_eq!( + mime_from_file_name("report.pdf"), + Some(MIME_APPLICATION_PDF) + ); + assert_eq!(mime_from_file_name("/docs/scan.PNG"), Some(MIME_IMAGE_PNG)); + assert_eq!(mime_from_file_name("a/b/photo.jpeg"), Some(MIME_IMAGE_JPEG)); + assert_eq!(mime_from_file_name("photo.JPG"), Some(MIME_IMAGE_JPEG)); + assert_eq!(mime_from_file_name("anim.gif"), Some(MIME_IMAGE_GIF)); + assert_eq!(mime_from_file_name("sticker.webp"), Some(MIME_IMAGE_WEBP)); + assert_eq!(mime_from_file_name("scan.tif"), Some(MIME_IMAGE_TIFF)); + assert_eq!(mime_from_file_name("scan.tiff"), Some(MIME_IMAGE_TIFF)); + assert_eq!(mime_from_file_name("bitmap.bmp"), Some(MIME_IMAGE_BMP)); } #[test] @@ -104,4 +140,75 @@ mod tests { assert_eq!(mime_from_file_name("/trailing/"), None); assert_eq!(mime_from_file_name(""), None); } + + #[test] + fn normalize_declared_mime_lowercases_and_strips_parameters() { + assert_eq!( + normalize_declared_mime("Application/PDF"), + Some(MIME_APPLICATION_PDF.to_string()) + ); + assert_eq!( + normalize_declared_mime("image/png; charset=utf-8"), + Some(MIME_IMAGE_PNG.to_string()) + ); + } + + #[test] + fn normalize_declared_mime_treats_generic_and_blank_as_absent() { + assert_eq!(normalize_declared_mime("application/octet-stream"), None); + assert_eq!(normalize_declared_mime("Binary/Octet-Stream"), None); + assert_eq!( + normalize_declared_mime("application/octet-stream; charset=binary"), + None + ); + assert_eq!(normalize_declared_mime(" "), None); + } + + #[test] + fn resolve_document_mime_prefers_specific_declared() { + assert_eq!( + resolve_document_mime(Some("image/png"), b"%PDF-1.4", Some("/x.pdf")), + MIME_IMAGE_PNG + ); + } + + #[test] + fn resolve_document_mime_normalizes_declared_casing_and_parameters() { + assert_eq!( + resolve_document_mime(Some("Application/PDF; charset=utf-8"), b"", Some("/x")), + MIME_APPLICATION_PDF + ); + } + + #[test] + fn resolve_document_mime_sniffs_when_declared_is_generic() { + assert_eq!( + resolve_document_mime(Some("Binary/Octet-Stream"), b"%PDF-1.4 payload", Some("/x")), + MIME_APPLICATION_PDF + ); + } + + #[test] + fn resolve_document_mime_uses_file_name_when_declared_and_bytes_ambiguous() { + assert_eq!( + resolve_document_mime(None, b"unrecognized bytes", Some("/docs/file.pdf")), + MIME_APPLICATION_PDF + ); + } + + #[test] + fn resolve_document_mime_falls_back_to_octet_stream() { + assert_eq!( + resolve_document_mime(None, b"unrecognized bytes", Some("/docs/file")), + MIME_APPLICATION_OCTET_STREAM + ); + assert_eq!( + resolve_document_mime( + Some("application/octet-stream"), + b"unrecognized bytes", + None + ), + MIME_APPLICATION_OCTET_STREAM + ); + } }