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!("构建 HTTP 客户端失败: {}", e)))?;
124        let response = client
125            .get(url)
126            .send()
127            .await
128            .map_err(|e| LoaderError::Other(format!("HTTP 请求失败 {}: {}", url, e)))?;
129        let status = response.status();
130        if !status.is_success() {
131            return Err(LoaderError::Other(format!("HTTP 错误 {}: {}", url, status)));
132        }
133        let html = response
134            .text()
135            .await
136            .map_err(|e| LoaderError::Other(format!("读取响应失败 {}: {}", url, e)))?;
137        Ok((url.to_string(), html))
138    }
139}
140
141#[async_trait]
142impl DocumentLoader for WebScraperLoader {
143    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
144        let mut documents = Vec::new();
145        let mut visited = HashSet::new();
146        let mut queue = vec![(self.url.clone(), 0usize)];
147        let mut failed_count: usize = 0;
148
149        while let Some((url, depth)) = queue.pop() {
150            if visited.contains(&url) || documents.len() >= self.max_pages {
151                continue;
152            }
153            visited.insert(url.clone());
154
155            let (fetched_url, html) = match Self::fetch_page(&url, self.timeout).await {
156                Ok(r) => r,
157                Err(e) => {
158                    failed_count += 1;
159                    if self.fail_on_error {
160                        return Err(e);
161                    }
162                    // 跳过失败页面,继续爬取其他(经日志门面暴露,便于宿主捕获)
163                    log::warn!("爬取 {} 失败 (第 {} 个失败): {}", url, failed_count, e);
164                    continue;
165                }
166            };
167
168            let text = Self::extract_text(&html);
169
170            let mut metadata = HashMap::new();
171            metadata.insert("format".to_string(), "html".to_string());
172            metadata.insert("source".to_string(), fetched_url.clone());
173
174            documents.push(Document {
175                content: text,
176                metadata,
177                id: None,
178            });
179
180            // 递归跟踪链接
181            if depth < self.max_depth {
182                let links = Self::extract_links(&html, &fetched_url);
183                for link in links {
184                    if !visited.contains(&link) {
185                        queue.push((link, depth + 1));
186                    }
187                }
188            }
189        }
190
191        if failed_count > 0 {
192            log::warn!(
193                "爬取完成,共 {} 个页面失败,{} 个页面成功",
194                failed_count,
195                documents.len()
196            );
197        }
198
199        Ok(documents)
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206
207    #[test]
208    fn test_extract_links() {
209        let html = "<html><body><a href=\"/about\">About</a><a href=\"https://example.com/contact\">Contact</a><a href=\"#top\">Top</a></body></html>";
210        let links = WebScraperLoader::extract_links(html, "https://example.com/");
211        assert!(links.contains(&"https://example.com/about".to_string()));
212        assert!(links.contains(&"https://example.com/contact".to_string()));
213        // # 链接应被过滤
214        assert!(!links.iter().any(|l| l.contains('#')));
215    }
216
217    #[test]
218    fn test_extract_domain() {
219        assert_eq!(
220            WebScraperLoader::extract_domain("https://example.com/path"),
221            "example.com"
222        );
223        assert_eq!(
224            WebScraperLoader::extract_domain("http://sub.example.com:8080/path"),
225            "sub.example.com:8080"
226        );
227    }
228
229    #[test]
230    fn test_resolve_url_absolute() {
231        let result =
232            WebScraperLoader::resolve_url("https://example.com/", "https://other.com/page");
233        assert_eq!(result, Some("https://other.com/page".to_string()));
234    }
235
236    #[test]
237    fn test_resolve_url_relative() {
238        let result = WebScraperLoader::resolve_url("https://example.com/dir/page", "other");
239        assert_eq!(result, Some("https://example.com/dir/other".to_string()));
240    }
241
242    #[test]
243    fn test_resolve_url_root_relative() {
244        let result = WebScraperLoader::resolve_url("https://example.com/dir/page", "/root");
245        assert_eq!(result, Some("https://example.com/root".to_string()));
246    }
247
248    #[test]
249    fn test_extract_text() {
250        let html = "<html><body><p>Hello World</p></body></html>";
251        let text = WebScraperLoader::extract_text(html);
252        assert!(text.contains("Hello World"));
253    }
254
255    #[test]
256    fn test_new_creates_single_page_scraper() {
257        let loader = WebScraperLoader::new("https://example.com");
258        assert_eq!(loader.max_depth, 0);
259        assert_eq!(loader.max_pages, 1);
260    }
261
262    #[test]
263    fn test_with_options() {
264        let loader = WebScraperLoader::new("https://example.com")
265            .with_max_depth(2)
266            .with_max_pages(10);
267        assert_eq!(loader.max_depth, 2);
268        assert_eq!(loader.max_pages, 10);
269    }
270}
lc_rag/loaders/web_scraper.rs

lc_rag/loaders/
web_scraper.rs
