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