1use 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
21pub struct HTMLLoader {
23 html: Option<String>,
24 url: Option<String>,
25}
26
27impl HTMLLoader {
28 pub fn new(html: impl Into<String>) -> Self {
30 Self {
31 html: Some(html.into()),
32 url: None,
33 }
34 }
35
36 pub fn from_url(url: impl Into<String>) -> Self {
38 Self {
39 html: None,
40 url: Some(url.into()),
41 }
42 }
43
44 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 text = text
52 .replace("&", "&")
53 .replace("<", "<")
54 .replace(">", ">")
55 .replace(" ", " ")
56 .replace(""", "\"")
57 .replace("'", "'");
58 WHITESPACE_RE.replace_all(&text, " ").trim().to_string()
60 }
61
62 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 & b < 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>"hello" 'world'</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 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 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>© 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}