From eb502824f0af99ca3e035581c0d428b182c701e9 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 09:32:46 -0700 Subject: [PATCH] fix stuff --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/auth-aws/src/signer.rs | 5 - .../crates/core/src/ocr/provider_config.rs | 23 +- litellm-rust/crates/llms/Cargo.toml | 1 + .../ocr/analyze_transformation.rs | 409 +++++++------ .../llms/src/aws_textract/ocr/common_utils.rs | 559 ++++++++++++++++-- .../src/aws_textract/ocr/transformation.rs | 201 +++---- .../crates/llms/src/base_llm/ocr/error.rs | 7 + 8 files changed, 850 insertions(+), 356 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 72f5e70eea5..860f01c4ad1 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2176,6 +2176,7 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_with", + "strum", "thiserror 2.0.19", "time", "tokio", diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs index 46d6fb5c391..49a3910c1d5 100644 --- a/litellm-rust/crates/auth-aws/src/signer.rs +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -9,8 +9,6 @@ use crate::{ is_sigv4_computed_header, resolve_credentials, sign_post, }; -/// SigV4 over the serialized body. Credentials are resolved up front, since -/// they do not depend on the body; the signature waits for the final bytes. #[derive(Clone, Debug)] pub struct SigV4Signer { region: String, @@ -33,8 +31,6 @@ impl SigV4Signer { Self { clock, ..self } } - /// A host with its own resolution chain hands credentials down in - /// `optional_params`; only derive them here when it supplied none. pub async fn resolve( region: String, service: &'static str, @@ -57,7 +53,6 @@ impl RequestSigner for SigV4Signer { &self, request: UnsignedRequest<'_>, ) -> Result, litellm_http::Error> { - // Sending a caller's copy next to the computed one is rejected by AWS. if let Some((name, _)) = request .headers .iter() diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index 34e2a77b6d1..d38d87b92cc 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -1,7 +1,7 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; use litellm_llms::{ aws_textract::ocr::{ - analyze_transformation::TextractAnalyzeDocumentConfig, + analyze_transformation::TextractAnalyzeDocumentConfig, common_utils::TextractOperation, transformation::TextractDetectTextConfig, }, azure_ai::ocr::{ @@ -178,10 +178,10 @@ pub(crate) fn resolve_provider_config( .parse::() .map_err(|_| Error::InvalidProvider(provider.custom_llm_provider.to_string()))?; let config = match ocr_provider { - OcrProvider::AwsTextract if provider.model.eq_ignore_ascii_case("analyze-document") => { - OcrConfigKind::AwsTextractAnalyze - } - OcrProvider::AwsTextract => OcrConfigKind::AwsTextract, + OcrProvider::AwsTextract => match TextractOperation::from_model(provider.model)? { + TextractOperation::DetectDocumentText => OcrConfigKind::AwsTextract, + TextractOperation::AnalyzeDocument => OcrConfigKind::AwsTextractAnalyze, + }, OcrProvider::Cohere => OcrConfigKind::Cohere, OcrProvider::Mistral => OcrConfigKind::Mistral, OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => { @@ -438,6 +438,19 @@ mod tests { assert_eq!(config, expected_config); } + #[rstest] + #[case::misspelled_operation("aws_textract/analyse-document")] + #[case::operation_name_from_the_api("aws_textract/AnalyzeDocument")] + fn textract_models_outside_its_two_operations_are_refused(#[case] model: &str) { + assert!(matches!( + resolve_provider_config(model, None), + Err(Error::InvalidModel { + provider: "aws_textract", + .. + }) + )); + } + #[rstest] #[case("aws_textract/detect-document-text", OcrConfigKind::AwsTextract)] #[case("aws_textract/analyze-document", OcrConfigKind::AwsTextractAnalyze)] diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 7afc4171ca8..0cc7af1836f 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -27,6 +27,7 @@ serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } serde_path_to_error = "0.1" serde_with.workspace = true +strum.workspace = true thiserror.workspace = true time.workspace = true tokio = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs index 65c7688ea0c..d476861e6e1 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs @@ -4,24 +4,24 @@ use litellm_core_utils::call_arguments::{CallArguments, parse_options}; use serde::{Deserialize, Serialize}; use super::common_utils::{ - Block, DocumentMetadata, HEALTH_CHECK_IMAGE_DATA_URI, TextractDocument, TextractEnvironment, - document_bytes, endpoint, environment, error_class, inline_document, lines_by_page, + Block, BlockType, FeatureType, LayoutType, TextractDocument, TextractEnvironment, + TextractOperation, TextractResponse, document_bytes, endpoint, environment, error_class, + health_check_document, inline_document, lines_by_page, ocr_response, }; use crate::base_llm::ocr::{ error::Error, handler::OcrClient, transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrPage, OcrRequestContext, - OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, decode_and_normalize_response, + BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrRequestContext, OcrResponseFormat, + PreparedOcrRequest, decode_and_normalize_response, }, }; -const ANALYZE_DOCUMENT_TARGET: &str = "Textract.AnalyzeDocument"; -const DEFAULT_FEATURE_TYPES: [&str; 2] = ["LAYOUT", "TABLES"]; +const DEFAULT_FEATURE_TYPES: [FeatureType; 2] = [FeatureType::Layout, FeatureType::Tables]; #[derive(Default, Deserialize)] pub struct AnalyzeDocumentOptions { - pub feature_types: Option>, + pub feature_types: Option>, } #[derive(Debug, Deserialize, Serialize)] @@ -29,18 +29,9 @@ pub struct AnalyzeDocumentRequest { #[serde(rename = "Document")] pub document: TextractDocument, #[serde(rename = "FeatureTypes")] - pub feature_types: Vec, + pub feature_types: Vec, } -#[derive(Deserialize)] -#[serde(rename_all = "PascalCase")] -pub struct AnalyzeDocumentResponse { - #[serde(default)] - blocks: Vec, - document_metadata: Option, -} - -/// Synchronous `AnalyzeDocument`: layout and tables rendered as markdown. #[derive(Clone, Copy, Debug, Default)] pub struct TextractAnalyzeDocumentConfig; @@ -54,10 +45,7 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig { } fn get_health_check_document(&self) -> OcrDocument { - OcrDocument::ImageUrl { - image_url: HEALTH_CHECK_IMAGE_DATA_URI.into(), - extra_fields: Default::default(), - } + health_check_document() } fn map_ocr_params( @@ -73,7 +61,7 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig { request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - environment(request, ANALYZE_DOCUMENT_TARGET).await + environment(request, TextractOperation::AnalyzeDocument).await } fn get_complete_url( @@ -94,12 +82,10 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig { ) -> Result { Ok(AnalyzeDocumentRequest { document: document_bytes(&document)?, - feature_types: optional_params.feature_types.clone().unwrap_or_else(|| { - DEFAULT_FEATURE_TYPES - .iter() - .map(|feature| feature.to_string()) - .collect() - }), + feature_types: optional_params + .feature_types + .clone() + .unwrap_or_else(|| DEFAULT_FEATURE_TYPES.to_vec()), }) } @@ -136,10 +122,12 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig { fn normalize_response( model: &str, - response: AnalyzeDocumentResponse, + response: TextractResponse, ) -> Result { let blocks = &response.blocks; - let has_layout = blocks.iter().any(is_layout); + let has_layout = blocks + .iter() + .any(|block| block.block_type.layout().is_some()); let page_markdown: Vec<(i64, String)> = if has_layout { let by_id: HashMap<&str, &Block> = blocks .iter() @@ -154,65 +142,63 @@ fn normalize_response( } else { lines_by_page(blocks) }; - let pages: Vec = page_markdown - .into_iter() - .map(|(page, markdown)| OcrPage { - index: page - 1, - markdown, - ..Default::default() - }) - .collect(); - let pages_processed = response - .document_metadata - .and_then(|metadata| metadata.pages) - .or_else(|| i64::try_from(pages.len()).ok()); - Ok(LiteLLMOcrResponse { - usage_info: Some(OcrUsageInfo { - pages_processed, - ..Default::default() - }), - ..LiteLLMOcrResponse::new(model, pages) - }) -} - -fn is_layout(block: &Block) -> bool { - block.block_type.starts_with("LAYOUT_") + Ok(ocr_response( + model, + page_markdown, + response.document_metadata, + )) } /// Layout blocks arrive in reading order. A list's items are repeated as -/// top-level `LAYOUT_TEXT` blocks, and a `LAYOUT_TABLE` only links to the -/// table's lines, so the nth layout table on a page takes the nth `TABLE`. +/// top-level `LAYOUT_TEXT` blocks. A `LAYOUT_TABLE` that links to its `TABLE` +/// renders it; one that only links to the table's lines takes the `TABLE` at +/// the same position on the page. fn layout_markdown(blocks: &[Block], page: i64, by_id: &HashMap<&str, &Block>) -> String { let on_page = || blocks.iter().filter(move |block| block.page() == page); let list_items: BTreeSet<&str> = on_page() - .filter(|block| block.block_type == "LAYOUT_LIST") + .filter(|block| block.block_type == BlockType::LayoutList) .flat_map(Block::children) .collect(); let tables: Vec<&Block> = on_page() - .filter(|block| block.block_type == "TABLE") + .filter(|block| block.block_type == BlockType::Table) .collect(); let table_ordinal: HashMap<&str, usize> = on_page() - .filter(|block| block.block_type == "LAYOUT_TABLE") + .filter(|block| block.block_type == BlockType::LayoutTable) .enumerate() .map(|(ordinal, block)| (block.id.as_str(), ordinal)) .collect(); + let table_of = |layout_table: &Block| { + layout_table + .children() + .filter_map(|id| by_id.get(id).copied()) + .find(|child| child.block_type == BlockType::Table) + .or_else(|| { + table_ordinal + .get(layout_table.id.as_str()) + .and_then(|ordinal| tables.get(*ordinal).copied()) + }) + }; let sections: Vec = on_page() - .filter(|block| is_layout(block) && !list_items.contains(block.id.as_str())) - .map(|block| match block.block_type.as_str() { - "LAYOUT_TITLE" => format!("# {}", text_of(block, by_id, " ")), - "LAYOUT_SECTION_HEADER" => format!("## {}", text_of(block, by_id, " ")), - "LAYOUT_LIST" => block + .filter(|block| !list_items.contains(block.id.as_str())) + .filter_map(|block| Some((block, block.block_type.layout()?))) + .map(|(block, layout)| match layout { + LayoutType::Title => format!("# {}", text_of(block, by_id, " ")), + LayoutType::SectionHeader => format!("## {}", text_of(block, by_id, " ")), + LayoutType::List => block .children() .filter_map(|id| by_id.get(id)) .map(|item| format!("- {}", strip_bullet(&text_of(item, by_id, " ")))) .collect::>() .join("\n"), - "LAYOUT_TABLE" => table_ordinal - .get(block.id.as_str()) - .and_then(|ordinal| tables.get(*ordinal)) - .map(|table| table_markdown(table, by_id)) - .unwrap_or_else(|| text_of(block, by_id, "\n")), - _ => text_of(block, by_id, " "), + LayoutType::Table => match table_of(block) { + Some(table) => table_markdown(table, by_id), + None => text_of(block, by_id, "\n"), + }, + LayoutType::KeyValue => text_of(block, by_id, "\n"), + LayoutType::Figure => String::new(), + LayoutType::Text | LayoutType::Header | LayoutType::Footer | LayoutType::PageNumber => { + text_of(block, by_id, " ") + } }) .filter(|section| !section.trim().is_empty()) .collect(); @@ -241,7 +227,7 @@ fn table_markdown(table: &Block, by_id: &HashMap<&str, &Block>) -> String { let cells: BTreeMap<(usize, usize), String> = table .children() .filter_map(|id| by_id.get(id)) - .filter(|cell| cell.block_type == "CELL") + .filter(|cell| cell.block_type == BlockType::Cell) .filter_map(|cell| { Some(( (cell.row_index?, cell.column_index?), @@ -269,23 +255,19 @@ fn table_markdown(table: &Block, by_id: &HashMap<&str, &Block>) -> String { #[cfg(test)] mod tests { + use rstest::{fixture, rstest}; use serde_json::{Value, json}; use super::*; - fn markdown(blocks: Value) -> Vec<(i64, String)> { - TextractAnalyzeDocumentConfig - .transform_ocr_response( - "analyze-document", - &serde_json::to_vec(&json!({"DocumentMetadata": {"Pages": 1}, "Blocks": blocks})) - .unwrap(), - OcrResponseFormat::Litellm, - ) - .unwrap() - .pages - .into_iter() - .map(|page| (page.index, page.markdown)) - .collect() + const MODEL: &str = "analyze-document"; + + #[fixture] + fn document() -> OcrDocument { + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,aGk=".into(), + extra_fields: Default::default(), + } } fn child(ids: &[&str]) -> Value { @@ -300,40 +282,57 @@ mod tests { json!({"Id": id, "BlockType": "WORD", "Text": text}) } + fn layout(id: &str, block_type: &str, children: &[&str]) -> Value { + json!({"Id": id, "BlockType": block_type, "Relationships": child(children)}) + } + + fn table(id: &str, cells: &[&str]) -> Value { + json!({"Id": id, "BlockType": "TABLE", "Relationships": child(cells)}) + } + fn cell(id: &str, row: usize, column: usize, words: &[&str]) -> Value { json!({"Id": id, "BlockType": "CELL", "RowIndex": row, "ColumnIndex": column, "Relationships": child(words)}) } - #[test] - fn layout_becomes_headings_paragraphs_and_a_list_without_repeating_its_items() { - let pages = markdown(json!([ + fn on_page(page: i64, mut block: Value) -> Value { + block["Page"] = json!(page); + block + } + + #[rstest] + #[case::headings_paragraphs_and_a_list_without_repeating_its_items( + json!([ line("l1", "Quarterly Report"), line("l2", "This report lists"), line("l3", "the invoices."), line("l4", "Line items"), line("l5", "- Pay within 30 days"), line("l6", "\u{2022} Quote the number"), - {"Id": "t", "BlockType": "LAYOUT_TITLE", "Relationships": child(&["l1"])}, - {"Id": "p", "BlockType": "LAYOUT_TEXT", "Relationships": child(&["l2", "l3"])}, - {"Id": "h", "BlockType": "LAYOUT_SECTION_HEADER", "Relationships": child(&["l4"])}, - {"Id": "ul", "BlockType": "LAYOUT_LIST", "Relationships": child(&["i1", "i2"])}, - {"Id": "i1", "BlockType": "LAYOUT_TEXT", "Relationships": child(&["l5"])}, - {"Id": "i2", "BlockType": "LAYOUT_TEXT", "Relationships": child(&["l6"])} - ])); - - assert_eq!( - pages, - vec![( - 0, - "# Quarterly Report\n\nThis report lists the invoices.\n\n## Line items\n\n- Pay within 30 days\n- Quote the number".to_string() - )] - ); - } - - #[test] - fn a_layout_table_is_rendered_from_the_table_cells_in_row_and_column_order() { - let pages = markdown(json!([ + layout("t", "LAYOUT_TITLE", &["l1"]), + layout("p", "LAYOUT_TEXT", &["l2", "l3"]), + layout("h", "LAYOUT_SECTION_HEADER", &["l4"]), + layout("ul", "LAYOUT_LIST", &["i1", "i2"]), + layout("i1", "LAYOUT_TEXT", &["l5"]), + layout("i2", "LAYOUT_TEXT", &["l6"]) + ]), + vec![( + 0, + "# Quarterly Report\n\nThis report lists the invoices.\n\n## Line items\n\n- Pay within 30 days\n- Quote the number" + )] + )] + #[case::header_footer_and_page_number_stay_in_reading_order( + json!([ + line("l1", "ACME Corp"), line("l2", "Body"), line("l3", "Confidential"), line("l4", "3"), + layout("hd", "LAYOUT_HEADER", &["l1"]), + layout("p", "LAYOUT_TEXT", &["l2"]), + layout("ft", "LAYOUT_FOOTER", &["l3"]), + layout("pn", "LAYOUT_PAGE_NUMBER", &["l4"]) + ]), + vec![(0, "ACME Corp\n\nBody\n\nConfidential\n\n3")] + )] + #[case::a_table_is_rendered_from_its_cells_in_row_and_column_order( + json!([ line("l1", "Invoice"), line("l2", "Total"), line("l3", "12345"), line("l4", "a|b"), word("w1", "Invoice"), word("w2", "Total"), word("w3", "12345"), word("w4", "a|b"), {"Id": "tb", "BlockType": "TABLE", "Relationships": [ @@ -342,85 +341,139 @@ mod tests { ]}, cell("c1", 1, 1, &["w1"]), cell("c2", 1, 2, &["w2"]), cell("c3", 2, 1, &["w3"]), cell("c4", 2, 2, &["w4"]), - {"Id": "lt", "BlockType": "LAYOUT_TABLE", "Relationships": child(&["l1", "l2", "l3", "l4"])} - ])); - - assert_eq!( - pages, - vec![( - 0, - "| Invoice | Total |\n| --- | --- |\n| 12345 | a\\|b |".to_string() - )] - ); - } - - #[test] - fn a_layout_table_without_table_blocks_keeps_its_lines() { - let pages = markdown(json!([ + layout("lt", "LAYOUT_TABLE", &["l1", "l2", "l3", "l4"]) + ]), + vec![(0, "| Invoice | Total |\n| --- | --- |\n| 12345 | a\\|b |")] + )] + #[case::a_layout_table_that_links_its_table_renders_that_one( + json!([ + word("w1", "first"), word("w2", "second"), + table("tb1", &["c1"]), cell("c1", 1, 1, &["w1"]), + table("tb2", &["c2"]), cell("c2", 1, 1, &["w2"]), + layout("lt", "LAYOUT_TABLE", &["tb2"]) + ]), + vec![(0, "| second |\n| --- |")] + )] + #[case::a_missing_cell_leaves_an_empty_column( + json!([ + word("w1", "a"), word("w2", "b"), word("w3", "c"), + table("tb", &["c1", "c2", "c3"]), + cell("c1", 1, 1, &["w1"]), cell("c2", 1, 2, &["w2"]), cell("c3", 2, 2, &["w3"]), + layout("lt", "LAYOUT_TABLE", &[]) + ]), + vec![(0, "| a | b |\n| --- | --- |\n| | c |")] + )] + #[case::a_layout_table_without_table_blocks_keeps_its_lines( + json!([ line("l1", "Invoice Total"), line("l2", "12345 67.89"), - {"Id": "lt", "BlockType": "LAYOUT_TABLE", "Relationships": child(&["l1", "l2"])} - ])); - - assert_eq!(pages, vec![(0, "Invoice Total\n12345 67.89".to_string())]); - } - - #[test] - fn a_response_without_layout_blocks_falls_back_to_lines() { - let pages = markdown(json!([ - line("l1", "first"), - word("w1", "first"), - line("l2", "second") - ])); - - assert_eq!(pages, vec![(0, "first\nsecond".to_string())]); - } - - #[test] - fn each_page_gets_its_own_markdown_and_its_own_tables() { - let pages = markdown(json!([ - {"Id": "a", "BlockType": "LINE", "Text": "one", "Page": 1}, - {"Id": "b", "BlockType": "LINE", "Text": "two", "Page": 2}, - {"Id": "w", "BlockType": "WORD", "Text": "cell", "Page": 2}, - {"Id": "t1", "BlockType": "LAYOUT_TEXT", "Page": 1, "Relationships": child(&["a"])}, - {"Id": "tb", "BlockType": "TABLE", "Page": 2, "Relationships": child(&["c"])}, - {"Id": "c", "BlockType": "CELL", "Page": 2, "RowIndex": 1, "ColumnIndex": 1, - "Relationships": child(&["w"])}, - {"Id": "lt", "BlockType": "LAYOUT_TABLE", "Page": 2, "Relationships": child(&["b"])} - ])); - - assert_eq!( - pages, - vec![(0, "one".to_string()), (1, "| cell |\n| --- |".to_string())] - ); - } - - #[test] - fn feature_types_default_to_layout_and_tables_and_can_be_overridden() { - let document = || OcrDocument::ImageUrl { - image_url: "data:image/png;base64,aGk=".into(), - extra_fields: Default::default(), - }; - let request = |options: Value| { - let arguments: CallArguments = serde_json::from_value(options).unwrap(); - let params = TextractAnalyzeDocumentConfig - .map_ocr_params(&arguments, "analyze-document") - .unwrap(); - serde_json::to_value( - TextractAnalyzeDocumentConfig - .transform_ocr_request("analyze-document", document(), ¶ms, &[]) + layout("lt", "LAYOUT_TABLE", &["l1", "l2"]) + ]), + vec![(0, "Invoice Total\n12345 67.89")] + )] + #[case::key_values_keep_one_line_each( + json!([ + line("l1", "Name: Ana"), + line("l2", "Date: 2024-01-01"), + layout("kv", "LAYOUT_KEY_VALUE", &["l1", "l2"]) + ]), + vec![(0, "Name: Ana\nDate: 2024-01-01")] + )] + #[case::a_figure_has_no_markdown( + json!([ + line("l1", "Caption"), + layout("f", "LAYOUT_FIGURE", &[]), + layout("p", "LAYOUT_TEXT", &["l1"]) + ]), + vec![(0, "Caption")] + )] + #[case::a_block_type_added_later_is_ignored( + json!([ + line("l1", "Body"), + layout("new", "LAYOUT_SIDEBAR", &["l1"]), + layout("p", "LAYOUT_TEXT", &["l1"]) + ]), + vec![(0, "Body")] + )] + #[case::without_layout_blocks_lines_are_used( + json!([line("l1", "first"), word("w1", "first"), line("l2", "second")]), + vec![(0, "first\nsecond")] + )] + #[case::each_page_gets_its_own_markdown_and_its_own_tables( + json!([ + on_page(1, line("a", "one")), + on_page(2, line("b", "two")), + on_page(2, word("w", "cell")), + on_page(1, layout("t1", "LAYOUT_TEXT", &["a"])), + on_page(2, table("tb", &["c"])), + on_page(2, cell("c", 1, 1, &["w"])), + on_page(2, layout("lt", "LAYOUT_TABLE", &["b"])) + ]), + vec![(0, "one"), (1, "| cell |\n| --- |")] + )] + fn blocks_become_markdown_pages(#[case] blocks: Value, #[case] expected: Vec<(i64, &str)>) { + let response = TextractAnalyzeDocumentConfig + .transform_ocr_response( + MODEL, + &serde_json::to_vec(&json!({"DocumentMetadata": {"Pages": 1}, "Blocks": blocks})) .unwrap(), + OcrResponseFormat::Litellm, ) - .unwrap() - }; + .unwrap(); + + let pages: Vec<(i64, &str)> = response + .pages + .iter() + .map(|page| (page.index, page.markdown.as_str())) + .collect(); + assert_eq!(pages, expected); + } + + #[rstest] + #[case::hyphen("- item", "item")] + #[case::asterisk("* item", "item")] + #[case::bullet("\u{2022} item", "item")] + #[case::middle_dot("\u{00b7}item", "item")] + #[case::no_bullet("item - with a dash", "item - with a dash")] + fn list_items_lose_their_own_bullet(#[case] item: &str, #[case] expected: &str) { + assert_eq!(strip_bullet(item), expected); + } + + #[rstest] + #[case::defaults_to_layout_and_tables(json!({}), json!(["LAYOUT", "TABLES"]))] + #[case::overridden(json!({"feature_types": ["FORMS", "SIGNATURES"]}), json!(["FORMS", "SIGNATURES"]))] + #[case::explicit_null_uses_the_default(json!({"feature_types": null}), json!(["LAYOUT", "TABLES"]))] + fn feature_types_reach_the_request( + document: OcrDocument, + #[case] arguments: Value, + #[case] expected: Value, + ) { + let arguments: CallArguments = serde_json::from_value(arguments).unwrap(); + let params = TextractAnalyzeDocumentConfig + .map_ocr_params(&arguments, MODEL) + .unwrap(); + + let request = TextractAnalyzeDocumentConfig + .transform_ocr_request(MODEL, document, ¶ms, &[]) + .unwrap(); assert_eq!( - request(json!({})), - json!({"Document": {"Bytes": "aGk="}, "FeatureTypes": ["LAYOUT", "TABLES"]}) + serde_json::to_value(request).unwrap(), + json!({"Document": {"Bytes": "aGk="}, "FeatureTypes": expected}) ); - assert_eq!( - request(json!({"feature_types": ["FORMS"]}))["FeatureTypes"], - json!(["FORMS"]) + } + + #[rstest] + #[case::undocumented_feature(json!({"feature_types": ["HANDWRITING"]}))] + #[case::lowercase_feature(json!({"feature_types": ["layout"]}))] + #[case::not_a_list(json!({"feature_types": "LAYOUT"}))] + fn feature_types_outside_the_documented_values_are_refused(#[case] arguments: Value) { + let arguments: CallArguments = serde_json::from_value(arguments).unwrap(); + + assert!( + TextractAnalyzeDocumentConfig + .map_ocr_params(&arguments, MODEL) + .is_err() ); } } diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs index bc1b5eb66d9..8268ad066a1 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs @@ -2,20 +2,53 @@ use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_auth_aws::{SigV4Signer, resolve_aws_region}; use litellm_http::outbound::RequestSigner; use serde::{Deserialize, Serialize}; +use strum::{EnumString, IntoStaticStr, VariantNames}; use crate::base_llm::ocr::{ document::{InlineDocument, inline_remote_document}, error::Error, transformation::{ - OCR_INLINE_MAX_BYTES, OcrDocument, OcrEnvironment, OcrRequestContext, PreparedOcrRequest, + LiteLLMOcrResponse, OcrDocument, OcrEnvironment, OcrPage, OcrRequestContext, OcrUsageInfo, + PreparedOcrRequest, }, }; const TEXTRACT_SERVICE: &str = "textract"; const AWS_JSON_CONTENT_TYPE: &str = "application/x-amz-json-1.1"; +const TARGET_HEADER: &str = "X-Amz-Target"; +const CONTENT_TYPE_HEADER: &str = "Content-Type"; const UNSUPPORTED_DOCUMENT: &str = "UnsupportedDocumentException"; +const SYNC_DOCUMENT_MAX_BYTES: usize = 10 * 1024 * 1024; -pub(super) const HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"; +const HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"; + +/// Textract has operations rather than models; the model slot of +/// `aws_textract/` names the one to call. +#[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, VariantNames, PartialEq, Eq)] +#[strum(serialize_all = "kebab-case", ascii_case_insensitive)] +pub enum TextractOperation { + DetectDocumentText, + AnalyzeDocument, +} + +impl TextractOperation { + pub const PROVIDER: &'static str = "aws_textract"; + + pub fn from_model(model: &str) -> Result { + model.parse().map_err(|_| Error::InvalidModel { + provider: Self::PROVIDER, + model: model.to_string(), + supported: Self::VARIANTS, + }) + } + + fn target(self) -> &'static str { + match self { + Self::DetectDocumentText => "Textract.DetectDocumentText", + Self::AnalyzeDocument => "Textract.AnalyzeDocument", + } + } +} #[derive(Debug, Deserialize, Serialize)] pub struct TextractDocument { @@ -23,12 +56,115 @@ pub struct TextractDocument { pub bytes: String, } +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum FeatureType { + Tables, + Forms, + Queries, + Signatures, + Layout, +} + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub(super) enum BlockType { + KeyValueSet, + Page, + Line, + Word, + Table, + Cell, + SelectionElement, + MergedCell, + Title, + Query, + QueryResult, + Signature, + TableTitle, + TableFooter, + LayoutText, + LayoutTitle, + LayoutHeader, + LayoutFooter, + LayoutSectionHeader, + LayoutPageNumber, + LayoutList, + LayoutFigure, + LayoutTable, + LayoutKeyValue, + #[serde(other)] + Unknown, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum LayoutType { + Text, + Title, + Header, + Footer, + SectionHeader, + PageNumber, + List, + Figure, + Table, + KeyValue, +} + +impl BlockType { + pub fn layout(self) -> Option { + match self { + Self::LayoutText => Some(LayoutType::Text), + Self::LayoutTitle => Some(LayoutType::Title), + Self::LayoutHeader => Some(LayoutType::Header), + Self::LayoutFooter => Some(LayoutType::Footer), + Self::LayoutSectionHeader => Some(LayoutType::SectionHeader), + Self::LayoutPageNumber => Some(LayoutType::PageNumber), + Self::LayoutList => Some(LayoutType::List), + Self::LayoutFigure => Some(LayoutType::Figure), + Self::LayoutTable => Some(LayoutType::Table), + Self::LayoutKeyValue => Some(LayoutType::KeyValue), + Self::KeyValueSet + | Self::Page + | Self::Line + | Self::Word + | Self::Table + | Self::Cell + | Self::SelectionElement + | Self::MergedCell + | Self::Title + | Self::Query + | Self::QueryResult + | Self::Signature + | Self::TableTitle + | Self::TableFooter + | Self::Unknown => None, + } + } +} + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub(super) enum RelationshipType { + Value, + Child, + ComplexFeatures, + MergedCell, + Title, + Answer, + Table, + TableTitle, + TableFooter, + #[serde(other)] + Unknown, +} + #[derive(Deserialize)] #[serde(rename_all = "PascalCase")] pub(super) struct Block { #[serde(default)] pub id: String, - pub block_type: String, + pub block_type: BlockType, pub text: Option, pub page: Option, pub row_index: Option, @@ -40,13 +176,12 @@ pub(super) struct Block { #[derive(Deserialize)] #[serde(rename_all = "PascalCase")] pub(super) struct Relationship { - pub r#type: String, + pub r#type: RelationshipType, #[serde(default)] pub ids: Vec, } impl Block { - /// The synchronous API omits `Page` because it only ever reads one. pub fn page(&self) -> i64 { self.page.unwrap_or(1) } @@ -54,7 +189,7 @@ impl Block { pub fn children(&self) -> impl Iterator { self.relationships .iter() - .filter(|relationship| relationship.r#type == "CHILD") + .filter(|relationship| relationship.r#type == RelationshipType::Child) .flat_map(|relationship| relationship.ids.iter().map(String::as_str)) } } @@ -65,6 +200,14 @@ pub(super) struct DocumentMetadata { pub pages: Option, } +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +pub struct TextractResponse { + #[serde(default)] + pub(super) blocks: Vec, + pub(super) document_metadata: Option, +} + pub struct TextractEnvironment { headers: Vec<(String, String)>, region: String, @@ -81,9 +224,16 @@ impl OcrEnvironment for TextractEnvironment { } } +pub(super) fn health_check_document() -> OcrDocument { + OcrDocument::ImageUrl { + image_url: HEALTH_CHECK_IMAGE_DATA_URI.into(), + extra_fields: Default::default(), + } +} + pub(super) async fn environment( request: &PreparedOcrRequest, - target: &'static str, + operation: TextractOperation, ) -> Result { let env_lookup = |name: &str| request.connection.secret(name); let region = @@ -102,21 +252,38 @@ pub(super) async fn environment( .await .map_err(litellm_auth::Error::from)?; Ok(TextractEnvironment { - headers: request - .connection - .extra_headers - .iter() - .cloned() - .chain([ - ("X-Amz-Target".into(), target.into()), - ("Content-Type".into(), AWS_JSON_CONTENT_TYPE.into()), - ]) - .collect(), + headers: operation_headers(&request.connection.extra_headers, operation), region, signer, }) } +/// A caller's copy of an operation header would reach the wire next to ours +/// while the signature covers only one value, which Textract rejects. +fn operation_headers( + extra_headers: &[(String, String)], + operation: TextractOperation, +) -> Vec<(String, String)> { + let operation = [ + (TARGET_HEADER, operation.target()), + (CONTENT_TYPE_HEADER, AWS_JSON_CONTENT_TYPE), + ]; + extra_headers + .iter() + .filter(|(name, _)| { + !operation + .iter() + .any(|(operation_name, _)| name.eq_ignore_ascii_case(operation_name)) + }) + .cloned() + .chain( + operation + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())), + ) + .collect() +} + pub(super) fn endpoint(request: &PreparedOcrRequest, environment: &TextractEnvironment) -> String { request .connection @@ -128,7 +295,7 @@ pub(super) fn endpoint(request: &PreparedOcrRequest, environment: &TextractEnvir pub(super) fn document_bytes(document: &OcrDocument) -> Result { let inline = InlineDocument::parse(document.source())?.ok_or(Error::InvalidDataUri)?; Ok(TextractDocument { - bytes: STANDARD.encode(inline.decode(OCR_INLINE_MAX_BYTES)?), + bytes: STANDARD.encode(inline.decode(SYNC_DOCUMENT_MAX_BYTES)?), }) } @@ -152,8 +319,9 @@ struct AwsError { message: String, } -/// Textract answers a multi-page PDF or TIFF with a bare "unsupported document -/// format", which reads like a corrupt file. Say what the limit is. +/// Textract answers both an unsupported format and a multi-page PDF or TIFF +/// with a bare "unsupported document format", which reads like a corrupt file. +/// Say what the synchronous API accepts. pub(super) fn error_class(body: String, status: u16, headers: Vec<(String, String)>) -> Error { let unsupported = serde_json::from_str::(&body) .ok() @@ -162,7 +330,7 @@ pub(super) fn error_class(body: String, status: u16, headers: Vec<(String, Strin status, body: match unsupported { Some(error) => format!( - "{UNSUPPORTED_DOCUMENT}: {}. aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; multi-page documents are not supported", + "{UNSUPPORTED_DOCUMENT}: {}. aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; other formats and multi-page documents are not supported", error.message ), None => body, @@ -178,7 +346,7 @@ pub(super) fn lines_by_page(blocks: &[Block]) -> Vec<(i64, String)> { .map(|page| { let lines: Vec<&str> = blocks .iter() - .filter(|block| block.block_type == "LINE" && block.page() == page) + .filter(|block| block.block_type == BlockType::Line && block.page() == page) .filter_map(|block| block.text.as_deref()) .collect(); (page, lines.join("\n")) @@ -187,61 +355,324 @@ pub(super) fn lines_by_page(blocks: &[Block]) -> Vec<(i64, String)> { .collect() } +pub(super) fn ocr_response( + model: &str, + page_markdown: Vec<(i64, String)>, + document_metadata: Option, +) -> LiteLLMOcrResponse { + let pages: Vec = page_markdown + .into_iter() + .map(|(page, markdown)| OcrPage { + index: page - 1, + markdown, + ..Default::default() + }) + .collect(); + let pages_processed = document_metadata + .and_then(|metadata| metadata.pages) + .or_else(|| i64::try_from(pages.len()).ok()); + LiteLLMOcrResponse { + usage_info: Some(OcrUsageInfo { + pages_processed, + ..Default::default() + }), + ..LiteLLMOcrResponse::new(model, pages) + } +} + #[cfg(test)] mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + use super::*; - #[test] - fn a_multi_page_rejection_names_the_single_page_limit_and_keeps_the_status() { - let error = error_class( - r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#.into(), - 400, - vec![("x-amzn-requestid".into(), "abc".into())], + const HINT: &str = "other formats and multi-page documents are not supported"; + + fn blocks(value: Value) -> Vec { + serde_json::from_value(value).unwrap() + } + + #[rstest] + #[case::detect("detect-document-text", TextractOperation::DetectDocumentText)] + #[case::analyze("analyze-document", TextractOperation::AnalyzeDocument)] + #[case::any_case("Analyze-Document", TextractOperation::AnalyzeDocument)] + fn a_model_names_its_operation(#[case] model: &str, #[case] expected: TextractOperation) { + assert_eq!(TextractOperation::from_model(model).unwrap(), expected); + } + + #[rstest] + #[case::misspelled("analyse-document")] + #[case::operation_name_from_the_api("AnalyzeDocument")] + #[case::operation_litellm_does_not_call("analyze-expense")] + #[case::empty("")] + fn a_model_outside_the_operations_is_refused_with_the_supported_names(#[case] model: &str) { + let error = TextractOperation::from_model(model).unwrap_err(); + + assert_eq!( + error.to_string(), + format!( + "invalid model: aws_textract has no model {model:?} - use one of: detect-document-text, analyze-document" + ) ); + assert_eq!(error.http_status_code(), Some(400)); + } + + #[rstest] + #[case::line("LINE", BlockType::Line)] + #[case::key_value_set("KEY_VALUE_SET", BlockType::KeyValueSet)] + #[case::layout_section_header("LAYOUT_SECTION_HEADER", BlockType::LayoutSectionHeader)] + #[case::layout_key_value("LAYOUT_KEY_VALUE", BlockType::LayoutKeyValue)] + #[case::added_by_textract_later("LAYOUT_SIDEBAR", BlockType::Unknown)] + fn block_type_reads_the_documented_names(#[case] wire: &str, #[case] expected: BlockType) { + let block: Block = serde_json::from_value(json!({"BlockType": wire})).unwrap(); + + assert_eq!(block.block_type, expected); + } + + #[rstest] + #[case::layout_title(BlockType::LayoutTitle, Some(LayoutType::Title))] + #[case::layout_table(BlockType::LayoutTable, Some(LayoutType::Table))] + #[case::table_is_not_layout(BlockType::Table, None)] + #[case::title_is_not_layout(BlockType::Title, None)] + #[case::unknown_is_not_layout(BlockType::Unknown, None)] + fn only_layout_block_types_have_a_layout_type( + #[case] block_type: BlockType, + #[case] expected: Option, + ) { + assert_eq!(block_type.layout(), expected); + } + + #[rstest] + #[case::child_only(json!([{"Type": "CHILD", "Ids": ["a", "b"]}]), vec!["a", "b"])] + #[case::other_relationships_are_skipped( + json!([ + {"Type": "TABLE_TITLE", "Ids": ["t"]}, + {"Type": "CHILD", "Ids": ["a"]}, + {"Type": "MERGED_CELL", "Ids": ["m"]}, + {"Type": "ADDED_LATER", "Ids": ["x"]}, + {"Type": "CHILD", "Ids": ["b"]} + ]), + vec!["a", "b"] + )] + #[case::no_relationships(json!([]), vec![])] + fn children_are_the_ids_of_child_relationships( + #[case] relationships: Value, + #[case] expected: Vec<&str>, + ) { + let block: Block = + serde_json::from_value(json!({"BlockType": "LINE", "Relationships": relationships})) + .unwrap(); + + assert_eq!(block.children().collect::>(), expected); + } + + #[rstest] + #[case::tables("TABLES", Some(FeatureType::Tables))] + #[case::forms("FORMS", Some(FeatureType::Forms))] + #[case::queries("QUERIES", Some(FeatureType::Queries))] + #[case::signatures("SIGNATURES", Some(FeatureType::Signatures))] + #[case::layout("LAYOUT", Some(FeatureType::Layout))] + #[case::lowercase_is_not_a_feature("layout", None)] + #[case::undocumented("HANDWRITING", None)] + fn feature_type_accepts_only_the_documented_values( + #[case] wire: &str, + #[case] expected: Option, + ) { + assert_eq!( + serde_json::from_value::(json!(wire)).ok(), + expected + ); + if let Some(feature) = expected { + assert_eq!(serde_json::to_value(feature).unwrap(), json!(wire)); + } + } + + #[rstest] + #[case::image_url( + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,aGVsbG8=".into(), + extra_fields: Default::default(), + }, + "aGVsbG8=" + )] + #[case::document_url( + OcrDocument::DocumentUrl { + document_url: "data:application/pdf;base64,YWJj".into(), + extra_fields: Default::default(), + }, + "YWJj" + )] + #[case::percent_encoded_data_uri_is_re_encoded_as_base64( + OcrDocument::DocumentUrl { + document_url: "data:,abc".into(), + extra_fields: Default::default(), + }, + "YWJj" + )] + fn document_bytes_are_the_base64_payload_without_the_data_uri_envelope( + #[case] document: OcrDocument, + #[case] expected: &str, + ) { + assert_eq!(document_bytes(&document).unwrap().bytes, expected); + } + + #[rstest] + #[case::remote_url("https://example.com/a.pdf".to_string(), Error::InvalidDataUri)] + #[case::invalid_base64("data:image/png;base64,@@@".to_string(), Error::InvalidDataUri)] + #[case::over_the_sync_limit( + format!("data:,{}", "a".repeat(SYNC_DOCUMENT_MAX_BYTES + 1)), + Error::InlineDocumentTooLarge + )] + fn document_bytes_refuse_what_the_sync_api_cannot_take( + #[case] document_url: String, + #[case] expected: Error, + ) { + let error = document_bytes(&OcrDocument::DocumentUrl { + document_url, + extra_fields: Default::default(), + }) + .unwrap_err(); + + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + } + + #[rstest] + #[case::bare_type( + r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#, + Some("Request has unsupported document format") + )] + #[case::namespaced_type( + r#"{"__type":"com.amazonaws.textract#UnsupportedDocumentException","Message":"bad"}"#, + Some("bad") + )] + #[case::lowercase_message( + r#"{"__type":"UnsupportedDocumentException","message":"bad"}"#, + Some("bad") + )] + #[case::other_exception(r#"{"__type":"AccessDeniedException","Message":"no"}"#, None)] + #[case::json_without_a_type(r#"{"Message":"no"}"#, None)] + #[case::not_json("bad gateway", None)] + fn only_an_unsupported_document_gains_the_sync_api_hint( + #[case] body: &str, + #[case] hinted_message: Option<&str>, + ) { + let response_headers = vec![("x-amzn-requestid".to_string(), "abc".to_string())]; let Error::Provider { status, - body, + body: reported, headers, - } = error + } = error_class(body.into(), 400, response_headers.clone()) else { panic!("expected a provider error"); }; + assert_eq!(status, 400); - assert!(body.contains("Request has unsupported document format")); - assert!(body.contains("single-page PDF or TIFF")); - assert_eq!(headers, vec![("x-amzn-requestid".into(), "abc".into())]); - } - - #[test] - fn a_namespaced_exception_type_is_recognized() { - let Error::Provider { body, .. } = error_class( - r#"{"__type":"com.amazonaws.textract#UnsupportedDocumentException","message":"bad"}"# - .into(), - 400, - Vec::new(), - ) else { - panic!("expected a provider error"); - }; - assert!(body.contains("multi-page documents are not supported")); - } - - #[test] - fn other_provider_errors_pass_through_untouched() { - for body in [ - r#"{"__type":"AccessDeniedException","Message":"no"}"#, - "bad gateway", - ] { - let Error::Provider { - body: reported, - status, - .. - } = error_class(body.into(), 403, Vec::new()) - else { - panic!("expected a provider error"); - }; - assert_eq!(reported, body); - assert_eq!(status, 403); + assert_eq!(headers, response_headers); + match hinted_message { + Some(message) => { + assert!(reported.contains(message), "{reported}"); + assert!(reported.contains(HINT), "{reported}"); + } + None => assert_eq!(reported, body), } } + + #[rstest] + #[case::no_caller_headers(vec![], vec![])] + #[case::unrelated_headers_are_kept(vec![("x-trace", "1")], vec![("x-trace", "1")])] + #[case::a_caller_content_type_is_replaced( + vec![("content-type", "application/json"), ("x-trace", "1")], + vec![("x-trace", "1")] + )] + #[case::a_caller_target_is_replaced( + vec![("X-AMZ-TARGET", "Textract.AnalyzeDocument")], + vec![] + )] + fn operation_headers_are_sent_once( + #[case] extra_headers: Vec<(&str, &str)>, + #[case] kept: Vec<(&str, &str)>, + ) { + let owned = |headers: Vec<(&str, &str)>| -> Vec<(String, String)> { + headers + .into_iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() + }; + + let headers = + operation_headers(&owned(extra_headers), TextractOperation::DetectDocumentText); + + let mut expected = owned(kept); + expected.extend(owned(vec![ + ("X-Amz-Target", "Textract.DetectDocumentText"), + ("Content-Type", "application/x-amz-json-1.1"), + ])); + assert_eq!(headers, expected); + } + + #[rstest] + #[case::words_are_not_repeated( + json!([ + {"BlockType": "PAGE"}, + {"BlockType": "LINE", "Text": "Invoice 12345"}, + {"BlockType": "WORD", "Text": "Invoice"}, + {"BlockType": "WORD", "Text": "12345"}, + {"BlockType": "LINE", "Text": "total 67.89"} + ]), + vec![(1, "Invoice 12345\ntotal 67.89")] + )] + #[case::pages_are_sorted_and_keep_line_order( + json!([ + {"BlockType": "LINE", "Text": "second", "Page": 2}, + {"BlockType": "LINE", "Text": "first", "Page": 1}, + {"BlockType": "LINE", "Text": "also second", "Page": 2} + ]), + vec![(1, "first"), (2, "second\nalso second")] + )] + #[case::a_page_without_lines_is_dropped( + json!([ + {"BlockType": "PAGE", "Page": 1}, + {"BlockType": "LINE", "Text": "only", "Page": 2} + ]), + vec![(2, "only")] + )] + #[case::no_blocks(json!([]), vec![])] + fn lines_are_grouped_by_page(#[case] input: Value, #[case] expected: Vec<(i64, &str)>) { + let pages = lines_by_page(&blocks(input)); + + let pages: Vec<(i64, &str)> = pages + .iter() + .map(|(page, markdown)| (*page, markdown.as_str())) + .collect(); + assert_eq!(pages, expected); + } + + #[rstest] + #[case::metadata_wins(Some(3), Some(3))] + #[case::metadata_without_pages_falls_back_to_the_page_count(None, Some(2))] + fn pages_are_zero_indexed_and_usage_reports_pages_processed( + #[case] metadata_pages: Option, + #[case] expected: Option, + ) { + let response = ocr_response( + "detect-document-text", + vec![(1, "first".into()), (3, "third".into())], + Some(DocumentMetadata { + pages: metadata_pages, + }), + ); + + let pages: Vec<(i64, &str)> = response + .pages + .iter() + .map(|page| (page.index, page.markdown.as_str())) + .collect(); + assert_eq!(pages, vec![(0, "first"), (2, "third")]); + assert_eq!(response.usage_info.unwrap().pages_processed, expected); + } } diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs index 148b7bed439..ad630a1ca4c 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs @@ -2,35 +2,25 @@ use litellm_core_utils::call_arguments::CallArguments; use serde::{Deserialize, Serialize}; use super::common_utils::{ - Block, DocumentMetadata, HEALTH_CHECK_IMAGE_DATA_URI, TextractDocument, TextractEnvironment, - document_bytes, endpoint, environment, error_class, inline_document, lines_by_page, + TextractDocument, TextractEnvironment, TextractOperation, TextractResponse, document_bytes, + endpoint, environment, error_class, health_check_document, inline_document, lines_by_page, + ocr_response, }; use crate::base_llm::ocr::{ error::Error, handler::OcrClient, transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrPage, OcrRequestContext, - OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, decode_and_normalize_response, + BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrRequestContext, OcrResponseFormat, + PreparedOcrRequest, decode_and_normalize_response, }, }; -const DETECT_DOCUMENT_TEXT_TARGET: &str = "Textract.DetectDocumentText"; - #[derive(Debug, Deserialize, Serialize)] pub struct DetectDocumentTextRequest { #[serde(rename = "Document")] pub document: TextractDocument, } -#[derive(Deserialize)] -#[serde(rename_all = "PascalCase")] -pub struct DetectDocumentTextResponse { - #[serde(default)] - blocks: Vec, - document_metadata: Option, -} - -/// Synchronous `DetectDocumentText`: plain lines from one image or single-page document. #[derive(Clone, Copy, Debug, Default)] pub struct TextractDetectTextConfig; @@ -40,10 +30,7 @@ impl BaseOcrConfig for TextractDetectTextConfig { type Environment = TextractEnvironment; fn get_health_check_document(&self) -> OcrDocument { - OcrDocument::ImageUrl { - image_url: HEALTH_CHECK_IMAGE_DATA_URI.into(), - extra_fields: Default::default(), - } + health_check_document() } fn map_ocr_params( @@ -59,7 +46,7 @@ impl BaseOcrConfig for TextractDetectTextConfig { request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - environment(request, DETECT_DOCUMENT_TEXT_TARGET).await + environment(request, TextractOperation::DetectDocumentText).await } fn get_complete_url( @@ -116,48 +103,36 @@ impl BaseOcrConfig for TextractDetectTextConfig { fn normalize_response( model: &str, - response: DetectDocumentTextResponse, + response: TextractResponse, ) -> Result { - let pages: Vec = lines_by_page(&response.blocks) - .into_iter() - .map(|(page, markdown)| OcrPage { - index: page - 1, - markdown, - ..Default::default() - }) - .collect(); - let pages_processed = response - .document_metadata - .and_then(|metadata| metadata.pages) - .or_else(|| i64::try_from(pages.len()).ok()); - Ok(LiteLLMOcrResponse { - usage_info: Some(OcrUsageInfo { - pages_processed, - ..Default::default() - }), - ..LiteLLMOcrResponse::new(model, pages) - }) + Ok(ocr_response( + model, + lines_by_page(&response.blocks), + response.document_metadata, + )) } #[cfg(test)] mod tests { - use serde_json::json; + use rstest::{fixture, rstest}; + use serde_json::{Value, json}; use super::*; - fn normalize(response: serde_json::Value) -> LiteLLMOcrResponse { - TextractDetectTextConfig - .transform_ocr_response( - "detect-document-text", - &serde_json::to_vec(&response).unwrap(), - OcrResponseFormat::Litellm, - ) - .unwrap() + const MODEL: &str = "detect-document-text"; + + #[fixture] + fn document(#[default("data:image/png;base64,aGVsbG8=")] source: &str) -> OcrDocument { + OcrDocument::DocumentUrl { + document_url: source.into(), + extra_fields: Default::default(), + } } - #[test] - fn lines_become_one_markdown_page_and_words_are_not_repeated() { - let response = normalize(json!({ + #[rstest] + #[case::one_page_without_page_numbers( + json!({ + "DetectDocumentTextModelVersion": "1.0", "DocumentMetadata": {"Pages": 1}, "Blocks": [ {"BlockType": "PAGE"}, @@ -166,35 +141,90 @@ mod tests { {"BlockType": "WORD", "Text": "12345"}, {"BlockType": "LINE", "Text": "total 67.89"} ] - })); - - assert_eq!(response.pages.len(), 1); - assert_eq!(response.pages[0].index, 0); - assert_eq!(response.pages[0].markdown, "Invoice 12345\ntotal 67.89"); - assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); - } - - #[test] - fn lines_are_grouped_by_their_page_in_page_order() { - let response = normalize(json!({ + }), + vec![(0, "Invoice 12345\ntotal 67.89")], + Some(1) + )] + #[case::pages_out_of_order( + json!({ "DocumentMetadata": {"Pages": 2}, "Blocks": [ {"BlockType": "LINE", "Text": "second", "Page": 2}, {"BlockType": "LINE", "Text": "first", "Page": 1}, {"BlockType": "LINE", "Text": "also second", "Page": 2} ] - })); + }), + vec![(0, "first"), (1, "second\nalso second")], + Some(2) + )] + #[case::missing_metadata_counts_the_pages_with_text( + json!({"Blocks": [{"BlockType": "LINE", "Text": "only"}]}), + vec![(0, "only")], + Some(1) + )] + #[case::blank_document(json!({"DocumentMetadata": {"Pages": 1}}), vec![], Some(1))] + fn response_lines_become_one_markdown_page_per_document_page( + #[case] raw_response: Value, + #[case] expected_pages: Vec<(i64, &str)>, + #[case] expected_pages_processed: Option, + ) { + let response = TextractDetectTextConfig + .transform_ocr_response( + MODEL, + &serde_json::to_vec(&raw_response).unwrap(), + OcrResponseFormat::Litellm, + ) + .unwrap(); let pages: Vec<(i64, &str)> = response .pages .iter() .map(|page| (page.index, page.markdown.as_str())) .collect(); - assert_eq!(pages, vec![(0, "first"), (1, "second\nalso second")]); + assert_eq!(pages, expected_pages); + assert_eq!(response.model, MODEL); + assert_eq!( + response.usage_info.unwrap().pages_processed, + expected_pages_processed + ); } - #[test] - fn a_multi_page_rejection_is_explained_to_the_caller() { + #[rstest] + fn the_request_is_only_the_document_bytes(document: OcrDocument) { + let request = TextractDetectTextConfig + .transform_ocr_request(MODEL, document, &(), &[]) + .unwrap(); + + assert_eq!( + serde_json::to_value(request).unwrap(), + json!({"Document": {"Bytes": "aGVsbG8="}}) + ); + } + + #[rstest] + fn a_remote_url_is_refused_by_the_sync_transform( + #[with("https://example.com/a.pdf")] document: OcrDocument, + ) { + let error = TextractDetectTextConfig + .transform_ocr_request(MODEL, document, &(), &[]) + .unwrap_err(); + + assert!(matches!(error, Error::InvalidDataUri)); + } + + #[rstest] + fn the_health_check_document_is_an_inline_image_the_request_accepts() { + let document = TextractDetectTextConfig.get_health_check_document(); + + assert!( + TextractDetectTextConfig + .transform_ocr_request(MODEL, document, &(), &[]) + .is_ok() + ); + } + + #[rstest] + fn provider_errors_go_through_the_shared_textract_error_class() { let error = TextractDetectTextConfig.get_error_class( r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#.into(), 400, @@ -207,41 +237,4 @@ mod tests { .contains("multi-page documents are not supported") ); } - - #[test] - fn the_request_carries_the_document_bytes_without_the_data_uri_envelope() { - let request = TextractDetectTextConfig - .transform_ocr_request( - "detect-document-text", - OcrDocument::ImageUrl { - image_url: "data:image/png;base64,aGVsbG8=".into(), - extra_fields: Default::default(), - }, - &(), - &[], - ) - .unwrap(); - - assert_eq!( - serde_json::to_value(request).unwrap(), - json!({"Document": {"Bytes": "aGVsbG8="}}) - ); - } - - #[test] - fn a_remote_url_is_refused_by_the_sync_transform() { - let error = TextractDetectTextConfig - .transform_ocr_request( - "detect-document-text", - OcrDocument::DocumentUrl { - document_url: "https://example.com/a.pdf".into(), - extra_fields: Default::default(), - }, - &(), - &[], - ) - .unwrap_err(); - - assert!(matches!(error, Error::InvalidDataUri)); - } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 08f217351ba..e09842e2856 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -76,6 +76,12 @@ pub enum Error { Unsupported(&'static str), #[error("invalid provider: {0}")] InvalidProvider(String), + #[error("invalid model: {provider} has no model {model:?} - use one of: {}", supported.join(", "))] + InvalidModel { + provider: &'static str, + model: String, + supported: &'static [&'static str], + }, #[error("invalid request: {0}")] InvalidRequest(String), #[error("invalid response: {0}")] @@ -155,6 +161,7 @@ impl Error { | Self::DotModel | Self::InvalidRequest(_) | Self::InvalidProvider(_) + | Self::InvalidModel { .. } | Self::Params(_) | Self::Headers(_) | Self::Http(_)