diff options
| -rw-r--r-- | Cargo.lock | 176 | ||||
| -rw-r--r-- | Cargo.toml | 2 | ||||
| -rw-r--r-- | src/rag/loader.rs | 70 | ||||
| -rw-r--r-- | src/rag/mod.rs | 3 |
4 files changed, 57 insertions, 194 deletions
@@ -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" @@ -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 }); |