lc_rag/loaders/
web_scraper.rs

1//! 网页爬取加载器
2//!
3//! 从 URL 爬取网页内容,提取正文文本,支持递归链接跟踪。
4//! 基于 HTMLLoader 的文本提取逻辑,增加链接发现与批量爬取能力。
5
6use std::collections::{HashMap, HashSet};
7use std::sync::LazyLock;
8
9use async_trait::async_trait;
10
11use super::{DocumentLoader, LoaderError};
12use lc_vector_stores::Document;
13
14// M9: Pre-compile regexes once instead of on every call.
15static HREF_RE: LazyLock<regex::Regex> =
16    LazyLock::new(|| regex::Regex::new(r#"href\s*=\s*["']([^"']+)["']"#).unwrap());
17static DOMAIN_RE: LazyLock<regex::Regex> =
18    LazyLock::new(|| regex::Regex::new(r"https?://([^/]+)").unwrap());
19static DOMAIN_PREFIX_RE: LazyLock<regex::Regex> =
20    LazyLock::new(|| regex::Regex::new(r"https?://[^/]+").unwrap());
21
22/// 网页爬取加载器
23///
24/// 从 URL 爬取网页,提取正文文本。可选递归跟踪同域链接。
25pub struct WebScraperLoader {
26    /// 起始 URL
27    url: String,
28    /// 最大递归深度(0 = 仅爬起始页)
29    max_depth: usize,
30    /// 最大爬取页面数
31    max_pages: usize,
32    /// 是否在爬取失败时返回错误(默认 false,跳过失败页面)
33    fail_on_error: bool,
34}
35
36impl WebScraperLoader {
37    /// 从 URL 创建加载器(仅爬取指定页面)
38    pub fn new(url: impl Into<String>) -> Self {
39        Self {
40            url: url.into(),
41            max_depth: 0,
42            max_pages: 1,
43            fail_on_error: false,
44        }
45    }
46
47    /// 设置最大递归深度
48    pub fn with_max_depth(mut self, depth: usize) -> Self {
49        self.max_depth = depth;
50        self
51    }
52
53    /// 设置最大爬取页面数
54    pub fn with_max_pages(mut self, pages: usize) -> Self {
55        self.max_pages = pages;
56        self
57    }
58
59    /// 设置爬取失败时是否返回错误(默认跳过失败页面)
60    pub fn with_fail_on_error(mut self, fail: bool) -> Self {
61        self.fail_on_error = fail;
62        self
63    }
64
65    /// 从 HTML 提取纯文本(复用 HTMLLoader 的逻辑)
66    fn extract_text(html: &str) -> String {
67        super::HTMLLoader::extract_text(html)
68    }
69
70    /// 从 HTML 提取链接
71    fn extract_links(html: &str, base_url: &str) -> Vec<String> {
72        let base_domain = Self::extract_domain(base_url);
73        HREF_RE
74            .captures_iter(html)
75            .filter_map(|cap| cap.get(1).map(|m| m.as_str().to_string()))
76            .filter(|link| !link.starts_with('#') && !link.starts_with("javascript:"))
77            .filter_map(|link| Self::resolve_url(base_url, &link))
78            .filter(|url| Self::extract_domain(url) == base_domain)
79            .collect()
80    }
81
82    /// 提取域名
83    fn extract_domain(url: &str) -> String {
84        DOMAIN_RE
85            .captures(url)
86            .and_then(|c| c.get(1).map(|m| m.as_str().to_string()))
87            .unwrap_or_default()
88    }
89
90    /// 解析相对 URL 为绝对 URL
91    fn resolve_url(base: &str, href: &str) -> Option<String> {
92        if href.starts_with("http://") || href.starts_with("https://") {
93            Some(href.to_string())
94        } else if href.starts_with('/') {
95            // 找到 scheme://domain 部分
96            let domain = DOMAIN_PREFIX_RE.find(base)?.as_str();
97            Some(format!("{}{}", domain, href))
98        } else {
99            // 相对路径
100            let base_dir = base.rfind('/').map(|i| &base[..=i]).unwrap_or(base);
101            Some(format!("{}{}", base_dir, href))
102        }
103    }
104
105    /// 爬取单个页面
106    async fn fetch_page(url: &str) -> Result<(String, String), LoaderError> {
107        let response = reqwest::get(url)
108            .await
109            .map_err(|e| LoaderError::Other(format!("HTTP 请求失败 {}: {}", url, e)))?;
110        let status = response.status();
111        if !status.is_success() {
112            return Err(LoaderError::Other(format!("HTTP 错误 {}: {}", url, status)));
113        }
114        let html = response
115            .text()
116            .await
117            .map_err(|e| LoaderError::Other(format!("读取响应失败 {}: {}", url, e)))?;
118        Ok((url.to_string(), html))
119    }
120}
121
122#[async_trait]
123impl DocumentLoader for WebScraperLoader {
124    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
125        let mut documents = Vec::new();
126        let mut visited = HashSet::new();
127        let mut queue = vec![(self.url.clone(), 0usize)];
128        let mut failed_count: usize = 0;
129
130        while let Some((url, depth)) = queue.pop() {
131            if visited.contains(&url) || documents.len() >= self.max_pages {
132                continue;
133            }
134            visited.insert(url.clone());
135
136            let (fetched_url, html) = match Self::fetch_page(&url).await {
137                Ok(r) => r,
138                Err(e) => {
139                    failed_count += 1;
140                    if self.fail_on_error {
141                        return Err(e);
142                    }
143                    // 跳过失败页面,继续爬取其他
144                    eprintln!(
145                        "警告: 爬取 {} 失败 (第 {} 个失败): {}",
146                        url, failed_count, e
147                    );
148                    continue;
149                }
150            };
151
152            let text = Self::extract_text(&html);
153
154            let mut metadata = HashMap::new();
155            metadata.insert("format".to_string(), "html".to_string());
156            metadata.insert("source".to_string(), fetched_url.clone());
157
158            documents.push(Document {
159                content: text,
160                metadata,
161                id: None,
162            });
163
164            // 递归跟踪链接
165            if depth < self.max_depth {
166                let links = Self::extract_links(&html, &fetched_url);
167                for link in links {
168                    if !visited.contains(&link) {
169                        queue.push((link, depth + 1));
170                    }
171                }
172            }
173        }
174
175        if failed_count > 0 {
176            eprintln!(
177                "警告: 爬取完成,共 {} 个页面失败,{} 个页面成功",
178                failed_count,
179                documents.len()
180            );
181        }
182
183        Ok(documents)
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn test_extract_links() {
193        let html = "<html><body><a href=\"/about\">About</a><a href=\"https://example.com/contact\">Contact</a><a href=\"#top\">Top</a></body></html>";
194        let links = WebScraperLoader::extract_links(html, "https://example.com/");
195        assert!(links.contains(&"https://example.com/about".to_string()));
196        assert!(links.contains(&"https://example.com/contact".to_string()));
197        // # 链接应被过滤
198        assert!(!links.iter().any(|l| l.contains('#')));
199    }
200
201    #[test]
202    fn test_extract_domain() {
203        assert_eq!(
204            WebScraperLoader::extract_domain("https://example.com/path"),
205            "example.com"
206        );
207        assert_eq!(
208            WebScraperLoader::extract_domain("http://sub.example.com:8080/path"),
209            "sub.example.com:8080"
210        );
211    }
212
213    #[test]
214    fn test_resolve_url_absolute() {
215        let result =
216            WebScraperLoader::resolve_url("https://example.com/", "https://other.com/page");
217        assert_eq!(result, Some("https://other.com/page".to_string()));
218    }
219
220    #[test]
221    fn test_resolve_url_relative() {
222        let result = WebScraperLoader::resolve_url("https://example.com/dir/page", "other");
223        assert_eq!(result, Some("https://example.com/dir/other".to_string()));
224    }
225
226    #[test]
227    fn test_resolve_url_root_relative() {
228        let result = WebScraperLoader::resolve_url("https://example.com/dir/page", "/root");
229        assert_eq!(result, Some("https://example.com/root".to_string()));
230    }
231
232    #[test]
233    fn test_extract_text() {
234        let html = "<html><body><p>Hello World</p></body></html>";
235        let text = WebScraperLoader::extract_text(html);
236        assert!(text.contains("Hello World"));
237    }
238
239    #[test]
240    fn test_new_creates_single_page_scraper() {
241        let loader = WebScraperLoader::new("https://example.com");
242        assert_eq!(loader.max_depth, 0);
243        assert_eq!(loader.max_pages, 1);
244    }
245
246    #[test]
247    fn test_with_options() {
248        let loader = WebScraperLoader::new("https://example.com")
249            .with_max_depth(2)
250            .with_max_pages(10);
251        assert_eq!(loader.max_depth, 2);
252        assert_eq!(loader.max_pages, 10);
253    }
254}
lc_rag/loaders/web_scraper.rs

lc_rag/loaders/
web_scraper.rs
