Skip to main content

aegis_tools/
doc_extract_pro.rs

1//! # doc_extract_pro (opt-in, external)
2//!
3//! Layer-1 heavy document extraction via the external
4//! [`opendataloader-pdf`](https://github.com/opendataloader-project/opendataloader-pdf)
5//! CLI. Handles what the pure-Rust `read_document` cannot: scanned PDFs (OCR),
6//! complex/borderless tables, formulas and chart descriptions.
7//!
8//! Off by default. Enable in config:
9//! ```toml
10//! [doc_extract]
11//! enabled = true
12//! path    = "opendataloader-pdf"   # binary on PATH or absolute path
13//! mode    = "fast"                 # "fast" | "hybrid"
14//! timeout_secs = 120
15//! ```
16//!
17//! Note: `opendataloader-pdf` requires a JVM (and, for `hybrid`, extra ML
18//! backends), so it is unsuitable for tiny 1c1g hosts — hence opt-in.
19
20use crate::registry::{Tool, ToolContext};
21use aegis_security::check_path;
22use anyhow::Result;
23use async_trait::async_trait;
24use serde_json::{json, Value};
25use std::process::Stdio;
26
27/// Opt-in wrapper around the external `opendataloader-pdf` CLI.
28pub struct DocExtractProTool {
29    /// Path/name of the `opendataloader-pdf` executable.
30    pub binary: String,
31    /// "fast" (default, deterministic) or "hybrid" (AI backend, OCR).
32    pub mode: String,
33    /// Per-invocation timeout in seconds.
34    pub timeout_secs: u64,
35}
36
37#[async_trait]
38impl Tool for DocExtractProTool {
39    fn name(&self) -> &str {
40        "doc_extract_pro"
41    }
42
43    fn description(&self) -> &str {
44        "High-accuracy PDF extraction (scanned PDFs/OCR, complex tables, formulas) via the external opendataloader-pdf tool. Outputs Markdown. Heavier than read_document; use for hard documents."
45    }
46
47    fn parameters(&self) -> Value {
48        json!({
49            "type": "object",
50            "properties": {
51                "path": { "type": "string", "description": "Path to the PDF (within the working directory)" },
52                "output_dir": { "type": "string", "description": "Directory to write markdown/json output (within the working directory). Default: alongside the input." }
53            },
54            "required": ["path"]
55        })
56    }
57
58    async fn execute(&self, args: Value, ctx: &ToolContext<'_>) -> Result<String> {
59        let path_str = args["path"].as_str().unwrap_or("").trim();
60        if path_str.is_empty() {
61            return Ok("Error: path is required".to_string());
62        }
63        let safe_path = check_path(path_str, &ctx.cwd)?;
64        if !safe_path.exists() {
65            anyhow::bail!("File not found: {path_str}");
66        }
67
68        // Output dir defaults to the input's parent, confined to the workspace.
69        let out_dir = match args["output_dir"].as_str() {
70            Some(d) if !d.trim().is_empty() => check_path(d.trim(), &ctx.cwd)?,
71            _ => safe_path
72                .parent()
73                .map(|p| p.to_path_buf())
74                .unwrap_or_else(|| ctx.cwd.clone()),
75        };
76        let _ = std::fs::create_dir_all(&out_dir);
77
78        let mut cmd = tokio::process::Command::new(&self.binary);
79        cmd.arg(safe_path.as_os_str())
80            .arg("--output-dir")
81            .arg(out_dir.as_os_str())
82            .arg("--format")
83            .arg("markdown");
84        if self.mode == "hybrid" {
85            cmd.arg("--hybrid").arg("docling-fast");
86        }
87        cmd.stdout(Stdio::piped()).stderr(Stdio::piped());
88
89        let output = tokio::time::timeout(
90            std::time::Duration::from_secs(self.timeout_secs),
91            cmd.output(),
92        )
93        .await
94        .map_err(|_| {
95            anyhow::anyhow!("doc_extract_pro timed out after {}s", self.timeout_secs)
96        })?
97        .map_err(|e| {
98            anyhow::anyhow!(
99                "failed to run '{}': {e}. Is opendataloader-pdf installed and on PATH?",
100                self.binary
101            )
102        })?;
103
104        let stderr = String::from_utf8_lossy(&output.stderr);
105        if !output.status.success() {
106            return Ok(format!(
107                "[doc_extract_pro error] exit {}\n{stderr}",
108                output.status.code().unwrap_or(-1)
109            ));
110        }
111
112        // opendataloader-pdf writes <name>.md into out_dir. Read it back.
113        let stem = safe_path
114            .file_stem()
115            .and_then(|s| s.to_str())
116            .unwrap_or("output");
117        let md_path = out_dir.join(format!("{stem}.md"));
118        if let Ok(md) = std::fs::read_to_string(&md_path) {
119            if !md.trim().is_empty() {
120                return Ok(md);
121            }
122        }
123
124        // Fallback: some builds print to stdout.
125        let stdout = String::from_utf8_lossy(&output.stdout);
126        if !stdout.trim().is_empty() {
127            Ok(stdout.into_owned())
128        } else {
129            Ok(format!(
130                "doc_extract_pro finished but no markdown was found at {}. stderr: {stderr}",
131                md_path.display()
132            ))
133        }
134    }
135}