aegis_tools/
document.rs

1//! # ReadDocumentTool
2//!
3//! Local, pure-Rust extraction of text/Markdown from documents:
4//! - **PDF** (`.pdf`)          — digital text via `pdf-extract` (no OCR)
5//! - **Excel** (`.xlsx/.xls/.ods`) — cells → Markdown tables via `calamine`
6//! - **Word** (`.docx`)        — paragraph text via `dotext`
7//! - **PowerPoint** (`.pptx`)  — per-slide text via `dotext`
8//!
9//! No external runtime is required (no JVM, Python or browser), so this fits the
10//! 1c1g / lite profile. Scanned PDFs, complex tables, OCR and formula recovery
11//! are out of scope here — those belong to the opt-in `doc_extract_pro` tool
12//! (external `opendataloader-pdf`).
13
14use crate::registry::{Tool, ToolContext};
15use aegis_security::check_path;
16use anyhow::Result;
17use async_trait::async_trait;
18use serde_json::{json, Value};
19use std::path::Path;
20
21/// Default output cap (characters). 0 disables truncation.
22const DEFAULT_MAX_CHARS: usize = 16_000;
23/// Refuse to load files bigger than this to protect memory on small hosts.
24const MAX_FILE_BYTES: u64 = 50 * 1024 * 1024; // 50 MiB
25
26/// Reads PDF / Word / Excel / PowerPoint into text or Markdown.
27pub struct ReadDocumentTool;
28
29impl ReadDocumentTool {
30    /// Create a new `ReadDocumentTool`.
31    pub fn new() -> Self {
32        ReadDocumentTool
33    }
34}
35
36impl Default for ReadDocumentTool {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42#[async_trait]
43impl Tool for ReadDocumentTool {
44    fn name(&self) -> &str {
45        "read_document"
46    }
47
48    fn description(&self) -> &str {
49        "Extract text/Markdown from a local PDF, Word (.docx), Excel (.xlsx/.xls/.ods) or PowerPoint (.pptx) file. Pure-Rust, no OCR (use doc_extract_pro for scanned PDFs)."
50    }
51
52    fn parameters(&self) -> Value {
53        json!({
54            "type": "object",
55            "properties": {
56                "path": { "type": "string", "description": "Path to the document (within the working directory)" },
57                "format": { "type": "string", "enum": ["markdown", "text"], "description": "Output format (default: markdown)" },
58                "pages": { "type": "string", "description": "PDF only: page range like '1-5,8' (1-indexed). Default: all pages." },
59                "sheet": { "type": "string", "description": "Excel only: sheet name to extract. Default: all sheets." },
60                "max_chars": { "type": "integer", "description": "Truncate output to this many characters (default 16000; 0 = no limit)" }
61            },
62            "required": ["path"]
63        })
64    }
65
66    async fn execute(&self, args: Value, ctx: &ToolContext<'_>) -> Result<String> {
67        let path_str = args["path"].as_str().unwrap_or("").trim();
68        if path_str.is_empty() {
69            return Ok("Error: path is required".to_string());
70        }
71        let as_markdown = args["format"].as_str().unwrap_or("markdown") == "markdown";
72        let pages = args["pages"].as_str().map(|s| s.to_string());
73        let sheet = args["sheet"].as_str().map(|s| s.to_string());
74        let max_chars = args["max_chars"]
75            .as_u64()
76            .map(|v| v as usize)
77            .unwrap_or(DEFAULT_MAX_CHARS);
78
79        // Path safety: confine to the working directory.
80        let safe_path = check_path(path_str, &ctx.cwd)?;
81        if !safe_path.exists() {
82            anyhow::bail!("File not found: {path_str}");
83        }
84        if let Ok(meta) = std::fs::metadata(&safe_path) {
85            if meta.len() > MAX_FILE_BYTES {
86                anyhow::bail!(
87                    "File too large ({} bytes, limit {}). Try a smaller file or use doc_extract_pro.",
88                    meta.len(),
89                    MAX_FILE_BYTES
90                );
91            }
92        }
93
94        let ext = safe_path
95            .extension()
96            .and_then(|e| e.to_str())
97            .unwrap_or("")
98            .to_lowercase();
99
100        // CPU-bound parsing on a blocking thread so we don't stall the runtime.
101        let path_owned = safe_path.to_path_buf();
102        let ext_c = ext.clone();
103        let result = tokio::task::spawn_blocking(move || {
104            extract_document(&path_owned, &ext_c, as_markdown, pages, sheet)
105        })
106        .await
107        .map_err(|e| anyhow::anyhow!("document parsing task failed: {e}"))?;
108
109        let content = match result {
110            Ok(c) => c,
111            Err(e) => return Ok(format!("Failed to read {path_str}: {e}")),
112        };
113
114        if content.trim().is_empty() {
115            return Ok(format!(
116                "No extractable text found in {path_str} (it may be a scanned/image-only document — try doc_extract_pro with OCR)."
117            ));
118        }
119
120        if max_chars > 0 && content.chars().count() > max_chars {
121            let truncated: String = content.chars().take(max_chars).collect();
122            Ok(format!(
123                "{truncated}\n\n... [content truncated at {max_chars} characters. Narrow with 'pages'/'sheet' or raise max_chars.]"
124            ))
125        } else {
126            Ok(content)
127        }
128    }
129}
130
131/// Dispatch by extension. Runs synchronously (call from a blocking thread).
132fn extract_document(
133    path: &Path,
134    ext: &str,
135    markdown: bool,
136    pages: Option<String>,
137    sheet: Option<String>,
138) -> Result<String> {
139    match ext {
140        "pdf" => extract_pdf(path, markdown, pages.as_deref()),
141        "xlsx" | "xls" | "xlsm" | "xlsb" | "ods" => extract_spreadsheet(path, sheet.as_deref()),
142        "docx" => extract_docx(path),
143        "pptx" => extract_pptx(path, markdown),
144        other => anyhow::bail!(
145            "Unsupported document type '.{other}'. Supported: pdf, docx, xlsx/xls/ods, pptx."
146        ),
147    }
148}
149
150/// Parse a 1-indexed page-range spec like "1-5,8,10-12" into a set predicate.
151fn page_selected(spec: &Option<&str>, page_1indexed: usize) -> bool {
152    let Some(spec) = spec else {
153        return true;
154    };
155    for part in spec.split(',') {
156        let part = part.trim();
157        if part.is_empty() {
158            continue;
159        }
160        if let Some((a, b)) = part.split_once('-') {
161            if let (Ok(a), Ok(b)) = (a.trim().parse::<usize>(), b.trim().parse::<usize>()) {
162                if page_1indexed >= a && page_1indexed <= b {
163                    return true;
164                }
165            }
166        } else if let Ok(n) = part.parse::<usize>() {
167            if page_1indexed == n {
168                return true;
169            }
170        }
171    }
172    false
173}
174
175fn extract_pdf(path: &Path, markdown: bool, pages: Option<&str>) -> Result<String> {
176    let by_pages = pdf_extract::extract_text_by_pages(path)
177        .map_err(|e| anyhow::anyhow!("PDF text extraction failed: {e}"))?;
178    let mut out = String::new();
179    for (i, page_text) in by_pages.iter().enumerate() {
180        let page_no = i + 1;
181        if !page_selected(&pages, page_no) {
182            continue;
183        }
184        let trimmed = page_text.trim();
185        if trimmed.is_empty() {
186            continue;
187        }
188        if markdown {
189            out.push_str(&format!("## Page {page_no}\n\n{trimmed}\n\n"));
190        } else {
191            out.push_str(trimmed);
192            out.push_str("\n\n");
193        }
194    }
195    Ok(out.trim_end().to_string())
196}
197
198fn extract_spreadsheet(path: &Path, sheet: Option<&str>) -> Result<String> {
199    use calamine::{open_workbook_auto, Data, Reader};
200
201    let mut workbook =
202        open_workbook_auto(path).map_err(|e| anyhow::anyhow!("cannot open spreadsheet: {e}"))?;
203    let names: Vec<String> = workbook.sheet_names().to_vec();
204    let mut out = String::new();
205
206    for name in names {
207        if let Some(want) = sheet {
208            if !name.eq_ignore_ascii_case(want) {
209                continue;
210            }
211        }
212        let range = match workbook.worksheet_range(&name) {
213            Ok(r) => r,
214            Err(e) => {
215                out.push_str(&format!("# Sheet: {name}\n\n(error reading sheet: {e})\n\n"));
216                continue;
217            }
218        };
219        out.push_str(&format!("# Sheet: {name}\n\n"));
220
221        let mut rows = range.rows().peekable();
222        // Header row (first row) → Markdown table header + separator.
223        if let Some(first) = rows.next() {
224            let header = row_to_cells(first);
225            out.push_str(&format!("| {} |\n", header.join(" | ")));
226            out.push_str(&format!(
227                "| {} |\n",
228                header.iter().map(|_| "---").collect::<Vec<_>>().join(" | ")
229            ));
230            for row in rows {
231                let cells = row_to_cells(row);
232                out.push_str(&format!("| {} |\n", cells.join(" | ")));
233            }
234        }
235        out.push('\n');
236    }
237
238    if sheet.is_some() && out.is_empty() {
239        anyhow::bail!("sheet '{}' not found", sheet.unwrap());
240    }
241    Ok(out.trim_end().to_string())
242}
243
244/// Render a calamine row of `Data` into Markdown-safe cell strings.
245fn row_to_cells(row: &[calamine::Data]) -> Vec<String> {
246    row.iter()
247        .map(|c| {
248            let s = match c {
249                calamine::Data::Empty => String::new(),
250                other => other.to_string(),
251            };
252            // Escape pipes so cells don't break the Markdown table.
253            s.replace('|', "\\|").replace('\n', " ").trim().to_string()
254        })
255        .collect()
256}
257
258fn extract_docx(path: &Path) -> Result<String> {
259    use dotext::MsDoc;
260    use std::io::Read;
261
262    let mut file =
263        dotext::Docx::open(path).map_err(|e| anyhow::anyhow!("cannot open .docx: {e}"))?;
264    let mut text = String::new();
265    file.read_to_string(&mut text)
266        .map_err(|e| anyhow::anyhow!("cannot read .docx: {e}"))?;
267    Ok(text.trim().to_string())
268}
269
270fn extract_pptx(path: &Path, markdown: bool) -> Result<String> {
271    use dotext::MsDoc;
272    use std::io::Read;
273
274    let mut file =
275        dotext::Pptx::open(path).map_err(|e| anyhow::anyhow!("cannot open .pptx: {e}"))?;
276    let mut text = String::new();
277    file.read_to_string(&mut text)
278        .map_err(|e| anyhow::anyhow!("cannot read .pptx: {e}"))?;
279    let text = text.trim();
280    if markdown {
281        // dotext separates slides with newlines; present as a single section.
282        Ok(format!("## Slides\n\n{text}"))
283    } else {
284        Ok(text.to_string())
285    }
286}
287
288#[cfg(test)]
289mod tests {
290    use super::*;
291
292    #[test]
293    fn test_page_selected_ranges() {
294        assert!(page_selected(&Some("1-5,8"), 3));
295        assert!(page_selected(&Some("1-5,8"), 8));
296        assert!(!page_selected(&Some("1-5,8"), 7));
297        assert!(page_selected(&None, 42)); // no spec = all pages
298        assert!(!page_selected(&Some("2"), 1));
299        assert!(page_selected(&Some("2"), 2));
300    }
301
302    #[test]
303    fn test_row_to_cells_escapes_pipes() {
304        let row = vec![
305            calamine::Data::String("a|b".to_string()),
306            calamine::Data::Empty,
307            calamine::Data::Int(42),
308        ];
309        let cells = row_to_cells(&row);
310        assert_eq!(cells[0], "a\\|b");
311        assert_eq!(cells[1], "");
312        assert_eq!(cells[2], "42");
313    }
314
315    #[test]
316    fn test_unsupported_ext() {
317        let err = extract_document(Path::new("/tmp/x.rtf"), "rtf", true, None, None)
318            .unwrap_err()
319            .to_string();
320        assert!(err.contains("Unsupported"));
321    }
322}
aegis_tools/document.rs

aegis_tools/
document.rs
