refactor(ocr-inputs): consolidate MIME precedence into pure core resolver with crate constants

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

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<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.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<Value> {
     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!({

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";

litellm-rust/crates/core/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod constants;
 pub mod error;
 pub mod ocr;
 pub mod providers;

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<String> {
+    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
+        );
+    }
 }