Skip to main content

kalosm_language/context/
document.rs

1use std::{convert::Infallible, future::Future};
2use url::Url;
3pub use whatlang::Lang;
4
5/// A document is a piece of text with a title.
6#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
7pub struct Document {
8    title: String,
9    body: String,
10    summary: Option<String>,
11    created_at: Option<chrono::DateTime<chrono::Utc>>,
12    updated_at: Option<chrono::DateTime<chrono::Utc>>,
13}
14
15impl Document {
16    /// Create a new document from a source.
17    pub async fn new<T: IntoDocument>(source: T) -> Result<Self, T::Error> {
18        source.into_document().await
19    }
20
21    /// Get the language of the document.
22    pub fn language(&self) -> Option<whatlang::Lang> {
23        whatlang::detect_lang(&self.body)
24    }
25
26    /// Create a new document from the raw parts.
27    pub fn from_parts(title: impl Into<String>, body: impl Into<String>) -> Self {
28        Self {
29            title: title.into(),
30            body: body.into(),
31            summary: None,
32            created_at: None,
33            updated_at: None,
34        }
35    }
36
37    /// Set the summary of the document.
38    pub fn set_summary(&mut self, summary: impl Into<String>) {
39        self.summary = Some(summary.into());
40    }
41
42    /// Set the created at time of the document.
43    pub fn set_created_at(&mut self, created_at: chrono::DateTime<chrono::Utc>) {
44        self.created_at = Some(created_at);
45    }
46
47    /// Set the updated at time of the document.
48    pub fn set_updated_at(&mut self, updated_at: chrono::DateTime<chrono::Utc>) {
49        self.updated_at = Some(updated_at);
50    }
51
52    /// Get the title of the document.
53    pub fn title(&self) -> &str {
54        &self.title
55    }
56
57    /// Get the body of the document.
58    pub fn body(&self) -> &str {
59        &self.body
60    }
61}
62
63impl From<String> for Document {
64    fn from(value: String) -> Self {
65        Self::from_parts("", value)
66    }
67}
68
69impl From<&str> for Document {
70    fn from(value: &str) -> Self {
71        Self::from_parts("", value)
72    }
73}
74
75impl AsRef<Document> for Document {
76    fn as_ref(&self) -> &Document {
77        self
78    }
79}
80
81impl std::fmt::Display for Document {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        write!(f, "{}\n{}", self.title, self.body)
84    }
85}
86
87/// A trait for types that can be converted into a document.
88pub trait IntoDocument {
89    /// The error type that can occur when converting the type into a [`Document`].
90    type Error: Send + Sync + 'static;
91
92    /// Convert the type into a document.
93    fn into_document(self) -> impl Future<Output = Result<Document, Self::Error>> + Send;
94}
95
96impl IntoDocument for String {
97    type Error = Infallible;
98
99    async fn into_document(self) -> Result<Document, Self::Error> {
100        Ok(Document::from_parts("", self))
101    }
102}
103
104impl IntoDocument for &String {
105    type Error = Infallible;
106
107    async fn into_document(self) -> Result<Document, Self::Error> {
108        Ok(Document::from_parts("", self.to_string()))
109    }
110}
111
112impl IntoDocument for &str {
113    type Error = Infallible;
114
115    async fn into_document(self) -> Result<Document, Self::Error> {
116        Ok(Document::from_parts("", self.to_string()))
117    }
118}
119
120impl IntoDocument for Document {
121    type Error = Infallible;
122
123    async fn into_document(self) -> Result<Document, Self::Error> {
124        Ok(self)
125    }
126}
127
128/// A document that can be added to a search index.
129pub trait IntoDocuments {
130    /// The error type that can occur when converting the document into [`Document`]s.
131    type Error: Send + Sync + 'static;
132
133    /// Convert the document into a [`Document`]
134    fn into_documents(self) -> impl Future<Output = Result<Vec<Document>, Self::Error>> + Send;
135}
136
137impl<T: IntoDocument + Send + Sync, I> IntoDocuments for I
138where
139    I: IntoIterator<Item = T> + Send + Sync,
140    <I as IntoIterator>::IntoIter: Send + Sync,
141{
142    type Error = T::Error;
143
144    async fn into_documents(self) -> Result<Vec<Document>, Self::Error> {
145        let mut documents = Vec::new();
146        for document in self {
147            documents.push(document.into_document().await?);
148        }
149        Ok(documents)
150    }
151}
152
153/// An error that can occur when extracting a document from a URL.
154#[derive(Debug, thiserror::Error)]
155pub enum ExtractDocumentError {
156    /// An error occurred when fetching the HTML.
157    #[error("Failed to fetch HTML: {0}")]
158    FetchHtml(#[from] reqwest::Error),
159    /// An error occurred when extracting the article.
160    #[error("Failed to extract article: {0}")]
161    ExtractArticle(#[from] readability::error::Error),
162    /// Failed to parse the URL.
163    #[error("Failed to parse URL: {0}")]
164    ParseUrl(#[from] url::ParseError),
165}
166
167pub(crate) async fn get_article(url: Url) -> Result<Document, ExtractDocumentError> {
168    let html = reqwest::get(url.clone()).await?.text().await?;
169    extract_article(&html)
170}
171
172pub(crate) fn extract_article(html: &str) -> Result<Document, ExtractDocumentError> {
173    let cleaned =
174        readability::extractor::extract(&mut html.as_bytes(), &Url::parse("https://example.com")?)
175            .unwrap();
176    Ok(Document::from_parts(cleaned.title, cleaned.text))
177}
178
179impl IntoDocument for Url {
180    type Error = ExtractDocumentError;
181
182    async fn into_document(self) -> Result<Document, Self::Error> {
183        get_article(self).await
184    }
185}