Skip to main content

lc_rag/loaders/
web_scraper.rs

1//! Web page scraper loader
2//!
3//! Crawls web page content from URLs, extracting the body text and supporting recursive link following.
4//! Built on HTMLLoader's text-extraction logic, adding link discovery and bulk crawling.
5
6use 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
15/// H8: default per-HTTP-request timeout — a hung target site will not block the crawler forever.
16const DEFAULT_HTTP_TIMEOUT: Duration = Duration::from_secs(30);
17
18// M9: Pre-compile regexes once instead of on every call.
19static 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
26/// Web page scraper loader
27///
28/// Crawls web pages from a URL, extracting the body text. Optionally follows same-domain links recursively.
29pub struct WebScraperLoader {
30    /// The starting URL
31    url: String,
32    /// Maximum recursion depth (0 = only the starting page)
33    max_depth: usize,
34    /// Maximum number of pages to crawl
35    max_pages: usize,
36    /// Whether to return an error when crawling fails (default false, skips failed pages)
37    fail_on_error: bool,
38    /// H8: per-HTTP-request timeout, preventing the crawler from blocking forever on a hung target site
39    timeout: Duration,
40}
41
42impl WebScraperLoader {
43    /// Creates a loader from a URL (crawls only the given page)
44    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    /// Sets the maximum recursion depth
55    pub fn with_max_depth(mut self, depth: usize) -> Self {
56        self.max_depth = depth;
57        self
58    }
59
60    /// Sets the maximum number of pages to crawl
61    pub fn with_max_pages(mut self, pages: usize) -> Self {
62        self.max_pages = pages;
63        self
64    }
65
66    /// Sets whether to return an error on crawl failure (default: skip failed pages)
67    pub fn with_fail_on_error(mut self, fail: bool) -> Self {
68        self.fail_on_error = fail;
69        self
70    }
71
72    /// Sets the per-HTTP-request timeout (H8, default 30s)
73    pub fn with_timeout(mut self, timeout: Duration) -> Self {
74        self.timeout = timeout;
75        self
76    }
77
78    /// Extracts plain text from HTML (reuses HTMLLoader's logic)
79    fn extract_text(html: &str) -> String {
80        super::HTMLLoader::extract_text(html)
81    }
82
83    /// Extracts links from HTML
84    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    /// Extracts the domain
96    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    /// Resolves a relative URL to an absolute URL
104    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            // Find the scheme://domain part
109            let domain = DOMAIN_PREFIX_RE.find(base)?.as_str();
110            Some(format!("{}{}", domain, href))
111        } else {
112            // Relative path
113            let base_dir = base.rfind('/').map(|i| &base[..=i]).unwrap_or(base);
114            Some(format!("{}{}", base_dir, href))
115        }
116    }
117
118    /// Crawls a single page
119    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                    // Skip failed pages and continue crawling the rest (exposed via the log facade so hosts can capture it)
166                    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            // Recursively follow links
184            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        // # links should be filtered
217        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}