scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

fix stuff

Browse source
This commit is contained in:
Yujong Lee 2026-09-19 09:32:46 -07:00
parent c404bed9f0
commit eb502824f0
8 changed files with 850 additions and 356 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-rust/Cargo.lock generated
View file

@ -2176,6 +2176,7 @@ dependencies = [
"serde_json",
"serde_path_to_error",
"serde_with",
"strum",
"thiserror 2.0.19",
"time",
"tokio",

5
litellm-rust/crates/auth-aws/src/signer.rs
View file

@ -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<Vec<(String, String)>, litellm_http::Error> {
// Sending a caller's copy next to the computed one is rejected by AWS.
if let Some((name, _)) = request
.headers
.iter()

23
litellm-rust/crates/core/src/ocr/provider_config.rs
View file

@ -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::<OcrProvider>()
.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)]

1
litellm-rust/crates/llms/Cargo.toml
View file

@ -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"] }

409
litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs
View file

@ -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<Vec<String>>,
pub feature_types: Option<Vec<FeatureType>>,
}
#[derive(Debug, Deserialize, Serialize)]
@ -29,18 +29,9 @@ pub struct AnalyzeDocumentRequest {
#[serde(rename = "Document")]
pub document: TextractDocument,
#[serde(rename = "FeatureTypes")]
pub feature_types: Vec<String>,
pub feature_types: Vec<FeatureType>,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct AnalyzeDocumentResponse {
#[serde(default)]
blocks: Vec<Block>,
document_metadata: Option<DocumentMetadata>,
}
/// 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<TextractEnvironment, Error> {
environment(request, ANALYZE_DOCUMENT_TARGET).await
environment(request, TextractOperation::AnalyzeDocument).await
}
fn get_complete_url(
@ -94,12 +82,10 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig {
) -> Result<AnalyzeDocumentRequest, Error> {
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<LiteLLMOcrResponse, Error> {
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<OcrPage> = 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<String> = 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::<Vec<_>>()
.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(), &params, &[])
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, &params, &[])
.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()
);
}
}

559
litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs
View file

@ -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/<model>` 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<Self, Error> {
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<LayoutType> {
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<String>,
pub page: Option<i64>,
pub row_index: Option<usize>,
@ -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<String>,
}
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<Item = &str> {
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<i64>,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct TextractResponse {
#[serde(default)]
pub(super) blocks: Vec<Block>,
pub(super) document_metadata: Option<DocumentMetadata>,
}
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<TextractEnvironment, Error> {
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<TextractDocument, Error> {
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::<AwsError>(&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<DocumentMetadata>,
) -> LiteLLMOcrResponse {
let pages: Vec<OcrPage> = 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<Block> {
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<LayoutType>,
) {
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::<Vec<_>>(), 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<FeatureType>,
) {
assert_eq!(
serde_json::from_value::<FeatureType>(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("<html>bad gateway</html>", 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"}"#,
"<html>bad gateway</html>",
] {
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<i64>,
#[case] expected: Option<i64>,
) {
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);
}
}

201
litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs
View file

@ -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<Block>,
document_metadata: Option<DocumentMetadata>,
}
/// 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<TextractEnvironment, Error> {
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<LiteLLMOcrResponse, Error> {
let pages: Vec<OcrPage> = 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<i64>,
) {
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));
}
}

7
litellm-rust/crates/llms/src/base_llm/ocr/error.rs
View file

@ -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(_)