refactor: load pdf/docx files (#592)

4 files changed, 57 insertions, 194 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 44b2734..31fab1b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -18,15 +18,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
 
 [[package]]
-name = "adobe-cmap-parser"
-version = "0.4.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "261a937a307ddc70a1605dec925987d7256adb00232161a1855ef9cc820bd8d5"
-dependencies = [
- "pom",
-]
-
-[[package]]
 name = "ahash"
 version = "0.8.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -84,7 +75,6 @@ dependencies = [
  "nu-ansi-term 0.50.0",
  "parking_lot",
  "path-absolutize",
- "pdf-extract",
  "pretty_assertions",
  "rand",
  "reedline",
@@ -106,6 +96,7 @@ dependencies = [
  "unicode-segmentation",
  "unicode-width",
  "urlencoding",
+ "which",
 ]
 
 [[package]]
@@ -735,79 +726,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3dca9240753cf90908d7e4aac30f630662b02aebaa1b58a3cadabdb23385b58b"
 
 [[package]]
-name = "encoding"
-version = "0.2.33"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6b0d943856b990d12d3b55b359144ff341533e516d94098b1d3fc1ac666d36ec"
-dependencies = [
- "encoding-index-japanese",
- "encoding-index-korean",
- "encoding-index-simpchinese",
- "encoding-index-singlebyte",
- "encoding-index-tradchinese",
-]
-
-[[package]]
-name = "encoding-index-japanese"
-version = "1.20141219.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "04e8b2ff42e9a05335dbf8b5c6f7567e5591d0d916ccef4e0b1710d32a0d0c91"
-dependencies = [
- "encoding_index_tests",
-]
-
-[[package]]
-name = "encoding-index-korean"
-version = "1.20141219.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4dc33fb8e6bcba213fe2f14275f0963fd16f0a02c878e3095ecfdf5bee529d81"
-dependencies = [
- "encoding_index_tests",
-]
-
-[[package]]
-name = "encoding-index-simpchinese"
-version = "1.20141219.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d87a7194909b9118fc707194baa434a4e3b0fb6a5a757c73c3adb07aa25031f7"
-dependencies = [
- "encoding_index_tests",
-]
-
-[[package]]
-name = "encoding-index-singlebyte"
-version = "1.20141219.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3351d5acffb224af9ca265f435b859c7c01537c0849754d3db3fdf2bfe2ae84a"
-dependencies = [
- "encoding_index_tests",
-]
-
-[[package]]
-name = "encoding-index-tradchinese"
-version = "1.20141219.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fd0e20d5688ce3cab59eb3ef3a2083a5c77bf496cb798dc6fcdb75f323890c18"
-dependencies = [
- "encoding_index_tests",
-]
-
-[[package]]
-name = "encoding_index_tests"
-version = "0.1.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a246d82be1c9d791c5dfde9a2bd045fc3cbba3fa2b11ad558f27d01712f00569"
-
-[[package]]
-name = "encoding_rs"
-version = "0.8.34"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b45de904aa0b010bce2ab45264d0631681847fa7b6f2eaa7dab7619943bc4f59"
-dependencies = [
- "cfg-if",
-]
-
-[[package]]
 name = "enum-as-inner"
 version = "0.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -855,15 +773,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a0474425d51df81997e2f90a21591180b38eccf27292d755f3e30750225c175b"
 
 [[package]]
-name = "euclid"
-version = "0.20.14"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2bb7ef65b3777a325d1eeefefab5b6d4959da54747e33bd6258e789640f307ad"
-dependencies = [
- "num-traits",
-]
-
-[[package]]
 name = "eventsource-stream"
 version = "0.2.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1459,12 +1368,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dd1bc4d24ad230d21fb898d1116b1801d7adfc449d42026475862ab48b11e70e"
 
 [[package]]
-name = "linked-hash-map"
-version = "0.5.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0717cef1bc8b636c6e1c1bbdefc09e6322da8a9321966e8928ef80d20f7f770f"
-
-[[package]]
 name = "linux-raw-sys"
 version = "0.4.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1503,23 +1406,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "lopdf"
-version = "0.32.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e775e4ee264e8a87d50a9efef7b67b4aa988cf94e75630859875fc347e6c872b"
-dependencies = [
- "encoding_rs",
- "flate2",
- "itoa",
- "linked-hash-map",
- "log",
- "md5",
- "nom",
- "time",
- "weezl",
-]
-
-[[package]]
 name = "mach2"
 version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1538,12 +1424,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "md5"
-version = "0.7.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "490cc448043f947bae3cbee9c203358d62dbee0db12107a74be5c30ccfd09771"
-
-[[package]]
 name = "memchr"
 version = "2.7.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1936,21 +1816,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "pdf-extract"
-version = "0.7.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3423481005e61b95855d53d7b6c0bcc514b4fbab45d165776d6e42c0a1642b22"
-dependencies = [
- "adobe-cmap-parser",
- "encoding",
- "euclid",
- "lopdf",
- "postscript",
- "type1-encoding-parser",
- "unicode-normalization",
-]
-
-[[package]]
 name = "percent-encoding"
 version = "2.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -2019,18 +1884,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "pom"
-version = "1.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "60f6ce597ecdcc9a098e7fddacb1065093a3d66446fa16c675e7e71d1b5c28e6"
-
-[[package]]
-name = "postscript"
-version = "0.14.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "78451badbdaebaf17f053fd9152b3ffb33b516104eacb45e7864aaa9c712f306"
-
-[[package]]
 name = "powerfmt"
 version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -3005,15 +2858,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
 
 [[package]]
-name = "type1-encoding-parser"
-version = "0.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d3d6cc09e1a99c7e01f2afe4953789311a1c50baebbdac5b477ecf78e2e92a5b"
-dependencies = [
- "pom",
-]
-
-[[package]]
 name = "typenum"
 version = "1.17.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -3337,10 +3181,16 @@ dependencies = [
 ]
 
 [[package]]
