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> = 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
29/// Web page scraper loader
30///
31/// Crawls web pages from a URL, extracting the body text. Optionally follows same-domain links recursively.
32pub struct WebScraperLoader {
33    /// The starting URL
34    url: String,
35    /// Maximum recursion depth (0 = only the starting page)
36    max_depth: usize,
37    /// Maximum number of pages to crawl
38    max_pages: usize,
39    /// Whether to return an error when crawling fails (default false, skips failed pages)
40    fail_on_error: bool,
41    /// H8: per-HTTP-request timeout, preventing the crawler from blocking forever on a hung target site
42    timeout: Duration,
43}
44
45impl WebScraperLoader {
46    /// Creates a loader from a URL (crawls only the given page)
47    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    /// Sets the maximum recursion depth
58    pub fn with_max_depth(mut self, depth: usize) -> Self {
59        self.max_depth = depth;
60        self
61    }
62
63    /// Sets the maximum number of pages to crawl
64    pub fn with_max_pages(mut self, pages: usize) -> Self {
65        self.max_pages = pages;
66        self
67    }
68
69    /// Sets whether to return an error on crawl failure (default: skip failed pages)
70    pub fn with_fail_on_error(mut self, fail: bool) -> Self {
71        self.fail_on_error = fail;
72        self
73    }
74
75    /// Sets the per-HTTP-request timeout (H8, default 30s)
76    pub fn with_timeout(mut self, timeout: Duration) -> Self {
77        self.timeout = timeout;
78        self
79    }
80
81    /// Extracts plain text from HTML (reuses HTMLLoader's logic)
82    fn extract_text(html: &str) -> String {
83        super::HTMLLoader::extract_text(html)
84    }
85
86    /// Extracts links from HTML
87    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    /// Extracts the domain
99    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    /// Resolves a relative URL to an absolute URL
107    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            // Find the scheme://domain part
112            let domain = DOMAIN_PREFIX_RE.find(base)?.as_str();
113            Some(format!("{}{}", domain, href))
114        } else {
115            // Relative path
116            let base_dir = base.rfind('/').map(|i| &base[..=i]).unwrap_or(base);
117            Some(format!("{}{}", base_dir, href))
118        }
119    }
120
121    /// Crawls a single page through the shared SSRF-hardened helper; returns the
122    /// post-redirect final URL alongside the body so metadata/link resolution use the
123    /// real destination.
124    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                    // Skip failed pages and continue crawling the rest (exposed via the log facade so hosts can capture it)
151                    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            // Recursively follow links
169            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        // # links should be filtered
202        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}