Skip to main content

autoagents_core/
document.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4use crate::embeddings::{Embed, EmbedError, TextEmbedder};
5
6/// Represents a piece of content along with optional metadata.
7/// Mirrors the LangChain document abstraction (`page_content` + `metadata`).
8#[derive(Debug, Clone, Serialize, Deserialize, Default)]
9pub struct Document {
10    pub page_content: String,
11    #[serde(default)]
12    pub metadata: Value,
13}
14
15impl Document {
16    pub fn new(page_content: impl Into<String>) -> Self {
17        Self {
18            page_content: page_content.into(),
19            metadata: Value::Object(Default::default()),
20        }
21    }
22
23    pub fn with_metadata(page_content: impl Into<String>, metadata: Value) -> Self {
24        Self {
25            page_content: page_content.into(),
26            metadata,
27        }
28    }
29}
30
31impl Embed for Document {
32    fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError> {
33        embedder.embed(self.page_content.clone());
34        Ok(())
35    }
36}
37
38#[cfg(test)]
39mod tests {
40    use super::*;
41    use serde_json::json;
42
43    #[test]
44    fn document_new_initializes_empty_metadata() {
45        let document = Document::new("hello");
46
47        assert_eq!(document.page_content, "hello");
48        assert_eq!(document.metadata, json!({}));
49    }
50
51    #[test]
52    fn document_with_metadata_preserves_metadata() {
53        let metadata = json!({
54            "source": "unit-test",
55            "page": 2
56        });
57        let document = Document::with_metadata("hello", metadata.clone());
58
59        assert_eq!(document.page_content, "hello");
60        assert_eq!(document.metadata, metadata);
61    }
62
63    #[test]
64    fn document_embed_pushes_page_content_into_embedder() {
65        let document = Document::new("embedded text");
66        let mut embedder = TextEmbedder::new();
67
68        document.embed(&mut embedder).expect("document embeds");
69
70        assert_eq!(embedder.parts(), &["embedded text"]);
71    }
72}