lc_rag/loaders/
html.rs

1//! HTML 文档加载器
2//!
3//! 支持从 HTML 字符串或 URL 加载文档,去除 script/style,剥离标签,解码实体,提取纯文本。
4
5use std::collections::HashMap;
6use std::sync::LazyLock;
7
8use async_trait::async_trait;
9use regex::Regex;
10
11use super::{DocumentLoader, LoaderError};
12use lc_vector_stores::Document;
13
14static SCRIPT_RE: LazyLock<Regex> =
15    LazyLock::new(|| Regex::new(r"(?s)<script.*?</script>").unwrap());
16static STYLE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"(?s)<style.*?</style>").unwrap());
17static TAG_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"<[^>]+>").unwrap());
18static WHITESPACE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\s+").unwrap());
19
20/// HTML 加载器:去除 script/style,剥离标签,解码实体,提取纯文本
21pub struct HTMLLoader {
22    html: Option<String>,
23    url: Option<String>,
24}
25
26impl HTMLLoader {
27    /// 从 HTML 字符串创建加载器
28    pub fn new(html: impl Into<String>) -> Self {
29        Self {
30            html: Some(html.into()),
31            url: None,
32        }
33    }
34
35    /// 从 URL 创建加载器(异步抓取 HTML 后解析)
36    pub fn from_url(url: impl Into<String>) -> Self {
37        Self {
38            html: None,
39            url: Some(url.into()),
40        }
41    }
42
43    /// 从 HTML 提取纯文本(纯函数,便于测试)
44    pub fn extract_text(html: &str) -> String {
45        let mut text = html.to_string();
46        text = SCRIPT_RE.replace_all(&text, "").to_string();
47        text = STYLE_RE.replace_all(&text, "").to_string();
48        text = TAG_RE.replace_all(&text, " ").to_string();
49        // 解码常见实体
50        text = text
51            .replace("&amp;", "&")
52            .replace("&lt;", "<")
53            .replace("&gt;", ">")
54            .replace("&nbsp;", " ")
55            .replace("&quot;", "\"")
56            .replace("&#39;", "'");
57        // 压缩空白
58        WHITESPACE_RE.replace_all(&text, " ").trim().to_string()
59    }
60
61    /// 从 URL 抓取 HTML 内容
62    async fn fetch_html(url: &str) -> Result<String, LoaderError> {
63        let response = reqwest::get(url)
64            .await
65            .map_err(|e| LoaderError::Other(format!("HTTP 请求失败: {}", e)))?;
66        let status = response.status();
67        if !status.is_success() {
68            return Err(LoaderError::Other(format!("HTTP 错误: {}", status)));
69        }
70        response
71            .text()
72            .await
73            .map_err(|e| LoaderError::Other(format!("读取响应失败: {}", e)))
74    }
75}
76
77#[async_trait]
78impl DocumentLoader for HTMLLoader {
79    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
80        let html = if let Some(ref html) = self.html {
81            html.clone()
82        } else if let Some(ref url) = self.url {
83            Self::fetch_html(url).await?
84        } else {
85            return Err(LoaderError::Other(
86                "HTMLLoader 未设置 html 或 url".to_string(),
87            ));
88        };
89
90        let text = Self::extract_text(&html);
91        let mut metadata = HashMap::new();
92        metadata.insert("format".to_string(), "html".to_string());
93        if let Some(ref url) = self.url {
94            metadata.insert("source".to_string(), url.clone());
95        }
96        Ok(vec![Document {
97            content: text,
98            metadata,
99            id: None,
100        }])
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn test_extract_text_removes_scripts_and_styles() {
110        let html = r#"<html><head><script>alert(1)</script><style>body{}</style></head><body><p>Hello</p></body></html>"#;
111        let text = HTMLLoader::extract_text(html);
112        assert!(text.contains("Hello"));
113        assert!(!text.contains("alert"));
114        assert!(!text.contains("body{}"));
115        assert!(!text.contains('<'));
116    }
117
118    #[test]
119    fn test_extract_text_decodes_entities() {
120        let html = "<p>a &amp; b &lt; c</p>";
121        let text = HTMLLoader::extract_text(html);
122        assert_eq!(text, "a & b < c");
123    }
124
125    #[test]
126    fn test_extract_text_decodes_more_entities() {
127        let html = "<p>&quot;hello&quot; &#39;world&#39;</p>";
128        let text = HTMLLoader::extract_text(html);
129        assert_eq!(text, "\"hello\" 'world'");
130    }
131
132    #[test]
133    fn test_extract_text_compresses_whitespace() {
134        let html = "<p>hello</p>\n\n<p>world</p>";
135        let text = HTMLLoader::extract_text(html);
136        assert_eq!(text, "hello world");
137    }
138
139    #[tokio::test]
140    async fn test_load_returns_document() {
141        let loader = HTMLLoader::new("<p>test</p>");
142        let docs = loader.load().await.unwrap();
143        assert_eq!(docs.len(), 1);
144        assert_eq!(docs[0].content, "test");
145        assert_eq!(docs[0].metadata.get("format"), Some(&"html".to_string()));
146    }
147
148    #[tokio::test]
149    async fn test_load_from_url_has_source_metadata() {
150        let loader = HTMLLoader::from_url("https://example.com");
151        // 不实际请求,只验证构造
152        assert!(loader.url.is_some());
153        assert!(loader.html.is_none());
154        assert_eq!(loader.url.as_deref(), Some("https://example.com"));
155    }
156
157    #[tokio::test]
158    async fn test_load_from_url_invalid_url() {
159        let loader = HTMLLoader::from_url("http://nonexistent.invalid.example");
160        let result = loader.load().await;
161        assert!(result.is_err());
162    }
163
164    #[tokio::test]
165    async fn test_load_from_html_with_source_in_metadata() {
166        let loader = HTMLLoader::new("<p>hello</p>");
167        let docs = loader.load().await.unwrap();
168        // 从 HTML 字符串加载时没有 source
169        assert!(!docs[0].metadata.contains_key("source"));
170    }
171
172    #[test]
173    fn test_extract_text_empty() {
174        assert_eq!(HTMLLoader::extract_text(""), "");
175    }
176
177    #[test]
178    fn test_extract_text_nested_tags() {
179        let html = "<div><p><b>bold</b> text</p></div>";
180        let text = HTMLLoader::extract_text(html);
181        assert_eq!(text, "bold text");
182    }
183
184    #[test]
185    fn test_extract_text_realistic_page() {
186        let html = r#"<!DOCTYPE html>
187<html lang="en">
188<head>
189    <meta charset="UTF-8">
190    <title>Test Page</title>
191    <script src="app.js"></script>
192    <style>body { margin: 0; }</style>
193</head>
194<body>
195    <h1>Welcome</h1>
196    <p>This is a <strong>test</strong> page.</p>
197    <footer>&copy; 2026</footer>
198</body>
199</html>"#;
200        let text = HTMLLoader::extract_text(html);
201        assert!(text.contains("Welcome"));
202        assert!(text.contains("test page"));
203        assert!(!text.contains("app.js"));
204        assert!(!text.contains("margin"));
205        assert!(!text.contains("DOCTYPE"));
206    }
207}
lc_rag/loaders/html.rs

lc_rag/loaders/
html.rs
