88 lines
2.3 KiB
Rust
88 lines
2.3 KiB
Rust
use std::{collections::HashMap, io::Cursor};
|
|
|
|
use serde_json::Value;
|
|
|
|
/**
|
|
* modified from https://github.com/Abraxas-365/langchain-rust/tree/v4.6.0/src/document_loaders
|
|
*/
|
|
use super::*;
|
|
#[derive(Debug, Clone)]
|
|
pub struct HtmlLoader<R> {
|
|
html: R,
|
|
url: Url,
|
|
}
|
|
|
|
impl HtmlLoader<Cursor<Vec<u8>>> {
|
|
pub fn from_string<S: Into<String>>(input: S, url: Url) -> Self {
|
|
let input = input.into();
|
|
let reader = Cursor::new(input.into_bytes());
|
|
Self::new(reader, url)
|
|
}
|
|
}
|
|
|
|
impl<R: Read> HtmlLoader<R> {
|
|
pub fn new(html: R, url: Url) -> Self {
|
|
Self { html, url }
|
|
}
|
|
}
|
|
|
|
impl<R: Read + Send + Sync + 'static> Loader for HtmlLoader<R> {
|
|
fn load(mut self) -> LoaderResult<Vec<Document>> {
|
|
let cleaned_html = readability::extractor::extract(&mut self.html, &self.url)?;
|
|
let doc =
|
|
Document::new(format!("{}\n{}", cleaned_html.title, cleaned_html.text)).with_metadata(
|
|
HashMap::from([("source".to_string(), Value::from(self.url.as_str()))]),
|
|
);
|
|
|
|
Ok(vec![doc])
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_html_loader() {
|
|
let input = "<p>Hello world!</p>";
|
|
|
|
let html_loader = HtmlLoader::new(
|
|
input.as_bytes(),
|
|
Url::parse("https://example.com/").unwrap(),
|
|
);
|
|
|
|
let documents = html_loader.load().unwrap();
|
|
|
|
let expected = "\nHello world!";
|
|
|
|
assert_eq!(documents.len(), 1);
|
|
assert_eq!(
|
|
documents[0].metadata.get("source").unwrap(),
|
|
&Value::from("https://example.com/")
|
|
);
|
|
assert_eq!(documents[0].page_content, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_html_load_from_path() {
|
|
let buffer = include_bytes!("../../../fixtures/sample.html");
|
|
let html_loader = HtmlLoader::new(
|
|
Cursor::new(buffer),
|
|
Url::parse("https://example.com/").unwrap(),
|
|
);
|
|
|
|
let documents = html_loader.load().unwrap();
|
|
|
|
let expected = "Example Domain\n\n This domain is for use in illustrative examples in \
|
|
documents. You may\n use this domain in literature without prior \
|
|
coordination or asking for\n permission.\n More information...";
|
|
|
|
assert_eq!(documents.len(), 1);
|
|
assert_eq!(
|
|
documents[0].metadata.get("source").unwrap(),
|
|
&Value::from("https://example.com/")
|
|
);
|
|
assert_eq!(documents[0].page_content, expected);
|
|
}
|
|
}
|