-name = "weezl"
-version = "0.1.8"
+name = "which"
+version = "6.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "53a85b86a771b1c87058196170769dd264f66c0782acf1ae6cc51bfd64b39082"
+checksum = "8211e4f58a2b2805adfbefbc07bab82958fc91e3836339b1ab7ae32465dce0d7"
+dependencies = [
+ "either",
+ "home",
+ "rustix",
+ "winsafe",
+]
 
 [[package]]
 name = "widestring"
@@ -3576,6 +3426,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "winsafe"
+version = "0.0.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904"
+
+[[package]]
 name = "wl-clipboard-rs"
 version = "0.8.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index f4d19f2..6b08e02 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -60,7 +60,7 @@ json-patch = { version = "2.0.0", default-features = false }
 bitflags = "2.5.0"
 path-absolutize = "3.1.1"
 hnsw_rs = "0.3.0"
-pdf-extract = "0.7.7"
+which = "6.0.1"
 
 [dependencies.reqwest]
 version = "0.12.0"
diff --git a/src/rag/loader.rs b/src/rag/loader.rs
index 106802a..ba44dac 100644
--- a/src/rag/loader.rs
+++ b/src/rag/loader.rs
@@ -1,49 +1,44 @@
-use super::RagDocument;
+use super::*;
 
-use anyhow::{bail, Context, Result};
+use anyhow::{bail, Result};
 use async_recursion::async_recursion;
-use std::{path::Path, process::Command};
-use tokio::fs;
+use lazy_static::lazy_static;
+use std::{fs::read_to_string, path::Path};
+use which::which;
 
