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 through the shared SSRF-hardened helper; returns the
119    /// post-redirect final URL alongside the body so metadata/link resolution use the
120    /// real destination.
121    async fn fetch_page(url: &str, timeout: Duration) -> Result<(String, String), LoaderError> {
122        super::guarded_fetch(url, timeout).await
123    }
124}
125
126#[async_trait]
127impl DocumentLoader for WebScraperLoader {
128    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
129        let mut documents = Vec::new();
130        let mut visited = HashSet::new();
131        let mut queue = vec![(self.url.clone(), 0usize)];
132        let mut failed_count: usize = 0;
133
134        while let Some((url, depth)) = queue.pop() {
135            if visited.contains(&url) || documents.len() >= self.max_pages {
136                continue;
137            }
138            visited.insert(url.clone());
139
140            let (fetched_url, html) = match Self::fetch_page(&url, self.timeout).await {
141                Ok(r) => r,
142                Err(e) => {
143                    failed_count += 1;
144                    if self.fail_on_error {
145                        return Err(e);
146                    }
147                    // Skip failed pages and continue crawling the rest (exposed via the log facade so hosts can capture it)
148                    log::warn!("Failed to crawl {} (failure #{}): {}", url, failed_count, e);
149                    continue;
150                }
151            };
152
153            let text = Self::extract_text(&html);
154
155            let mut metadata = HashMap::new();
156            metadata.insert("format".to_string(), "html".to_string().into());
157            metadata.insert("source".to_string(), fetched_url.clone().into());
158
159            documents.push(Document {
160                content: text,
161                metadata,
162                id: None,
163            });
164
165            // Recursively follow links
166            if depth < self.max_depth {
167                let links = Self::extract_links(&html, &fetched_url);
168                for link in links {
169                    if !visited.contains(&link) {
170                        queue.push((link, depth + 1));
171                    }
172                }
173            }
174        }
175
176        if failed_count > 0 {
177            log::warn!(
178                "Crawling finished: {} pages failed, {} pages succeeded",
179                failed_count,
180                documents.len()
181            );
182        }
183
184        Ok(documents)
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    #[test]
193    fn test_extract_links() {
194        let html = "<html><body><a href=\"/about\">About</a><a href=\"https://example.com/contact\">Contact</a><a href=\"#top\">Top</a></body></html>";
195        let links = WebScraperLoader::extract_links(html, "https://example.com/");
196        assert!(links.contains(&"https://example.com/about".to_string()));
197        assert!(links.contains(&"https://example.com/contact".to_string()));
198        // # links should be filtered
199        assert!(!links.iter().any(|l| l.contains('#')));
200    }
201
202    #[test]
203    fn test_extract_domain() {
204        assert_eq!(
205            WebScraperLoader::extract_domain("https://example.com/path"),
206            "example.com"
207        );
208        assert_eq!(
209            WebScraperLoader::extract_domain("http://sub.example.com:8080/path"),
210            "sub.example.com:8080"
211        );
212    }
213
214    #[test]
215    fn test_resolve_url_absolute() {
216        let result =
217            WebScraperLoader::resolve_url("https://example.com/", "https://other.com/page");
218        assert_eq!(result, Some("https://other.com/page".to_string()));
219    }
220
221    #[test]
222    fn test_resolve_url_relative() {
223        let result = WebScraperLoader::resolve_url("https://example.com/dir/page", "other");
224        assert_eq!(result, Some("https://example.com/dir/other".to_string()));
225    }
226
227    #[test]
228    fn test_resolve_url_root_relative() {
229        let result = WebScraperLoader::resolve_url("https://example.com/dir/page", "/root");
230        assert_eq!(result, Some("https://example.com/root".to_string()));
231    }
232
233    #[test]
234    fn test_extract_text() {
235        let html = "<html><body><p>Hello World</p></body></html>";
236        let text = WebScraperLoader::extract_text(html);
237        assert!(text.contains("Hello World"));
238    }
239
240    #[test]
241    fn test_new_creates_single_page_scraper() {
242        let loader = WebScraperLoader::new("https://example.com");
243        assert_eq!(loader.max_depth, 0);
244        assert_eq!(loader.max_pages, 1);
245    }
246
247    #[test]
248    fn test_with_options() {
249        let loader = WebScraperLoader::new("https://example.com")
250            .with_max_depth(2)
251            .with_max_pages(10);
252        assert_eq!(loader.max_depth, 2);
253        assert_eq!(loader.max_pages, 10);
254    }
255}