uni_xervo/traits/
docs.rs

1//! Document understanding types and traits — VLM-based extraction and OCR.
2
3use crate::error::Result;
4use crate::traits::{ImageInput, ModelInfo};
5use async_trait::async_trait;
6
7/// Options for a [`DocumentExtractionModel::extract`] call.
8#[derive(Debug, Clone)]
9pub struct DocExtractOptions {
10    /// Desired primary output format. The `plain_markdown` field on
11    /// [`DocExtractResult`] is always populated; this controls the
12    /// in-block `content` string format for non-text block kinds.
13    pub output: DocOutputFormat,
14    /// Whether to extract table blocks (some models skip tables for speed).
15    pub include_tables: bool,
16    /// Whether to extract math formulas (some models skip formulas for speed).
17    pub include_formulas: bool,
18    /// Whether to populate per-block bounding boxes (extra work for some VLMs).
19    pub include_bboxes: bool,
20}
21
22/// Output format for the in-block `content` string of [`DocBlock`].
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24#[non_exhaustive]
25pub enum DocOutputFormat {
26    /// Markdown — the default for text-shaped blocks.
27    Markdown,
28    /// JSON — structured representation for downstream parsing.
29    Json,
30    /// HTML — useful for tables and rich layout preservation.
31    Html,
32}
33
34/// Per-page result of a document extraction call.
35#[derive(Debug, Clone)]
36pub struct DocExtractResult {
37    /// Structured blocks in reading order.
38    pub blocks: Vec<DocBlock>,
39    /// Concatenated Markdown view (convenience field, always populated).
40    pub plain_markdown: String,
41}
42
43/// One block within a [`DocExtractResult`].
44#[derive(Debug, Clone)]
45pub struct DocBlock {
46    /// Kind of block.
47    pub kind: DocBlockKind,
48    /// Block content. Markdown for text / heading / table; LaTeX for formula;
49    /// description for figure.
50    pub content: String,
51    /// Bounding box in page coordinates `[x0, y0, x1, y1]`, when requested
52    /// via [`DocExtractOptions::include_bboxes`] and supported by the model.
53    pub bbox: Option<[f32; 4]>,
54    /// Monotonically increasing index within the page, in reading order.
55    pub reading_order: u32,
56}
57
58/// Kind of block returned by a [`DocumentExtractionModel`].
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60#[non_exhaustive]
61pub enum DocBlockKind {
62    /// Paragraph or running text.
63    Text,
64    /// Heading (any level).
65    Heading,
66    /// List item.
67    List,
68    /// Tabular data.
69    Table,
70    /// Figure or image with a caption / description.
71    Figure,
72    /// Math formula.
73    Formula,
74    /// Figure / table caption.
75    Caption,
76    /// Page footer (e.g. page number, footnote).
77    Footer,
78    /// Page header (e.g. running title).
79    Header,
80}
81
82/// Per-image result of an OCR call.
83#[derive(Debug, Clone)]
84pub struct OcrResult {
85    /// Per-word or per-line text blocks with bounding boxes.
86    pub blocks: Vec<OcrBlock>,
87    /// Concatenated plain text (convenience field, always populated).
88    pub plain_text: String,
89}
90
91/// One text region within an [`OcrResult`].
92#[derive(Debug, Clone)]
93pub struct OcrBlock {
94    /// Recognized text.
95    pub text: String,
96    /// Bounding box in image coordinates `[x0, y0, x1, y1]`.
97    pub bbox: [f32; 4],
98    /// Provider-reported confidence in `[0.0, 1.0]`.
99    pub confidence: f32,
100}
101
102/// A vision-language model that extracts structured blocks from
103/// document page images.
104///
105/// Targets MinerU-2.5, Granite-Docling, olmOCR-2, and similar VLMs that
106/// understand page layout, tables, figures, and math. The caller is
107/// responsible for rendering PDF pages to images upstream — uni-xervo
108/// providers do not own PDF rendering.
109#[async_trait]
110pub trait DocumentExtractionModel: ModelInfo {
111    /// Extract structured blocks from one or more page images.
112    ///
113    /// Returns one [`DocExtractResult`] per page, in input order.
114    ///
115    /// # Errors
116    /// Returns an error if the provider rejects an input or fails internally.
117    async fn extract(
118        &self,
119        pages: Vec<ImageInput>,
120        options: DocExtractOptions,
121    ) -> Result<Vec<DocExtractResult>>;
122
123    /// Optional warmup hook. The default is a no-op.
124    ///
125    /// # Errors
126    /// Returns an error if the underlying model fails to initialize.
127    async fn warmup(&self) -> Result<()> {
128        Ok(())
129    }
130}
131
132/// A model that recognizes text in images (optical character recognition).
133///
134/// OCR is the document-blind counterpart to
135/// [`DocumentExtractionModel`] — read text out of an image without
136/// interpreting layout. Suited to EasyOCR, PaddleOCR, Tesseract-style models.
137#[async_trait]
138pub trait OcrModel: ModelInfo {
139    /// Recognize text in a batch of images.
140    ///
141    /// Returns one [`OcrResult`] per input image.
142    ///
143    /// # Errors
144    /// Returns an error if the provider rejects an input or fails internally.
145    async fn recognize(&self, images: Vec<ImageInput>) -> Result<Vec<OcrResult>>;
146
147    /// Optional warmup hook. The default is a no-op.
148    ///
149    /// # Errors
150    /// Returns an error if the underlying model fails to initialize.
151    async fn warmup(&self) -> Result<()> {
152        Ok(())
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159    use std::sync::Arc;
160
161    #[test]
162    fn document_extraction_model_is_dyn_safe() {
163        fn _accept(_: Arc<dyn DocumentExtractionModel>) {}
164    }
165
166    #[test]
167    fn ocr_model_is_dyn_safe() {
168        fn _accept(_: Arc<dyn OcrModel>) {}
169    }
170}
uni_xervo/traits/docs.rs

uni_xervo/traits/
docs.rs