-pub async fn load(path: &str, extension: &str) -> Result<Vec<RagDocument>> {
+lazy_static! {
+    static ref EXIST_PANDOC: bool = which("pandoc").is_ok();
+    static ref EXIST_PDFTOTEXT: bool = which("pdftotext").is_ok();
+}
+
+pub fn load(path: &str, extension: &str) -> Result<Vec<RagDocument>> {
     match extension {
-        "docx" | "epub" | "ipynb" => load_pandoc(path)
-            .await
-            .context("Failed to load with pandoc"),
-        "pdf" => load_pdf(path).await,
-        _ => load_plain(path).await,
+        "docx" | "epub" => load_with_pandoc(path),
+        "pdf" => load_with_pdftotext(path),
+        _ => load_plain(path),
     }
 }
 
-async fn load_plain(path: &str) -> Result<Vec<RagDocument>> {
-    let contents = fs::read_to_string(path).await?;
+fn load_plain(path: &str) -> Result<Vec<RagDocument>> {
+    let contents = read_to_string(path)?;
     let document = RagDocument::new(contents);
     Ok(vec![document])
 }
 
-async fn load_pdf(path: &str) -> Result<Vec<RagDocument>> {
-    let contents = pdf_extract::extract_text(path)?;
+fn load_with_pdftotext(path: &str) -> Result<Vec<RagDocument>> {
+    if !*EXIST_PDFTOTEXT {
+        bail!("Need to install pdftotext (part of the poppler package) to load the file.")
+    }
+    let contents = run_external_tool("pdftotext", &[path, "-"])?;
     let document = RagDocument::new(contents);
     Ok(vec![document])
 }
 
-async fn load_pandoc(path: &str) -> Result<Vec<RagDocument>> {
-    let output = Command::new("pandoc")
-        .arg("--to")
-        .arg("plain")
-        .arg(path)
-        .output()?;
-
-    if !output.status.success() {
-        let stderr = String::from_utf8_lossy(&output.stderr);
-        bail!(
-            "Pandoc conversion failed with exit code {:?}: {}",
-            output.status.code(),
-            stderr
-        );
+fn load_with_pandoc(path: &str) -> Result<Vec<RagDocument>> {
+    if !*EXIST_PANDOC {
+        bail!("Need to install pandoc to load the file.")
     }
-
-    let contents = std::str::from_utf8(&output.stdout)?;
+    let contents = run_external_tool("pandoc", &["--to", "plain", path])?;
     let document = RagDocument::new(contents);
     Ok(vec![document])
 }
@@ -89,7 +84,7 @@ pub async fn list_files(
     if !entry_path.is_dir() {
         bail!("Not a directory: {:?}", entry_path);
     }
-    let mut reader = fs::read_dir(entry_path).await?;
+    let mut reader = tokio::fs::read_dir(entry_path).await?;
     while let Some(entry) = reader.next_entry().await? {
         let path = entry.path();
         if path.is_file() {
@@ -119,6 +114,19 @@ fn is_valid_extension(suffixes: Option<&Vec<String>>, path: &Path) -> bool {
     true
 }
 
+fn run_external_tool(cmd: &str, args: &[&str]) -> Result<String> {
+    let (success, stdout, stderr) = run_command_with_output(cmd, args, None)?;
+    if success {
+        return Ok(stdout);
+    }
+    let err = if !stderr.is_empty() {
+        stderr
+    } else {
+        format!("`{cmd}` exited with non-zero.")
+    };
+    bail!("{err}")
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/rag/mod.rs b/src/rag/mod.rs
index 28ea6fc..011b574 100644
--- a/src/rag/mod.rs
+++ b/src/rag/mod.rs
@@ -189,8 +189,7 @@ impl Rag {
             let separator = autodetect_separator(&extension);
             let splitter = Splitter::new(self.data.chunk_size, self.data.chunk_overlap, separator);
             let documents = load(&path, &extension)
-                .await
-                .with_context(|| format!("Failed to load text at '{path}'"))?;
+                .with_context(|| format!("Failed to load file at '{path}'"))?;
             let documents =
                 splitter.split_documents(&documents, &SplitterChunkHeaderOptions::default());
             rag_files.push(RagFile { path, documents });