aegis_tools/
doc_extract_pro.rs1use 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
27pub struct DocExtractProTool {
29 pub binary: String,
31 pub mode: String,
33 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 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 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 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}