Skip to main content

lc_rag/loaders/
html.rs

1//! HTML document loader
2//!
3//! Loads documents from an HTML string or URL: strips script/style, removes tags,
4//! decodes entities, and extracts plain text.
5
6use std::collections::HashMap;
7use std::sync::LazyLock;
8
9use async_trait::async_trait;
10use regex::Regex;
11
12use super::{DocumentLoader, LoaderError};
13use lc_vector_stores::Document;
14
15static SCRIPT_RE: LazyLock<Regex> =
16    LazyLock::new(|| Regex::new(r"(?s)<script.*?</script>").unwrap());
17static STYLE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"(?s)<style.*?</style>").unwrap());
18static TAG_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"<[^>]+>").unwrap());
19static WHITESPACE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\s+").unwrap());
20
21/// HTML loader: strips script/style, removes tags, decodes entities, extracts plain text
22pub struct HTMLLoader {
23    html: Option<String>,
24    url: Option<String>,
25}
26
27impl HTMLLoader {
28    /// Creates a loader from an HTML string
29    pub fn new(html: impl Into<String>) -> Self {
30        Self {
31            html: Some(html.into()),
32            url: None,
33        }
34    }
35
36    /// Creates a loader from a URL (fetches the HTML asynchronously, then parses it)
37    pub fn from_url(url: impl Into<String>) -> Self {
38        Self {
39            html: None,
40            url: Some(url.into()),
41        }
42    }
43
44    /// Extracts plain text from HTML (a pure function, convenient for testing)
45    pub fn extract_text(html: &str) -> String {
46        let mut text = html.to_string();
47        text = SCRIPT_RE.replace_all(&text, "").to_string();
48        text = STYLE_RE.replace_all(&text, "").to_string();
49        text = TAG_RE.replace_all(&text, " ").to_string();
50        // Decode common entities
51        text = text
52            .replace("&amp;", "&")
53            .replace("&lt;", "<")
54            .replace("&gt;", ">")
55            .replace("&nbsp;", " ")
56            .replace("&quot;", "\"")
57            .replace("&#39;", "'");
58        // Compress whitespace
59        WHITESPACE_RE.replace_all(&text, " ").trim().to_string()
60    }
61
62    /// Fetches HTML content from a URL
63    async fn fetch_html(url: &str) -> Result<String, LoaderError> {
64        let response = reqwest::get(url)
65            .await
66            .map_err(|e| LoaderError::Other(format!("HTTP request failed: {}", e)))?;
67        let status = response.status();
68        if !status.is_success() {
69            return Err(LoaderError::Other(format!("HTTP error: {}", status)));
70        }
71        response
72            .text()
73            .await
74            .map_err(|e| LoaderError::Other(format!("failed to read response: {}", e)))
75    }
76}
77
78#[async_trait]
79impl DocumentLoader for HTMLLoader {
80    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
81        let html = if let Some(ref html) = self.html {
82            html.clone()
83        } else if let Some(ref url) = self.url {
84            Self::fetch_html(url).await?
85        } else {
86            return Err(LoaderError::Other(
87                "HTMLLoader has neither html nor url set".to_string(),
88            ));
89        };
90
91        let text = Self::extract_text(&html);
92        let mut metadata = HashMap::new();
93        metadata.insert("format".to_string(), "html".to_string().into());
94        if let Some(ref url) = self.url {
95            metadata.insert("source".to_string(), url.clone().into());
96        }
97        Ok(vec![Document {
98            content: text,
99            metadata,
100            id: None,
101        }])
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn test_extract_text_removes_scripts_and_styles() {
111        let html = r#"<html><head><script>alert(1)</script><style>body{}</style></head><body><p>Hello</p></body></html>"#;
112        let text = HTMLLoader::extract_text(html);
113        assert!(text.contains("Hello"));
114        assert!(!text.contains("alert"));
115        assert!(!text.contains("body{}"));
116        assert!(!text.contains('<'));
117    }
118
119    #[test]
120    fn test_extract_text_decodes_entities() {
121        let html = "<p>a &amp; b &lt; c</p>";
122        let text = HTMLLoader::extract_text(html);
123        assert_eq!(text, "a & b < c");
124    }
125
126    #[test]
127    fn test_extract_text_decodes_more_entities() {
128        let html = "<p>&quot;hello&quot; &#39;world&#39;</p>";
129        let text = HTMLLoader::extract_text(html);
130        assert_eq!(text, "\"hello\" 'world'");
131    }
132
133    #[test]
134    fn test_extract_text_compresses_whitespace() {
135        let html = "<p>hello</p>\n\n<p>world</p>";
136        let text = HTMLLoader::extract_text(html);
137        assert_eq!(text, "hello world");
138    }
139
140    #[tokio::test]
141    async fn test_load_returns_document() {
142        let loader = HTMLLoader::new("<p>test</p>");
143        let docs = loader.load().await.unwrap();
144        assert_eq!(docs.len(), 1);
145        assert_eq!(docs[0].content, "test");
146        assert_eq!(
147            docs[0].metadata.get("format"),
148            Some(&serde_json::Value::String("html".to_string()))
149        );
150    }
151
152    #[tokio::test]
153    async fn test_load_from_url_has_source_metadata() {
154        let loader = HTMLLoader::from_url("https://example.com");
155        // No actual request is made; only verify construction
156        assert!(loader.url.is_some());
157        assert!(loader.html.is_none());
158        assert_eq!(loader.url.as_deref(), Some("https://example.com"));
159    }
160
161    #[tokio::test]
162    async fn test_load_from_url_invalid_url() {
163        let loader = HTMLLoader::from_url("http://nonexistent.invalid.example");
164        let result = loader.load().await;
165        assert!(result.is_err());
166    }
167
168    #[tokio::test]
169    async fn test_load_from_html_with_source_in_metadata() {
170        let loader = HTMLLoader::new("<p>hello</p>");
171        let docs = loader.load().await.unwrap();
172        // No source when loading from an HTML string
173        assert!(!docs[0].metadata.contains_key("source"));
174    }
175
176    #[test]
177    fn test_extract_text_empty() {
178        assert_eq!(HTMLLoader::extract_text(""), "");
179    }
180
181    #[test]
182    fn test_extract_text_nested_tags() {
183        let html = "<div><p><b>bold</b> text</p></div>";
184        let text = HTMLLoader::extract_text(html);
185        assert_eq!(text, "bold text");
186    }
187
188    #[test]
189    fn test_extract_text_realistic_page() {
190        let html = r#"<!DOCTYPE html>
191<html lang="en">
192<head>
193    <meta charset="UTF-8">
194    <title>Test Page</title>
195    <script src="app.js"></script>
196    <style>body { margin: 0; }</style>
197</head>
198<body>
199    <h1>Welcome</h1>
200    <p>This is a <strong>test</strong> page.</p>
201    <footer>&copy; 2026</footer>
202</body>
203</html>"#;
204        let text = HTMLLoader::extract_text(html);
205        assert!(text.contains("Welcome"));
206        assert!(text.contains("test page"));
207        assert!(!text.contains("app.js"));
208        assert!(!text.contains("margin"));
209        assert!(!text.contains("DOCTYPE"));
210    }
211}