From 39cec3ef4e7ed7cece737cff3a2e6feec32d9e54 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 12 Sep 2026 10:26:53 -0700 Subject: [PATCH] fix(ocr): align conversion contracts --- .../codecs/document_intelligence/params.rs | 24 +++++++++++ .../document_intelligence/transformation.rs | 5 +-- .../src/ocr/codecs/mistral/transformation.rs | 1 + .../core/src/ocr/codecs/mistral/types.rs | 9 +++- litellm-rust/crates/core/src/ocr/document.rs | 41 +++++++++++++++++++ .../tests/azure_document_intelligence_ocr.rs | 8 ++++ litellm-rust/crates/core/tests/ocr.rs | 3 +- 7 files changed, 85 insertions(+), 6 deletions(-) diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs index 85d1dafa542..9389f93b8e3 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs @@ -163,6 +163,30 @@ mod tests { ); } + #[rstest] + #[case(json!([0, 1, 2]), Some("1,2,3"))] + #[case(json!([2, 0, 0, 1]), Some("1,2,3"))] + #[case(json!([]), None)] + #[case(json!("3-9"), Some("3-9"))] + #[case(json!("1-3, 5"), Some("1-3,5"))] + #[case(json!(["1", "3-5"]), Some("1,3-5"))] + fn page_mapping_matches_python(#[case] input: Value, #[case] expected: Option<&str>) { + assert_eq!( + map(json!({"pages": input})).unwrap().pages.as_deref(), + expected + ); + } + + #[rstest] + #[case(json!("a,b"))] + #[case(json!([-1]))] + #[case(json!([true, false]))] + #[case(json!([1, "2"]))] + #[case(json!(5))] + fn invalid_page_mapping_matches_python(#[case] input: Value) { + assert!(map(json!({"pages": input})).is_err()); + } + #[rstest] #[case(json!(["keyValuePairs"]), "keyValuePairs")] #[case(json!(["keyValuePairs", "languages"]), "keyValuePairs,languages")] diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs index 2b848fcfb7a..4b56e8d4034 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs @@ -46,10 +46,7 @@ pub(crate) fn transform_ocr_response( let mut extra_fields = Map::new(); extra_fields.insert("content".into(), option_value(result.content)); extra_fields.insert("tables".into(), option_value(result.tables)); - extra_fields.insert( - "key_value_pairs".into(), - option_value(result.key_value_pairs), - ); + extra_fields.insert("keyValuePairs".into(), option_value(result.key_value_pairs)); Ok(LiteLLMOcrResponse { pages, model: model.into(), diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs index 5bd7e555a1e..b27ce60628e 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs @@ -133,6 +133,7 @@ mod tests { #[rstest] #[case("pages", json!([0, 2]))] + #[case("pages", json!("0,2-4"))] #[case("include_image_base64", json!(true))] #[case("image_limit", json!(2))] #[case("image_min_size", json!(100))] diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs index 0e601cd8319..e0bc8a267d2 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs @@ -3,10 +3,17 @@ use serde_json::{Map, Value}; use crate::ocr::types::OcrDocument; +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum MistralOcrPages { + Range(String), + Indices(Vec), +} + #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub(crate) struct MistralOcrParams { #[serde(skip_serializing_if = "Option::is_none")] - pub pages: Option>, + pub pages: Option, #[serde(skip_serializing_if = "Option::is_none")] pub include_image_base64: Option, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index 31fb15ccd4a..817eefda13a 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -206,6 +206,47 @@ mod tests { ); } + #[test] + fn file_name_mime_mapping_matches_python() { + for (name, expected) in [ + ("document.pdf", "application/pdf"), + ("image.png", "image/png"), + ("photo.jpg", "image/jpeg"), + ("photo.jpeg", "image/jpeg"), + ("animation.gif", "image/gif"), + ("image.webp", "image/webp"), + ("scan.tiff", "image/tiff"), + ("scan.tif", "image/tiff"), + ("bitmap.bmp", "image/bmp"), + ("DOCUMENT.PDF", "application/pdf"), + ("IMAGE.PNG", "image/png"), + ("file.unknown-extension", "application/octet-stream"), + ] { + assert_eq!(mime_type_for_name(name), expected); + } + } + + #[test] + fn upload_mime_mapping_matches_python() { + assert_eq!( + upload_mime_type(Some("report.pdf"), Some("application/octet-stream")), + "application/pdf" + ); + assert_eq!(upload_mime_type(Some("image.png"), None), "image/png"); + assert_eq!(upload_mime_type(None, None), "application/octet-stream"); + assert_eq!( + upload_mime_type(Some("doc.pdf"), Some("application/pdf; charset=utf-8")), + "application/pdf" + ); + assert_eq!( + upload_mime_type( + Some("img.png"), + Some("image/png; charset=utf-8; boundary=something") + ), + "image/png" + ); + } + #[test] fn file_encoding_enforces_decoded_size_limit() { let bytes = vec![b'a'; OCR_INLINE_MAX_BYTES + 1]; diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 7c635ebae62..463cee994a4 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -125,6 +125,14 @@ async fn immediate_response_normalizes_pages_and_preserves_native() { json!({"width":816,"height":1056,"dpi":96}) ); assert_eq!(result.usage_info, Some(json!({"pages_processed":1}))); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells":[]}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key":{"content":"A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); assert_eq!(result.provider_native_response, Some(operation)); } diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 97868e9e660..cbc7f949ce7 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -58,7 +58,7 @@ async fn facade_executes_direct_mistral_once() { let result = perform_ocr(wire_request( "mistral/model", &base, - json!({"extract_header":true,"unknown":"ignored"}), + json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), )) .await .unwrap(); @@ -79,6 +79,7 @@ async fn facade_executes_direct_mistral_once() { json!({ "model":"model", "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "pages":"0,2-4", "extract_header":true }) );