lc_rag/loaders/
web_scraper.rs1use std::collections::{HashMap, HashSet};
7use std::sync::LazyLock;
8use std::time::Duration;
9
10use async_trait::async_trait;
11
12use super::{DocumentLoader, LoaderError};
13use lc_vector_stores::Document;
14
15const DEFAULT_HTTP_TIMEOUT: Duration = Duration::from_secs(30);
17
18static HREF_RE: LazyLock<regex::Regex> = LazyLock::new(|| {
20 regex::Regex::new(r#"href\s*=\s*["']([^"']+)["']"#).expect("static regex literal must compile")
21});
22static DOMAIN_RE: LazyLock<regex::Regex> = LazyLock::new(|| {
23 regex::Regex::new(r"https?://([^/]+)").expect("static regex literal must compile")
24});
25static DOMAIN_PREFIX_RE: LazyLock<regex::Regex> = LazyLock::new(|| {
26 regex::Regex::new(r"https?://[^/]+").expect("static regex literal must compile")
27});
28
29pub struct WebScraperLoader {
33 url: String,
35 max_depth: usize,
37 max_pages: usize,
39 fail_on_error: bool,
41 timeout: Duration,
43}
44
45impl WebScraperLoader {
46 pub fn new(url: impl Into<String>) -> Self {
48 Self {
49 url: url.into(),
50 max_depth: 0,
51 max_pages: 1,
52 fail_on_error: false,
53 timeout: DEFAULT_HTTP_TIMEOUT,
54 }
55 }
56
57 pub fn with_max_depth(mut self, depth: usize) -> Self {
59 self.max_depth = depth;
60 self
61 }
62
63 pub fn with_max_pages(mut self, pages: usize) -> Self {
65 self.max_pages = pages;
66 self
67 }
68
69 pub fn with_fail_on_error(mut self, fail: bool) -> Self {
71 self.fail_on_error = fail;
72 self
73 }
74
75 pub fn with_timeout(mut self, timeout: Duration) -> Self {
77 self.timeout = timeout;
78 self
79 }
80
81 fn extract_text(html: &str) -> String {
83 super::HTMLLoader::extract_text(html)
84 }
85
86 fn extract_links(html: &str, base_url: &str) -> Vec<String> {
88 let base_domain = Self::extract_domain(base_url);
89 HREF_RE
90 .captures_iter(html)
91 .filter_map(|cap| cap.get(1).map(|m| m.as_str().to_string()))
92 .filter(|link| !link.starts_with('#') && !link.starts_with("javascript:"))
93 .filter_map(|link| Self::resolve_url(base_url, &link))
94 .filter(|url| Self::extract_domain(url) == base_domain)
95 .collect()
96 }
97
98 fn extract_domain(url: &str) -> String {
100 DOMAIN_RE
101 .captures(url)
102 .and_then(|c| c.get(1).map(|m| m.as_str().to_string()))
103 .unwrap_or_default()
104 }
105
106 fn resolve_url(base: &str, href: &str) -> Option<String> {
108 if href.starts_with("http://") || href.starts_with("https://") {
109 Some(href.to_string())
110 } else if href.starts_with('/') {
111 let domain = DOMAIN_PREFIX_RE.find(base)?.as_str();
113 Some(format!("{}{}", domain, href))
114 } else {
115 let base_dir = base.rfind('/').map(|i| &base[..=i]).unwrap_or(base);
117 Some(format!("{}{}", base_dir, href))
118 }
119 }
120
121 async fn fetch_page(url: &str, timeout: Duration) -> Result<(String, String), LoaderError> {
125 super::guarded_fetch(url, timeout).await
126 }
127}
128
129#[async_trait]
130impl DocumentLoader for WebScraperLoader {
131 async fn load(&self) -> Result<Vec<Document>, LoaderError> {
132 let mut documents = Vec::new();
133 let mut visited = HashSet::new();
134 let mut queue = vec![(self.url.clone(), 0usize)];
135 let mut failed_count: usize = 0;
136
137 while let Some((url, depth)) = queue.pop() {
138 if visited.contains(&url) || documents.len() >= self.max_pages {
139 continue;
140 }
141 visited.insert(url.clone());
142
143 let (fetched_url, html) = match Self::fetch_page(&url, self.timeout).await {
144 Ok(r) => r,
145 Err(e) => {
146 failed_count += 1;
147 if self.fail_on_error {
148 return Err(e);
149 }
150 log::warn!("Failed to crawl {} (failure #{}): {}", url, failed_count, e);
152 continue;
153 }
154 };
155
156 let text = Self::extract_text(&html);
157
158 let mut metadata = HashMap::new();
159 metadata.insert("format".to_string(), "html".to_string().into());
160 metadata.insert("source".to_string(), fetched_url.clone().into());
161
162 documents.push(Document {
163 content: text,
164 metadata,
165 id: None,
166 });
167
168 if depth < self.max_depth {
170 let links = Self::extract_links(&html, &fetched_url);
171 for link in links {
172 if !visited.contains(&link) {
173 queue.push((link, depth + 1));
174 }
175 }
176 }
177 }
178
179 if failed_count > 0 {
180 log::warn!(
181 "Crawling finished: {} pages failed, {} pages succeeded",
182 failed_count,
183 documents.len()
184 );
185 }
186
187 Ok(documents)
188 }
189}
190
191#[cfg(test)]
192mod tests {
193 use super::*;
194
195 #[test]
196 fn test_extract_links() {
197 let html = "<html><body><a href=\"/about\">About</a><a href=\"https://example.com/contact\">Contact</a><a href=\"#top\">Top</a></body></html>";
198 let links = WebScraperLoader::extract_links(html, "https://example.com/");
199 assert!(links.contains(&"https://example.com/about".to_string()));
200 assert!(links.contains(&"https://example.com/contact".to_string()));
201 assert!(!links.iter().any(|l| l.contains('#')));
203 }
204
205 #[test]
206 fn test_extract_domain() {
207 assert_eq!(
208 WebScraperLoader::extract_domain("https://example.com/path"),
209 "example.com"
210 );
211 assert_eq!(
212 WebScraperLoader::extract_domain("http://sub.example.com:8080/path"),
213 "sub.example.com:8080"
214 );
215 }
216
217 #[test]
218 fn test_resolve_url_absolute() {
219 let result =
220 WebScraperLoader::resolve_url("https://example.com/", "https://other.com/page");
221 assert_eq!(result, Some("https://other.com/page".to_string()));
222 }
223
224 #[test]
225 fn test_resolve_url_relative() {
226 let result = WebScraperLoader::resolve_url("https://example.com/dir/page", "other");
227 assert_eq!(result, Some("https://example.com/dir/other".to_string()));
228 }
229
230 #[test]
231 fn test_resolve_url_root_relative() {
232 let result = WebScraperLoader::resolve_url("https://example.com/dir/page", "/root");
233 assert_eq!(result, Some("https://example.com/root".to_string()));
234 }
235
236 #[test]
237 fn test_extract_text() {
238 let html = "<html><body><p>Hello World</p></body></html>";
239 let text = WebScraperLoader::extract_text(html);
240 assert!(text.contains("Hello World"));
241 }
242
243 #[test]
244 fn test_new_creates_single_page_scraper() {
245 let loader = WebScraperLoader::new("https://example.com");
246 assert_eq!(loader.max_depth, 0);
247 assert_eq!(loader.max_pages, 1);
248 }
249
250 #[test]
251 fn test_with_options() {
252 let loader = WebScraperLoader::new("https://example.com")
253 .with_max_depth(2)
254 .with_max_pages(10);
255 assert_eq!(loader.max_depth, 2);
256 assert_eq!(loader.max_pages, 10);
257 }
258}