lc_rag/loaders/
web_scraper.rs

1//! 网页爬取加载器
2//!
3//! 从 URL 爬取网页内容,提取正文文本,支持递归链接跟踪。
4//! 基于 HTMLLoader 的文本提取逻辑,增加链接发现与批量爬取能力。
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: 默认单次 HTTP 请求超时——目标站挂起时爬虫不会永久阻塞。
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/// 网页爬取加载器
27///
28/// 从 URL 爬取网页,提取正文文本。可选递归跟踪同域链接。
29pub struct WebScraperLoader {
30    /// 起始 URL
31    url: String,
32    /// 最大递归深度(0 = 仅爬起始页)
33    max_depth: usize,
34    /// 最大爬取页面数
35    max_pages: usize,
36    /// 是否在爬取失败时返回错误(默认 false,跳过失败页面)
37    fail_on_error: bool,
38    /// H8: 单次 HTTP 请求超时,防目标站挂起导致爬虫永久阻塞
39    timeout: Duration,
40}
41
42impl WebScraperLoader {
43    /// 从 URL 创建加载器(仅爬取指定页面)
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    /// 设置最大递归深度
55    pub fn with_max_depth(mut self, depth: usize) -> Self {
56        self.max_depth = depth;
57        self
58    }
59
60    /// 设置最大爬取页面数
61    pub fn with_max_pages(mut self, pages: usize) -> Self {
62        self.max_pages = pages;
63        self
64    }
65
66    /// 设置爬取失败时是否返回错误(默认跳过失败页面)
67    pub fn with_fail_on_error(mut self, fail: bool) -> Self {
68        self.fail_on_error = fail;
69        self
70    }
71
72    /// 设置单次 HTTP 请求超时(H8,默认 30s)
73    pub fn with_timeout(mut self, timeout: Duration) -> Self {
74        self.timeout = timeout;
75        self
76    }
77
78    /// 从 HTML 提取纯文本(复用 HTMLLoader 的逻辑)
79    fn extract_text(html: &str) -> String {
80        super::HTMLLoader::extract_text(html)
81    }
82
83    /// 从 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    /// 提取域名
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    /// 解析相对 URL 为绝对 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            // 找到 scheme://domain 部分
109            let domain = DOMAIN_PREFIX_RE.find(base)?.as_str();
110            Some(format!("{}{}", domain, href))
111        } else {
112            // 相对路径
113            let base_dir = base.rfind('/').map(|i| &base[..=i]).unwrap_or(base);
114            Some(format!("{}{}", base_dir, href))
115        }
116    }
117
118    /// 爬取单个页面
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                    // 跳过失败页面,继续爬取其他(经日志门面暴露,便于宿主捕获)
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            // 递归跟踪链接
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        // # 链接应被过滤
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}
lc_rag/loaders/web_scraper.rs

lc_rag/loaders/
web_scraper.rs
