1use 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
16const 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
25pub struct HTMLLoader {
27 html: Option<String>,
28 url: Option<String>,
29}
30
31impl HTMLLoader {
32 pub fn new(html: impl Into<String>) -> Self {
34 Self {
35 html: Some(html.into()),
36 url: None,
37 }
38 }
39
40 pub fn from_url(url: impl Into<String>) -> Self {
42 Self {
43 html: None,
44 url: Some(url.into()),
45 }
46 }
47
48 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 text = text
56 .replace("&", "&")
57 .replace("<", "<")
58 .replace(">", ">")
59 .replace(" ", " ")
60 .replace(""", "\"")
61 .replace("'", "'");
62 WHITESPACE_RE.replace_all(&text, " ").trim().to_string()
64 }
65
66 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 & b < 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>"hello" 'world'</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 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 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>© 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}