autoagents_core/
document.rs1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4use crate::embeddings::{Embed, EmbedError, TextEmbedder};
5
6#[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}