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> =
20 LazyLock::new(|| regex::Regex::new(r#"href\s*=\s*["']([^"']+)["']"#).unwrap());
21static DOMAIN_RE: LazyLock<regex::Regex> =
22 LazyLock::new(|| regex::Regex::new(r"https?://([^/]+)").unwrap());
23static DOMAIN_PREFIX_RE: LazyLock<regex::Regex> =
24 LazyLock::new(|| regex::Regex::new(r"https?://[^/]+").unwrap());
25
26pub struct WebScraperLoader {
30 url: String,
32 max_depth: usize,
34 max_pages: usize,
36 fail_on_error: bool,
38 timeout: Duration,
40}
41
42impl WebScraperLoader {
43 pub fn new(url: impl Into<String>) -> Self {
45 Self {
46 url: url.into(),
47 max_depth: 0,
48 max_pages: 1,
49 fail_on_error: false,
50 timeout: DEFAULT_HTTP_TIMEOUT,
51 }
52 }
53
54 pub fn with_max_depth(mut self, depth: usize) -> Self {
56 self.max_depth = depth;
57 self
58 }
59
60 pub fn with_max_pages(mut self, pages: usize) -> Self {
62 self.max_pages = pages;
63 self
64 }
65
66 pub fn with_fail_on_error(mut self, fail: bool) -> Self {
68 self.fail_on_error = fail;
69 self
70 }
71
72 pub fn with_timeout(mut self, timeout: Duration) -> Self {
74 self.timeout = timeout;
75 self
76 }
77
78 fn extract_text(html: &str) -> String {
80 super::HTMLLoader::extract_text(html)
81 }
82
83 fn extract_links(html: &str, base_url: &str) -> Vec<String> {
85 let base_domain = Self::extract_domain(base_url);
86 HREF_RE
87 .captures_iter(html)
88 .filter_map(|cap| cap.get(1).map(|m| m.as_str().to_string()))
89 .filter(|link| !link.starts_with('#') && !link.starts_with("javascript:"))
90 .filter_map(|link| Self::resolve_url(base_url, &link))
91 .filter(|url| Self::extract_domain(url) == base_domain)
92 .collect()
93 }
94
95 fn extract_domain(url: &str) -> String {
97 DOMAIN_RE
98 .captures(url)
99 .and_then(|c| c.get(1).map(|m| m.as_str().to_string()))
100 .unwrap_or_default()
101 }
102
103 fn resolve_url(base: &str, href: &str) -> Option<String> {
105 if href.starts_with("http://") || href.starts_with("https://") {
106 Some(href.to_string())
107 } else if href.starts_with('/') {
108 let domain = DOMAIN_PREFIX_RE.find(base)?.as_str();
110 Some(format!("{}{}", domain, href))
111 } else {
112 let base_dir = base.rfind('/').map(|i| &base[..=i]).unwrap_or(base);
114 Some(format!("{}{}", base_dir, href))
115 }
116 }
117
118 async fn fetch_page(url: &str, timeout: Duration) -> Result<(String, String), LoaderError> {
120 let client = reqwest::Client::builder()
121 .timeout(timeout)
122 .build()
123 .map_err(|e| LoaderError::Other(format!("failed to build HTTP client: {}", e)))?;
124 let response = client
125 .get(url)
126 .send()
127 .await
128 .map_err(|e| LoaderError::Other(format!("HTTP request failed {}: {}", url, e)))?;
129 let status = response.status();
130 if !status.is_success() {
131 return Err(LoaderError::Other(format!(
132 "HTTP error {}: {}",
133 url, status
134 )));
135 }
136 let html = response
137 .text()
138 .await
139 .map_err(|e| LoaderError::Other(format!("failed to read response {}: {}", url, e)))?;
140 Ok((url.to_string(), html))
141 }
142}
143
144#[async_trait]
145impl DocumentLoader for WebScraperLoader {
146 async fn load(&self) -> Result<Vec<Document>, LoaderError> {
147 let mut documents = Vec::new();
148 let mut visited = HashSet::new();
149 let mut queue = vec![(self.url.clone(), 0usize)];
150 let mut failed_count: usize = 0;
151
152 while let Some((url, depth)) = queue.pop() {
153 if visited.contains(&url) || documents.len() >= self.max_pages {
154 continue;
155 }
156 visited.insert(url.clone());
157
158 let (fetched_url, html) = match Self::fetch_page(&url, self.timeout).await {
159 Ok(r) => r,
160 Err(e) => {
161 failed_count += 1;
162 if self.fail_on_error {
163 return Err(e);
164 }
165 log::warn!("Failed to crawl {} (failure #{}): {}", url, failed_count, e);
167 continue;
168 }
169 };
170
171 let text = Self::extract_text(&html);
172
173 let mut metadata = HashMap::new();
174 metadata.insert("format".to_string(), "html".to_string().into());
175 metadata.insert("source".to_string(), fetched_url.clone().into());
176
177 documents.push(Document {
178 content: text,
179 metadata,
180 id: None,
181 });
182
183 if depth < self.max_depth {
185 let links = Self::extract_links(&html, &fetched_url);
186 for link in links {
187 if !visited.contains(&link) {
188 queue.push((link, depth + 1));
189 }
190 }
191 }
192 }
193
194 if failed_count > 0 {
195 log::warn!(
196 "Crawling finished: {} pages failed, {} pages succeeded",
197 failed_count,
198 documents.len()
199 );
200 }
201
202 Ok(documents)
203 }
204}
205
206#[cfg(test)]
207mod tests {
208 use super::*;
209
210 #[test]
211 fn test_extract_links() {
212 let html = "<html><body><a href=\"/about\">About</a><a href=\"https://example.com/contact\">Contact</a><a href=\"#top\">Top</a></body></html>";
213 let links = WebScraperLoader::extract_links(html, "https://example.com/");
214 assert!(links.contains(&"https://example.com/about".to_string()));
215 assert!(links.contains(&"https://example.com/contact".to_string()));
216 assert!(!links.iter().any(|l| l.contains('#')));
218 }
219
220 #[test]
221 fn test_extract_domain() {
222 assert_eq!(
223 WebScraperLoader::extract_domain("https://example.com/path"),
224 "example.com"
225 );
226 assert_eq!(
227 WebScraperLoader::extract_domain("http://sub.example.com:8080/path"),
228 "sub.example.com:8080"
229 );
230 }
231
232 #[test]
233 fn test_resolve_url_absolute() {
234 let result =
235 WebScraperLoader::resolve_url("https://example.com/", "https://other.com/page");
236 assert_eq!(result, Some("https://other.com/page".to_string()));
237 }
238
239 #[test]
240 fn test_resolve_url_relative() {
241 let result = WebScraperLoader::resolve_url("https://example.com/dir/page", "other");
242 assert_eq!(result, Some("https://example.com/dir/other".to_string()));
243 }
244
245 #[test]
246 fn test_resolve_url_root_relative() {
247 let result = WebScraperLoader::resolve_url("https://example.com/dir/page", "/root");
248 assert_eq!(result, Some("https://example.com/root".to_string()));
249 }
250
251 #[test]
252 fn test_extract_text() {
253 let html = "<html><body><p>Hello World</p></body></html>";
254 let text = WebScraperLoader::extract_text(html);
255 assert!(text.contains("Hello World"));
256 }
257
258 #[test]
259 fn test_new_creates_single_page_scraper() {
260 let loader = WebScraperLoader::new("https://example.com");
261 assert_eq!(loader.max_depth, 0);
262 assert_eq!(loader.max_pages, 1);
263 }
264
265 #[test]
266 fn test_with_options() {
267 let loader = WebScraperLoader::new("https://example.com")
268 .with_max_depth(2)
269 .with_max_pages(10);
270 assert_eq!(loader.max_depth, 2);
271 assert_eq!(loader.max_pages, 10);
272 }
273}