use std::collections::HashMap;
use std::fmt::Write as _;
use std::fs;
use std::ops::Range;
use std::path::Path;
use anyhow::{Context, Result, bail};
use pdf_oxide::{Destination, OutlineItem, PdfDocument, RegionRole};
use serde::Serialize;
use crate::engine::document::{
Asset, BoundingBox, Document, DocumentFormat, Node, NodeKind, SourceAnchor,
};
use crate::engine::output::ImageMode;
use crate::engine::vision::{Analysis, DetectedObject, Request};
#[derive(Clone, Copy, Debug, Serialize)]
pub(crate) struct PdfInfo {
pub(crate) pages: usize,
}
#[derive(Debug, Serialize)]
pub(crate) struct ReadPdfOutput {
format: &'static str,
pages: usize,
markdown: String,
images: Vec<PdfImageOutput>,
}
#[derive(Debug, Serialize)]
struct PdfImageOutput {
page: usize,
width: u32,
height: u32,
mime: &'static str,
mode: ImageMode,
#[serde(skip_serializing_if = "Option::is_none")]
caption: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
objects: Option<Vec<DetectedObject>>,
#[serde(skip_serializing_if = "Option::is_none")]
ocr: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
encoding: Option<&'static str>,
#[serde(skip_serializing_if = "Option::is_none")]
data: Option<String>,
}
struct StructuralSectionState {
node_id: String,
node_index: usize,
headings: Vec<(u8, String)>,
}
#[derive(Clone, Copy)]
struct RegionParents<'a> {
page_id: &'a str,
base_parent_id: &'a str,
}
pub(crate) fn probe(bytes: &[u8]) -> Result<PdfInfo> {
let document = PdfDocument::from_bytes(bytes.to_vec()).context("parse PDF")?;
let pages = document.page_count().context("read PDF page count")?;
Ok(PdfInfo { pages })
}
pub(crate) fn extract_document(
path: &Path,
bytes: &[u8],
id: String,
assets_dir: &Path,
) -> Result<Document> {
let pdf = PdfDocument::from_bytes(bytes.to_vec()).context("parse PDF")?;
let pages = pdf.page_count().context("read PDF page count")?;
let mut document = document_outline(path, &pdf, id, pages)?;
let mut headings = Vec::new();
let mut sections = HashMap::new();
for page_index in 0..pages {
let page = page_index + 1;
let structured = pdf
.extract_structured(page_index)
.with_context(|| format!("extract structure from PDF page {page}"))?;
append_page_nodes(
&document.id,
page,
&structured,
&mut headings,
&mut sections,
&mut document.nodes,
);
append_page_assets(
&pdf,
&document.id,
page_index,
page,
assets_dir,
&mut document.assets,
)?;
}
Ok(document)
}
fn append_page_assets(
pdf: &PdfDocument,
document_id: &str,
page_index: usize,
page: usize,
assets_dir: &Path,
assets: &mut Vec<Asset>,
) -> Result<()> {
let handles = pdf
.page_image_handles(page_index)
.with_context(|| format!("enumerate images on PDF page {page}"))?;
if handles.is_empty() {
return Ok(());
}
fs::create_dir_all(assets_dir)
.with_context(|| format!("create PDF asset directory {}", assets_dir.display()))?;
for handle in handles {
let image = match handle.decode() {
Ok(image) => image,
Err(error) => {
log::warn!(
"skip undecodable PDF image {} on page {page}: {error}",
handle.paint_order
);
continue;
}
};
let png = match image.to_png_bytes() {
Ok(png) if !png.is_empty() => png,
Ok(_) => continue,
Err(error) => {
log::warn!(
"skip PDF image {} on page {page}: {error}",
handle.paint_order
);
continue;
}
};
let id = asset_id(document_id, page, handle.paint_order);
let path = assets_dir.join(format!("{id}.png"));
crate::engine::repo::write_atomic(&path, &png)
.with_context(|| format!("write PDF asset {}", path.display()))?;
assets.push(Asset {
id,
mime: "image/png".to_owned(),
path,
source_anchor: Some(SourceAnchor {
page: Some(page),
destination: None,
bbox: Some(BoundingBox {
x: handle.bbox.x,
y: handle.bbox.y,
width: handle.bbox.width,
height: handle.bbox.height,
}),
}),
});
}
Ok(())
}
fn asset_id(document_id: &str, page: usize, paint_order: usize) -> String {
let digest = blake3::hash(format!("{document_id}:page:{page}:image:{paint_order}").as_bytes());
format!("a_{}", &digest.to_hex()[..16])
}
pub(crate) fn extract_outline_document(
path: &std::path::Path,
bytes: &[u8],
id: String,
) -> Result<Document> {
let pdf = PdfDocument::from_bytes(bytes.to_vec()).context("parse PDF")?;
let pages = pdf.page_count().context("read PDF page count")?;
document_outline(path, &pdf, id, pages)
}
fn document_outline(
path: &std::path::Path,
pdf: &PdfDocument,
id: String,
pages: usize,
) -> Result<Document> {
let outline = pdf
.get_outline()
.context("read PDF outline")?
.unwrap_or_default();
let mut nodes = Vec::new();
append_outline_nodes(&id, &outline, None, "", &mut nodes);
let mut document = Document {
id,
format: DocumentFormat::Pdf,
source: std::path::PathBuf::new(),
title: String::new(),
pages,
nodes,
assets: Vec::new(),
};
document.rebind_source(path);
Ok(document)
}
fn append_outline_nodes(
document_id: &str,
items: &[OutlineItem],
parent_id: Option<&str>,
parent_position: &str,
nodes: &mut Vec<Node>,
) {
for (index, item) in items.iter().enumerate() {
let position = if parent_position.is_empty() {
index.to_string()
} else {
format!("{parent_position}.{index}")
};
let id = node_id(document_id, &position);
let source_anchor = item.dest.as_ref().map(|destination| match destination {
Destination::PageIndex(page) => SourceAnchor {
page: Some(page + 1),
destination: None,
bbox: None,
},
Destination::Named(destination) => SourceAnchor {
page: None,
destination: Some(destination.clone()),
bbox: None,
},
});
nodes.push(Node {
id: id.clone(),
parent_id: parent_id.map(str::to_owned),
kind: NodeKind::Section,
title: Some(item.title.clone()),
text: None,
level: None,
column: None,
source_anchor,
});
append_outline_nodes(document_id, &item.children, Some(&id), &position, nodes);
}
}
fn append_page_nodes(
document_id: &str,
page: usize,
structured: &pdf_oxide::StructuredPage,
headings: &mut Vec<(u8, String)>,
sections: &mut HashMap<usize, StructuralSectionState>,
nodes: &mut Vec<Node>,
) {
let page_id = node_id(document_id, &format!("page:{page}"));
nodes.push(Node {
id: page_id.clone(),
parent_id: None,
kind: NodeKind::Page,
title: Some(format!("Page {page}")),
text: None,
level: None,
column: None,
source_anchor: Some(SourceAnchor {
page: Some(page),
destination: None,
bbox: Some(BoundingBox {
x: 0.0,
y: 0.0,
width: structured.page_width,
height: structured.page_height,
}),
}),
});
for (position, region) in structured.regions.iter().enumerate() {
if let Some(section_id) = region.section_id {
let state = sections.entry(section_id).or_insert_with(|| {
let node_id = node_id(document_id, &format!("section:{section_id}"));
let node_index = nodes.len();
let title = match region.kind {
RegionRole::StructuralHeading { .. } => region.text.trim().to_owned(),
_ => format!("Section {section_id}"),
};
nodes.push(Node {
id: node_id.clone(),
parent_id: None,
kind: NodeKind::StructuralSection,
title: Some(title),
text: None,
level: None,
column: None,
source_anchor: Some(SourceAnchor {
page: Some(page),
destination: None,
bbox: Some(BoundingBox {
x: region.bbox.x,
y: region.bbox.y,
width: region.bbox.width,
height: region.bbox.height,
}),
}),
});
StructuralSectionState {
node_id,
node_index,
headings: Vec::new(),
}
});
if matches!(region.kind, RegionRole::StructuralHeading { .. })
&& nodes[state.node_index]
.title
.as_deref()
.is_some_and(|title| title.starts_with("Section "))
{
nodes[state.node_index].title = Some(region.text.trim().to_owned());
}
append_region_node(
document_id,
page,
position,
region,
RegionParents {
page_id: &page_id,
base_parent_id: &state.node_id,
},
&mut state.headings,
nodes,
);
} else {
append_region_node(
document_id,
page,
position,
region,
RegionParents {
page_id: &page_id,
base_parent_id: &page_id,
},
headings,
nodes,
);
}
}
}
fn append_region_node(
document_id: &str,
page: usize,
position: usize,
region: &pdf_oxide::StructuredRegion,
parents: RegionParents<'_>,
headings: &mut Vec<(u8, String)>,
nodes: &mut Vec<Node>,
) {
let text = region.text.trim();
if text.is_empty() {
return;
}
let (kind, level) = match region.kind {
RegionRole::BodyBlock => (NodeKind::Paragraph, None),
RegionRole::StructuralHeading { level } => (NodeKind::Heading, Some(level)),
RegionRole::MarginalLabel => (NodeKind::MarginalLabel, None),
RegionRole::Header => (NodeKind::Header, None),
RegionRole::Footer => (NodeKind::Footer, None),
RegionRole::PageNumber => (NodeKind::PageNumber, None),
RegionRole::Artifact => (NodeKind::Artifact, None),
};
let id = node_id(document_id, &format!("page:{page}:region:{position}"));
let parent_id = if let Some(level) = level {
while headings
.last()
.is_some_and(|(parent_level, _)| *parent_level >= level)
{
headings.pop();
}
headings
.last()
.map_or_else(|| parents.base_parent_id.to_owned(), |(_, id)| id.clone())
} else if matches!(kind, NodeKind::Paragraph | NodeKind::MarginalLabel) {
headings
.last()
.map_or_else(|| parents.base_parent_id.to_owned(), |(_, id)| id.clone())
} else {
parents.page_id.to_owned()
};
let (title, text) = if kind == NodeKind::Heading {
(Some(text.to_owned()), None)
} else {
(None, Some(text.to_owned()))
};
nodes.push(Node {
id: id.clone(),
parent_id: Some(parent_id),
kind,
title,
text,
level,
column: region.column_index,
source_anchor: Some(SourceAnchor {
page: Some(page),
destination: None,
bbox: Some(BoundingBox {
x: region.bbox.x,
y: region.bbox.y,
width: region.bbox.width,
height: region.bbox.height,
}),
}),
});
if let Some(level) = level {
headings.push((level, id));
}
}
fn node_id(document_id: &str, position: &str) -> String {
let digest = blake3::hash(format!("{document_id}\0{position}").as_bytes()).to_string();
format!("n_{}", &digest[..16])
}
fn selected_page_range(page: Option<usize>, pages: usize) -> Result<Range<usize>> {
let Some(page) = page else {
return Ok(0..pages);
};
if page == 0 {
bail!("PDF page must be greater than zero");
}
if page > pages {
bail!("PDF page {page} is past the end of the document ({pages} pages)");
}
Ok(page - 1..page)
}
pub(crate) fn read(
bytes: &[u8],
mode: ImageMode,
page: Option<usize>,
mut analyze: impl FnMut(&[u8], Request) -> Analysis,
) -> Result<ReadPdfOutput> {
let document = PdfDocument::from_bytes(bytes.to_vec()).context("parse PDF")?;
let pages = document.page_count().context("read PDF page count")?;
let mut markdown = String::new();
let mut images = Vec::new();
let page_range = selected_page_range(page, pages)?;
for page_index in page_range {
let page = page_index + 1;
let text = document
.extract_text(page_index)
.with_context(|| format!("extract text from PDF page {page}"))?;
if !markdown.is_empty() {
markdown.push('\n');
}
writeln!(markdown, "<!-- readseek:page {page} -->").unwrap();
let text = text.trim();
if !text.is_empty() {
markdown.push_str(text);
markdown.push('\n');
}
let handles = document
.page_image_handles(page_index)
.with_context(|| format!("enumerate images on PDF page {page}"))?;
for handle in handles {
let width = handle.width;
let height = handle.height;
let image = match handle.decode() {
Ok(image) => image,
Err(error) => {
log::warn!("PDF page {page} image skipped: {error}");
continue;
}
};
let png = match image.to_png_bytes() {
Ok(png) => png,
Err(error) => {
log::warn!("PDF page {page} image skipped: {error}");
continue;
}
};
let request = Request {
caption: matches!(mode, ImageMode::All | ImageMode::Caption),
objects: matches!(mode, ImageMode::All | ImageMode::Objects),
ocr: matches!(mode, ImageMode::All | ImageMode::Ocr),
};
let analysis = if mode == ImageMode::None {
Analysis::default()
} else {
analyze(&png, request)
};
let prepared = match (mode == ImageMode::None)
.then(|| crate::engine::image::preprocess(&png))
.transpose()
{
Ok(prepared) => prepared,
Err(error) => {
log::warn!("PDF page {page} image skipped: {error}");
continue;
}
};
let (mime, encoding, data) = if let Some(prepared) = prepared {
(prepared.mime, Some("base64"), Some(prepared.data))
} else {
("image/png", None, None)
};
let (caption, objects, ocr) = match mode {
ImageMode::None => (None, None, None),
ImageMode::All => (analysis.caption, analysis.objects, analysis.ocr),
ImageMode::Caption => (analysis.caption, None, None),
ImageMode::Objects => (None, analysis.objects, None),
ImageMode::Ocr => (None, None, analysis.ocr),
};
images.push(PdfImageOutput {
page,
width,
height,
mime,
mode,
caption,
objects,
ocr,
encoding,
data,
});
}
}
Ok(ReadPdfOutput {
format: "pdf",
pages,
markdown,
images,
})
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use pdf_oxide::geometry::Rect;
use pdf_oxide::{RegionRole, StructuredPage, StructuredRegion};
use super::{StructuralSectionState, append_page_nodes, selected_page_range};
use crate::engine::document::{Node, NodeKind};
#[test]
fn selects_one_based_pdf_pages() {
assert_eq!(selected_page_range(None, 4).unwrap(), 0..4);
assert_eq!(selected_page_range(Some(3), 4).unwrap(), 2..3);
assert!(selected_page_range(Some(0), 4).is_err());
assert!(selected_page_range(Some(5), 4).is_err());
}
fn region(kind: RegionRole, text: &str, section_id: Option<usize>) -> StructuredRegion {
StructuredRegion {
kind,
text: text.to_owned(),
bbox: Rect::new(1.0, 2.0, 3.0, 4.0),
spans: Vec::new(),
column_index: None,
section_id,
}
}
fn page(page_index: usize, regions: Vec<StructuredRegion>) -> StructuredPage {
StructuredPage {
page_index,
page_width: 100.0,
page_height: 200.0,
regions,
}
}
#[test]
fn structural_section_spans_pages() {
let mut nodes = Vec::new();
let mut headings = Vec::new();
let mut sections: HashMap<usize, StructuralSectionState> = HashMap::new();
append_page_nodes(
"document",
1,
&page(0, vec![region(RegionRole::BodyBlock, "first", Some(7))]),
&mut headings,
&mut sections,
&mut nodes,
);
append_page_nodes(
"document",
2,
&page(1, vec![region(RegionRole::BodyBlock, "second", Some(7))]),
&mut headings,
&mut sections,
&mut nodes,
);
let section = nodes
.iter()
.find(|node| node.kind == NodeKind::StructuralSection)
.expect("structural section");
let paragraphs: Vec<&Node> = nodes
.iter()
.filter(|node| node.kind == NodeKind::Paragraph)
.collect();
assert_eq!(paragraphs.len(), 2);
assert!(
paragraphs
.iter()
.all(|node| node.parent_id.as_deref() == Some(section.id.as_str()))
);
}
#[test]
fn untagged_heading_spans_pages() {
let mut nodes = Vec::new();
let mut headings = Vec::new();
let mut sections = HashMap::new();
append_page_nodes(
"document",
1,
&page(
0,
vec![region(
RegionRole::StructuralHeading { level: 1 },
"Heading",
None,
)],
),
&mut headings,
&mut sections,
&mut nodes,
);
append_page_nodes(
"document",
2,
&page(1, vec![region(RegionRole::BodyBlock, "body", None)]),
&mut headings,
&mut sections,
&mut nodes,
);
let heading = nodes
.iter()
.find(|node| node.kind == NodeKind::Heading)
.expect("heading");
let body = nodes
.iter()
.find(|node| node.kind == NodeKind::Paragraph)
.expect("body");
assert_eq!(body.parent_id.as_deref(), Some(heading.id.as_str()));
}
}