Repository: StarlightSearch/EmbedAnything (language: python)

Issue:
Split HTML parsing from website parsing
It seems like this library is capable of parsing HTML, but right now it'll only allow you to do so through a URL.

It'd be cool if HTML parsing was something we could do separately, instead of requiring a URL and a web request!


Discussion / hints:
If this is something you're interested in adding, I'm happy to open a PR for the Rust side of things!
Hey sure.. this is a good idea to parse downloaded documentation etc. Feel free to add an HTML processor in the `file_processor` module. 

You can open a PR. Also try to add an example use case. You can create a simple test HTM in test_files.

Fix patch:
diff --git a/rust/src/file_processor/html_processor.rs b/rust/src/file_processor/html_processor.rs
new file mode 100644
index 0000000..a19b78e
--- /dev/null
+++ b/rust/src/file_processor/html_processor.rs
@@ -0,0 +1,216 @@
+use anyhow::Result;
+use std::collections::{HashMap, HashSet};
+use std::rc::Rc;
+use scraper::{Html, Selector};
+use serde_json::json;
+use url::Url;
+use crate::embeddings::embed::{EmbedData, Embedder};
+use crate::embeddings::get_text_metadata;
+use crate::text_loader::{SplittingStrategy, TextLoader};
+
+#[derive(Debug)]
+pub struct HtmlDocument {
+    pub origin: Option<String>,
+    pub title: Option<String>,
+    pub headers: Option<Vec<String>>,
+    pub paragraphs: Option<Vec<String>>,
+    pub codes: Option<Vec<String>>,
+    pub links: Option<HashSet<String>>,
+}
+
+impl HtmlDocument {
+    pub async fn embed_webpage(
+        &self,
+        embeder: &Embedder,
+        chunk_size: usize,
+        batch_size: Option<usize>,
+    ) -> Result<Vec<EmbedData>> {
+        let mut embed_data = Vec::new();
+
+        if let Some(paragraphs) = &self.paragraphs {
+            embed_data.extend(
+                self.embed_tag("p", paragraphs, embeder, chunk_size, batch_size)
+                    .await?,
+            );
+        }
+
+        if let Some(headers) = &self.headers {
+            embed_data.extend(
+                self.embed_tag("h1", headers, embeder, chunk_size, batch_size)
+                    .await?,
+            );
+        }
+
+        if let Some(codes) = &self.codes {
+            embed_data.extend(
+                self.embed_tag("code", codes, embeder, chunk_size, batch_size)
+                    .await?,
+            );
+        }
+
+        Ok(embed_data)
+    }
+
+    pub async fn embed_tag(
+        &self,
+        tag: &str,
+        tag_content: &[String],
+        embeder: &Embedder,
+        chunk_size: usize,
+        batch_size: Option<usize>,
+    ) -> Result<Vec<EmbedData>> {
+        let mut embed_data = Vec::new();
+
+        for content in tag_content {
+            let textloader = TextLoader::new(chunk_size);
+            let chunks =
+                match textloader.split_into_chunks(content, SplittingStrategy::Sentence, None) {
+                    Some(chunks) => chunks,
+                    None => continue,
+                };
+
+            if chunks.is_empty() {
+                continue;
+            }
+
+            let tag_type = match tag {
+                "h1" => "header",
+                "h2" => "subheader",
+                "h3" => "subsubheader",
+                "p" => "paragraph",
+                "code" => "code",
+                _ => "paragraph",
+            };
+
+            let metadata = json!({
+                "url": self.origin,
+                "type": tag_type,
+                "full_text": content,
+            });
+
+            let metadata_hashmap: HashMap<String, String> = serde_json::from_value(metadata)?;
+
+            let encodings = embeder.embed(&chunks, batch_size).await?;
+            let embeddings =
+                get_text_metadata(&Rc::new(encodings), &chunks, &Some(metadata_hashmap))?;
+            embed_data.extend(embeddings);
+        }
+
+        Ok(embed_data)
+    }
+}
+
+/// A Struct for processing HTML files.
+pub struct HtmlProcessor;
+
+impl HtmlProcessor {
+    pub fn new() -> Self {
+        Self {}
+    }
+
+    /// Extracts the contents of an HTML file.
+    /// 
+    /// # Arguments
+    /// 
+    /// * `file_path` - The path to the HTML file.
+    /// * `origin` - The original URL of the HTML page, if any. This is required for extracting links.
+    /// 
+    /// # Returns
+    /// 
+    /// Returns a `Result` containing the extracted content as an `HtmlDocument` if successful,
+    /// or an `Error` if an error occurred during any part of the process.
+    pub fn process_html_file(&self, file_path: impl AsRef<std::path::Path>, origin: Option<impl Into<String>>) -> Result<HtmlDocument> {
+        // check if https is in the website. If not, add it.
+        let bytes = std::fs::read(file_path)?;
+        let out = String::from_utf8_lossy(&bytes);
+        self.process_html(out, origin)
+    }
+
+    /// Extracts the contents of an HTML text.
+    ///
+    /// # Arguments
+    ///
+    /// * `html` - The HTML text to be extracted.
+    /// * `origin` - The original URL of the HTML, if any. This is required for extracting links.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Result` containing the extracted content as an `HtmlDocument` if successful,
+    /// or an `Error` if an error occurred during any part of the process.
+    pub fn process_html(&self, html: impl Into<String>, origin: Option<impl Into<String>>) -> Result<HtmlDocument> {
+        // check if https is in the website. If not, add it.
+        let document = Html::parse_document(&html.into());
+        let headers = self.get_text_from_tag("h1,h2,h3", &document)?;
+        let paragraphs = self.get_text_from_tag("p", &document)?;
+        let codes = self.get_text_from_tag("code", &document)?;
+        let origin = origin.map(Into::into);
+        let links = match &origin {
+            Some(origin) => Some(self.extract_links(&origin.clone(), &document)?),
+            None => None
+        };
+        let title = self.get_title(&document)?;
+        let web_page = HtmlDocument {
+            origin,
+            title,
+            headers: Some(headers),
+            paragraphs: Some(paragraphs),
+            codes: Some(codes),
+            links,
+        };
+
+        Ok(web_page)
+    }
+
+    fn get_text_from_tag(&self, tag: &str, document: &Html) -> Result<Vec<String>> {
+        let selector = Selector::parse(tag).expect("invalid selector for tag");
+        Ok(document
+            .select(&selector)
+            .map(|element| element.text().collect::<String>().trim().to_string())
+            .collect())
+    }
+
+    fn extract_links(&self, website: &str, document: &Html) -> Result<HashSet<String>> {
+        let mut links = HashSet::new();
+        let base_url = Url::parse(&website)?;
+
+        for element in document.select(&Selector::parse("a").expect("invalid selector for link")) {
+            if let Some(href) = element.value().attr("href") {
+                let mut link_url = base_url.join(href)?;
+                // Normalize URLs, remove fragments and ensure they are absolute.
+                link_url.set_fragment(None);
+                links.insert(link_url.to_string());
+            }
+        }
+
+        Ok(links)
+    }
+
+    fn get_title(&self, document: &Html) -> Result<Option<String>> {
+        if let Some(title_element) = document.select(&Selector::parse("title").expect("invalid selector for title")).next() {
+            Ok(Some(title_element.text().collect::<String>()))
+        } else {
+            Ok(None)
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_process_html_file() {
+        let html_processor = HtmlProcessor::new();
+        let html_file = "test_files/test.html";
+        let result = html_processor.process_html_file(html_file, Some("https://example.com/"));
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_process_html_file_err() {
+        let html_processor = HtmlProcessor::new();
+        let html_file = "test_files/some_file_that_doesnt_exist.html";
+        let result = html_processor.process_html_file(html_file, Some("https://example.com/"));
+        assert!(result.is_err());
+    }
+}
diff --git a/rust/src/file_processor/mod.rs b/rust/src/file_processor/mod.rs
index ae30a7e..807170a 100644
--- a/rust/src/file_processor/mod.rs
+++ b/rust/src/file_processor/mod.rs
@@ -10,4 +10,7 @@ pub mod txt_processor;
 /// This module contains the processor to process web links.
 pub mod website_processor;
 
+/// This module contains the file processor for HTML files.
+pub mod html_processor;
+
 pub mod audio;
diff --git a/rust/src/file_processor/website_processor.rs b/rust/src/file_processor/website_processor.rs
index 80a1d39..7ac22fd 100644
--- a/rust/src/file_processor/website_processor.rs
+++ b/rust/src/file_processor/website_processor.rs
@@ -4,15 +4,14 @@ use std::{
 };
 
 use anyhow::Result;
-use scraper::{Html, Selector};
 use serde_json::json;
-use url::Url;
 
 use crate::{
     embeddings::{
         embed::{EmbedData, Embedder},
         get_text_metadata,
     },
+    file_processor::html_processor::HtmlProcessor,
     text_loader::{SplittingStrategy, TextLoader},
 };
 
@@ -98,7 +97,7 @@ impl WebPage {
 
             let metadata_hashmap: HashMap<String, String> = serde_json::from_value(metadata)?;
 
-            let encodings = embeder.embed(&chunks, batch_size).await.unwrap();
+            let encodings = embeder.embed(&chunks, batch_size).await?;
             let embeddings =
                 get_text_metadata(&Rc::new(encodings), &chunks, &Some(metadata_hashmap))?;
             embed_data.extend(embeddings);
@@ -127,11 +126,13 @@ impl Default for WebsiteProcessor {
     }
 }
 
-pub struct WebsiteProcessor;
+pub struct WebsiteProcessor {
+    html_processor: HtmlProcessor
+}
 
 impl WebsiteProcessor {
     pub fn new() -> Self {
-        Self {}
+        Self { html_processor: HtmlProcessor::new() }
     }
 
     pub fn process_website(&self, website: &str) -> Result<WebPage> {
@@ -143,55 +144,19 @@ impl WebsiteProcessor {
         };
 
         let response = reqwest::blocking::get(website)?.text()?;
-        let document = Html::parse_document(&response);
-        let headers = self.get_text_from_tag("h1,h2,h3", &document)?;
-        let paragraphs = self.get_text_from_tag("p", &document)?;
-        let codes = self.get_text_from_tag("code", &document)?;
-        let links = self.extract_links(website, &document)?;
-        let title = self.get_title(&document)?;
+        let html_document = self.html_processor.process_html(response, Some(website))?;
+
         let web_page = WebPage {
             url: website.to_string(),
-            title,
-            headers: Some(headers),
-            paragraphs: Some(paragraphs),
-            codes: Some(codes),
-            links: Some(links),
+            title: html_document.title,
+            headers: html_document.headers,
+            paragraphs: html_document.paragraphs,
+            codes: html_document.codes,
+            links: html_document.links,
         };
 
         Ok(web_page)
     }
-
-    fn get_text_from_tag(&self, tag: &str, document: &Html) -> Result<Vec<String>, anyhow::Error> {
-        let selector = Selector::parse(tag).unwrap();
-        Ok(document
-            .select(&selector)
-            .map(|element| element.text().collect::<String>().trim().to_string())
-            .collect())
-    }
-
-    fn extract_links(&self, website: &str, document: &Html) -> Result<HashSet<String>> {
-        let mut links = HashSet::new();
-        let base_url = Url::parse(website)?;
-
-        for element in document.select(&Selector::parse("a").unwrap()) {
-            if let Some(href) = element.value().attr("href") {
-                let mut link_url = base_url.join(href)?;
-                // Normalize URLs, remove fragments and ensure they are absolute.
-                link_url.set_fragment(None);
-                links.insert(link_url.to_string());
-            }
-        }
-
-        Ok(links)
-    }
-
-    fn get_title(&self, document: &Html) -> Result<Option<String>> {
-        if let Some(title_element) = document.select(&Selector::parse("title").unwrap()).next() {
-            Ok(Some(title_element.text().collect::<String>()))
-        } else {
-            Ok(None)
-        }
-    }
 }
 
 #[cfg(test)]
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index b19dcef..638e81b 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -55,7 +55,6 @@ use tokio::sync::mpsc; // Add this at the top of your file
 /// println!("{:?}", embeddings);
 /// ```
 /// This will output the embeddings of the queries using the OpenAI embedding model.
-
 pub async fn embed_query(
     query: Vec<String>,
     embeder: &Embedder,
@@ -71,6 +70,7 @@ pub async fn embed_query(
 
     Ok(embeddings)
 }
+
 /// Embeds the text from a file using the specified embedding model.
 ///
 /// # Arguments
@@ -99,7 +99,6 @@ pub async fn embed_query(
 /// let embeddings = embed_file(file_name, embeder, config).unwrap();
 /// ```
 /// This will output the embeddings of the file using the OpenAI embedding model.
-
 pub async fn embed_file<T: AsRef<std::path::Path>, F>(
     file_name: T,
     embeder: &Embedder,
@@ -172,7 +171,6 @@ where
 ///     }
 /// };
 /// ```
-
 pub async fn embed_webpage<F>(
     url: String,
     embeder: &Embedder,
@@ -208,6 +206,63 @@ where
     }
 }
 
+/// Embeds an HTML document using the specified embedding model.
+///
+/// # Arguments
+///
+/// * `file_name` - The path of the HTML document to embed.
+/// * `origin` - The original URL of the document. If specified, links can be resolved and metadata
+/// points to the site.
+/// * `embedder` - The embedding model to use. Supported options are "OpenAI", "Jina", and "Bert".
+///
+/// # Returns
+///
+/// The embeddings of the HTML document.
+///
+/// # Errors
+///
+/// Returns an error if the specified embedding model is invalid.
+///
+/// # Example
+///
+/// ```
+/// embed_html(
+///     "test_files/test.html",
+///     "https://example.com/",
+///     &Embedder::Text(TextEmbedder::Jina(JinaEmbedder::default())),
+///     Some(&config),
+///     None,
+/// )
+/// ```
+pub async fn embed_html(
+    file_name: impl AsRef<std::path::Path>,
+    origin: Option<impl Into<String>>,
+    embedder: &Embedder,
+    config: Option<&TextEmbedConfig>,
+    // Callback function
+    adapter: Option<impl FnOnce(Vec<EmbedData>)>,
+) -> Result<Option<Vec<EmbedData>>> {
+    let html_processor = file_processor::html_processor::HtmlProcessor::new();
+    let html = html_processor.process_html_file(file_name.as_ref(), origin)?;
+
+    let binding = TextEmbedConfig::default();
+    let config = config.unwrap_or(&binding);
+    let chunk_size = config.chunk_size.unwrap_or(256);
+    let batch_size = config.batch_size;
+
+    let embeddings = html
+        .embed_webpage(embedder, chunk_size, batch_size)
+        .await?;
+
+    // Send embeddings to vector database
+    if let Some(adapter) = adapter {
+        adapter(embeddings);
+        Ok(None)
+    } else {
+        Ok(Some(embeddings))
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 async fn emb_text<T: AsRef<std::path::Path>, F>(
     file: T,


Test patch:
diff --git a/test_files/test.html b/test_files/test.html
new file mode 100644
index 0000000..65d39ef
--- /dev/null
+++ b/test_files/test.html
@@ -0,0 +1,10 @@
+<!DOCTYPE html>
+<html>
+<body>
+
+<h1>My First Heading</h1>
+
+<p>My first paragraph.</p>
+
+</body>
+</html>
\ No newline at end of file