diff options
| -rw-r--r-- | Cargo.lock | 144 | ||||
| -rw-r--r-- | Cargo.toml | 1 | ||||
| -rw-r--r-- | config.example.yaml | 1 | ||||
| -rw-r--r-- | src/config/mod.rs | 16 | ||||
| -rw-r--r-- | src/rag/loader.rs | 23 | ||||
| -rw-r--r-- | src/rag/mod.rs | 4 | ||||
| -rw-r--r-- | src/utils/mod.rs | 11 | ||||
| -rw-r--r-- | src/utils/request.rs | 309 |
8 files changed, 483 insertions, 26 deletions
@@ -24,6 +24,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" dependencies = [ "cfg-if", + "getrandom", "once_cell", "version_check", "zerocopy 0.7.35", @@ -80,6 +81,7 @@ dependencies = [ "reedline", "reqwest", "reqwest-eventsource", + "scraper", "serde", "serde_json", "serde_yaml", @@ -624,6 +626,29 @@ dependencies = [ ] [[package]] +name = "cssparser" +version = "0.31.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3df4f93e5fbbe73ec01ec8d3f68bba73107993a5b1e7519273c32db9b0d5be" +dependencies = [ + "cssparser-macros", + "dtoa-short", + "itoa", + "phf 0.11.2", + "smallvec", +] + +[[package]] +name = "cssparser-macros" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13b588ba4ac1a99f7f2964d24b3d896ddc6bf847ee3855dbd4366f058cfcd331" +dependencies = [ + "quote", + "syn", +] + +[[package]] name = "deranged" version = "0.3.11" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -644,6 +669,17 @@ dependencies = [ ] [[package]] +name = "derive_more" +version = "0.99.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f33878137e4dafd7fa914ad4e259e18a4e8e532b9617a2d0150262bf53abfce" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] name = "diff" version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -697,12 +733,33 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "75b325c5dbd37f80359721ad39aca5a29fb04c89279657cffdda8736d0c0b9d2" [[package]] +name = "dtoa" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcbb2bf8e87535c23f7a8a321e364ce21462d0ff10cb6407820e8e96dfff6653" + +[[package]] +name = "dtoa-short" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd1511a7b6a56299bd043a9c167a6d2bfb37bf84a6dfceaba651168adfb43c87" +dependencies = [ + "dtoa", +] + +[[package]] name = "dyn-clone" version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d6ef0072f8a535281e4876be788938b528e9a1d43900b82c2569af7da799125" [[package]] +name = "ego-tree" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a68a4904193147e0a8dec3314640e6db742afd5f6e634f428a6af230d9b3591" + +[[package]] name = "either" version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1442,8 +1499,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "16ce3abbeba692c8b8441d036ef91aea6df8da2c6b6e21c7e14d3c18e526be45" dependencies = [ "log", - "phf", - "phf_codegen", + "phf 0.11.2", + "phf_codegen 0.11.2", "string_cache", "string_cache_codegen", "tendril", @@ -1886,15 +1943,35 @@ dependencies = [ [[package]] name = "phf" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fabbf1ead8a5bcbc20f5f8b939ee3f5b0f6f281b6ad3468b84656b658b455259" +dependencies = [ + "phf_shared 0.10.0", +] + +[[package]] +name = "phf" version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ade2d8b8f33c7333b51bcf0428d37e217e9f32192ae4772156f65063b8ce03dc" dependencies = [ + "phf_macros", "phf_shared 0.11.2", ] [[package]] name = "phf_codegen" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fb1c3a8bc4dd4e5cfce29b44ffc14bedd2ee294559a294e2a4d4c9e9a6a13cd" +dependencies = [ + "phf_generator 0.10.0", + "phf_shared 0.10.0", +] + +[[package]] +name = "phf_codegen" version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8d39688d359e6b34654d328e262234662d16cc0f60ec8dcbe5e718709342a5a" @@ -1924,6 +2001,19 @@ dependencies = [ ] [[package]] +name = "phf_macros" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3444646e286606587e49f3bcf1679b8cef1dc2c5ecc29ddacaffc305180d464b" +dependencies = [ + "phf_generator 0.11.2", + "phf_shared 0.11.2", + "proc-macro2", + "quote", + "syn", +] + +[[package]] name = "phf_shared" version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -2448,6 +2538,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] +name = "scraper" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b90460b31bfe1fc07be8262e42c665ad97118d4585869de9345a84d501a9eaf0" +dependencies = [ + "ahash", + "cssparser", + "ego-tree", + "html5ever", + "indexmap", + "once_cell", + "selectors", + "tendril", +] + +[[package]] name = "security-framework" version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -2471,6 +2577,25 @@ dependencies = [ ] [[package]] +name = "selectors" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4eb30575f3638fc8f6815f448d50cb1a2e255b0897985c8c59f4d37b72a07b06" +dependencies = [ + "bitflags 2.6.0", + "cssparser", + "derive_more", + "fxhash", + "log", + "new_debug_unreachable", + "phf 0.10.1", + "phf_codegen 0.10.0", + "precomputed-hash", + "servo_arc", + "smallvec", +] + +[[package]] name = "serde" version = "1.0.204" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -2529,6 +2654,15 @@ dependencies = [ ] [[package]] +name = "servo_arc" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d036d71a959e00c77a63538b90a6c2390969f9772b096ea837205c6bd0491a44" +dependencies = [ + "stable_deref_trait", +] + +[[package]] name = "sha2" version = "0.10.8" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -2639,6 +2773,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" [[package]] +name = "stable_deref_trait" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" + +[[package]] name = "string_cache" version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -60,6 +60,7 @@ hnsw_rs = "0.3.0" rayon = "1.10.0" uuid = { version = "1.9.1", features = ["v4"] } html2text = "0.12.5" +scraper = { version = "0.20.0", default-features = false, features = ["deterministic"] } sys-locale = "0.3.1" [dependencies.reqwest] diff --git a/config.example.yaml b/config.example.yaml index 5e82f4e..32e7f83 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -65,7 +65,6 @@ document_loaders: # Note: Use `$1` for input file and `$2` for output file. If `$2` is omitted, use stdout as output. pdf: 'pdftotext $1 -' # Load .pdf file, see https://poppler.freedesktop.org to set up pdftotext docx: 'pandoc --to plain $1' # Load .docx file, see https://pandoc.org to set up pandoc - recursive_url: 'rag-crawler $1 $2' # Load websites, see https://github.com/sigoden/rag-crawler to set up rag-crawler # ---- apperence ---- highlight: true # Controls syntax highlighting diff --git a/src/config/mod.rs b/src/config/mod.rs index 1c9ae9d..8db7085 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1740,16 +1740,12 @@ impl Config { } fn setup_document_loaders(&mut self) { - [ - ("pdf", "pdftotext $1 -"), - ("docx", "pandoc --to plain $1"), - (RECURSIVE_URL_LOADER, "rag-crawler $1 $2"), - ] - .into_iter() - .for_each(|(k, v)| { - let (k, v) = (k.to_string(), v.to_string()); - self.document_loaders.entry(k).or_insert(v); - }); + [("pdf", "pdftotext $1 -"), ("docx", "pandoc --to plain $1")] + .into_iter() + .for_each(|(k, v)| { + let (k, v) = (k.to_string(), v.to_string()); + self.document_loaders.entry(k).or_insert(v); + }); } } diff --git a/src/rag/loader.rs b/src/rag/loader.rs index c764d78..f6c3b8e 100644 --- a/src/rag/loader.rs +++ b/src/rag/loader.rs @@ -11,15 +11,20 @@ pub async fn load_recursive_url( path: &str, ) -> Result<Vec<(String, RagMetadata)>> { let extension = RECURSIVE_URL_LOADER; - let loader_command = loaders - .get(extension) - .with_context(|| format!("Document loader '{extension}' not configured"))?; - let contents = run_loader_command(path, extension, loader_command)?; - let pages: Vec<WebPage> = serde_json::from_str(&contents).context(r#"The crawler response is invalid. It should follow the JSON format: `[{"path":"...", "text":"..."}]`."#)?; + let pages: Vec<Page> = match loaders.get(extension) { + Some(loader_command) => { + let contents = run_loader_command(path, extension, loader_command)?; + serde_json::from_str(&contents).context(r#"The crawler response is invalid. It should follow the JSON format: `[{"path":"...", "text":"..."}]`."#)? + } + None => { + let options = CrawlOptions::preset(path); + crawl_website(path, options).await? + } + }; let output = pages .into_iter() .map(|v| { - let WebPage { path, text } = v; + let Page { path, text } = v; let mut metadata: RagMetadata = Default::default(); metadata.insert(PATH_METADATA.into(), path); metadata.insert(EXTENSION_METADATA.into(), "md".into()); @@ -29,12 +34,6 @@ pub async fn load_recursive_url( Ok(output) } -#[derive(Debug, Deserialize)] -struct WebPage { - path: String, - text: String, -} - pub async fn load_path( loaders: &HashMap<String, String>, path: &str, diff --git a/src/rag/mod.rs b/src/rag/mod.rs index 266a859..6dbb6be 100644 --- a/src/rag/mod.rs +++ b/src/rag/mod.rs @@ -371,6 +371,10 @@ impl Rag { self.data.add(next_file_id, files, document_ids, embeddings); self.data.document_paths = document_paths; + if self.data.files.is_empty() { + bail!("No RAG files"); + } + progress(&spinner, "Building store".into()); self.hnsw = self.data.build_hnsw(); self.bm25 = self.data.build_bm25(); diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 6f89fbe..d9d6617 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -121,6 +121,17 @@ pub fn fuzzy_match(text: &str, pattern: &str) -> bool { pattern_index == pattern_chars.len() } +pub fn pretty_error(err: &anyhow::Error) -> String { + let mut output = vec![]; + output.push(err.to_string()); + output.push("Caused by:".to_string()); + for (i, cause) in err.chain().skip(1).enumerate() { + output.push(format!(" {i}: {cause}")); + } + output.push(String::new()); + output.join("\n") +} + pub fn error_text(input: &str) -> String { nu_ansi_term::Style::new() .fg(nu_ansi_term::Color::Red) diff --git a/src/utils/request.rs b/src/utils/request.rs index 4a4e7c9..873cbcc 100644 --- a/src/utils/request.rs +++ b/src/utils/request.rs @@ -1,15 +1,28 @@ use super::*; -use anyhow::{bail, Result}; +use anyhow::{anyhow, bail, Context, Result}; +use fancy_regex::Regex; +use futures_util::{stream, StreamExt}; use http::header::CONTENT_TYPE; +use reqwest::Url; +use scraper::{Html, Selector}; +use serde::Deserialize; +use serde_json::Value; use std::{collections::HashMap, time::Duration}; +use std::{collections::HashSet, sync::Arc}; use tokio::io::AsyncWriteExt; +use tokio::sync::Semaphore; pub const URL_LOADER: &str = "url"; pub const RECURSIVE_URL_LOADER: &str = "recursive_url"; + pub const MEDIA_URL_EXTENSION: &str = "media_url"; pub const DEFAULT_EXTENSION: &str = "txt"; +const MAX_CRAWLS: usize = 5; +const BREAK_ON_ERROR: bool = false; +const USER_AGENT: &str = "curl/8.6.0"; + lazy_static::lazy_static! { static ref CLIENT: Result<reqwest::Client> = { let builder = reqwest::ClientBuilder::new().timeout(Duration::from_secs(30)); @@ -17,6 +30,21 @@ lazy_static::lazy_static! { let client = builder.build()?; Ok(client) }; + + static ref PRESET: Vec<(Regex, CrawlOptions)> = vec![ + ( + Regex::new(r"github.com/([^/]+)/([^/]+)/tree/([^/]+)").unwrap(), + CrawlOptions { exclude: vec!["changelog".into(), "changes".into(), "license".into()], ..Default::default() } + ), + ( + Regex::new(r"github.com/([^/]+)/([^/]+)/wiki").unwrap(), + CrawlOptions { exclude: vec!["_history".into()], extract: Some("#wiki-body".into()), ..Default::default() } + + ) + ]; + + static ref EXTENSION_RE: Regex = Regex::new(r"\.[^.]+$").unwrap(); + static ref GITHUB_REPO_RE: Regex = Regex::new(r"^https://github\.com/([^/]+)/([^/]+)/tree/([^/]+)").unwrap(); } pub async fn fetch( @@ -120,3 +148,282 @@ pub async fn fetch( }; Ok(result) } + +#[derive(Debug, Clone, Default)] +pub struct CrawlOptions { + extract: Option<String>, + exclude: Vec<String>, + no_log: bool, +} + +impl CrawlOptions { + pub fn preset(start_url: &str) -> CrawlOptions { + for (re, options) in PRESET.iter() { + if let Ok(true) = re.is_match(start_url) { + return options.clone(); + } + } + CrawlOptions::default() + } +} + +pub async fn crawl_website(start_url: &str, options: CrawlOptions) -> Result<Vec<Page>> { + let start_url = Url::parse(start_url)?; + let mut paths = vec![start_url.path().to_string()]; + let normalized_start_url = normalize_start_url(&start_url); + if !options.no_log { + println!( + "Start crawling url={start_url} exclude={} extract={}", + options.exclude.join(","), + options.extract.as_deref().unwrap_or_default() + ); + } + + if let Ok(true) = GITHUB_REPO_RE.is_match(start_url.as_str()) { + paths = crawl_gh_tree(&start_url, &options.exclude) + .await + .with_context(|| "Failed to craw github repo".to_string())?; + } + + let semaphore = Arc::new(Semaphore::new(MAX_CRAWLS)); + let mut result_pages = Vec::new(); + + let mut index = 0; + while index < paths.len() { + let batch = paths[index..std::cmp::min(index + MAX_CRAWLS, paths.len())].to_vec(); + + let tasks: Vec<_> = batch + .iter() + .map(|path| { + let options = options.clone(); + let permit = semaphore.clone().acquire_owned(); // acquire a permit for concurrency control + let normalized_start_url = normalized_start_url.clone(); + let path = path.clone(); + + async move { + let _permit = permit.await?; + let url = normalized_start_url + .join(&path) + .map_err(|_| anyhow!("Invalid crawl page at {}", path))?; + let mut page = crawl_page(&normalized_start_url, &path, options) + .await + .with_context(|| format!("Failed to crawl page {}", url.as_str()))?; + page.0 = url.as_str().to_string(); + Ok(page) + } + }) + .collect(); + + let results = stream::iter(tasks) + .buffer_unordered(MAX_CRAWLS) + .collect::<Vec<_>>() + .await; + + let mut new_paths = Vec::new(); + + for res in results { + match res { + Ok((path, text, links)) => { + if !options.no_log { + println!("Crawled {path}"); + } + if !text.is_empty() { + result_pages.push(Page { path, text }); + } + for link in links { + if !paths.iter().any(|p| match_link(p, &link)) { + new_paths.push(link); + } + } + } + Err(err) => { + if BREAK_ON_ERROR { + return Err(err); + } else if !options.no_log { + println!("{}", error_text(&pretty_error(&err))); + } + } + } + } + paths.extend(new_paths); + + index += batch.len(); + } + + Ok(result_pages) +} + +#[derive(Debug, Deserialize)] +pub struct Page { + pub path: String, + pub text: String, +} + +async fn crawl_gh_tree(start_url: &Url, exclude: &[String]) -> Result<Vec<String>> { + let path_segs: Vec<&str> = start_url.path().split('/').collect(); + if path_segs.len() < 4 { + bail!("Invalid gh tree {}", start_url.as_str()); + } + let client = match *CLIENT { + Ok(ref client) => client, + Err(ref err) => bail!("{err}"), + }; + let owner = path_segs[1]; + let repo = path_segs[2]; + let branch = path_segs[4]; + let root_path = path_segs[5..].join("/"); + + let url = format!( + "https://api.github.com/repos/{}/{}/git/ref/heads/{}", + owner, repo, branch + ); + + let res_body: Value = client + .get(&url) + .header("User-Agent", USER_AGENT) + .header("Accept", "application/vnd.github+json") + .header("X-GitHub-Api-Version", "2022-11-28") + .send() + .await? + .json() + .await?; + + let sha = res_body["object"]["sha"] + .as_str() + .ok_or_else(|| anyhow!("Not found branch or tag"))?; + + let url = format!( + "https://api.github.com/repos/{}/{}/git/trees/{}?recursive=true", + owner, repo, sha + ); + + let res_body: Value = client + .get(&url) + .header("User-Agent", USER_AGENT) + .header("Accept", "application/vnd.github+json") + .header("X-GitHub-Api-Version", "2022-11-28") + .send() + .await? + .json() + .await?; + let tree = res_body["tree"] + .as_array() + .ok_or_else(|| anyhow!("Invalid github repo tree"))?; + let paths = tree + .iter() + .flat_map(|v| { + let path = v["path"].as_str()?; + if (path.ends_with(".md") || path.ends_with(".MD")) + && path.starts_with(&root_path) + && !should_exclude_link(path, exclude) + { + Some(format!( + "https://raw.githubusercontent.com/{owner}/{repo}/{branch}/{path}" + )) + } else { + None + } + }) + .collect(); + + Ok(paths) +} + +async fn crawl_page( + start_url: &Url, + path: &str, + options: CrawlOptions, +) -> Result<(String, String, Vec<String>)> { + let client = match *CLIENT { + Ok(ref client) => client, + Err(ref err) => bail!("{err}"), + }; + let location = start_url.join(path)?; + let response = client + .get(location.as_str()) + .header("User-Agent", USER_AGENT) + .send() + .await?; + let body = response.text().await?; + + if let Ok(true) = GITHUB_REPO_RE.is_match(start_url.as_str()) { + return Ok((path.to_string(), body, vec![])); + } + + let mut links = HashSet::new(); + let document = Html::parse_document(&body); + let selector = Selector::parse("a").map_err(|err| anyhow!("Invalid link selector, {}", err))?; + + for element in document.select(&selector) { + if let Some(href) = element.value().attr("href") { + let href = Url::parse(href).ok().or_else(|| location.join(href).ok()); + match href { + None => continue, + Some(href) => { + if href.as_str().starts_with(location.as_str()) + && !should_exclude_link(href.path(), &options.exclude) + { + links.insert(href.path().to_string()); + } + } + } + } + } + + let text = if let Some(selector) = &options.extract { + let selector = Selector::parse(selector) + .map_err(|err| anyhow!("Invalid extract selector, {}", err))?; + document + .select(&selector) + .map(|v| html_to_md(&v.html())) + .collect::<Vec<String>>() + .join("\n\n") + } else { + html_to_md(&body) + }; + + Ok((path.to_string(), text, links.into_iter().collect())) +} + +fn html_to_md(html: &str) -> String { + html2text::from_read(html.as_bytes(), usize::MAX) +} + +fn should_exclude_link(link: &str, exclude: &[String]) -> bool { + if link.contains("#") { + return true; + } + let parts: Vec<&str> = link.trim_end_matches('/').split('/').collect(); + let name = parts.last().unwrap_or(&"").to_lowercase(); + + for exclude_name in exclude { + let yes = match EXTENSION_RE.is_match(exclude_name) { + Ok(true) => exclude_name.to_lowercase() == name.to_lowercase(), + _ => exclude_name.to_lowercase() == EXTENSION_RE.replace(&name, "").to_lowercase(), + }; + if yes { + return true; + } + } + false +} + +fn normalize_start_url(start_url: &Url) -> Url { + let mut start_url = start_url.clone(); + start_url.set_query(None); + start_url.set_fragment(None); + let new_path = match start_url.path().rfind('/') { + Some(last_slash_index) => start_url.path()[..last_slash_index + 1].to_string(), + None => start_url.path().to_string(), + }; + start_url.set_path(&new_path); + start_url +} + +fn match_link(path: &str, link: &str) -> bool { + path == link + || path + == link + .trim_end_matches("/index.html") + .trim_end_matches("/index.htm") +} |