Skip to main content

fierros_core/
document.rs

1use crate::Metadata;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
5pub struct Document {
6    pub id: String,
7    pub text: String,
8    pub metadata: Metadata,
9}
10
11impl Document {
12    pub fn new(id: impl Into<String>, text: impl Into<String>) -> Self {
13        Self {
14            id: id.into(),
15            text: text.into(),
16            metadata: Metadata::new(),
17        }
18    }
19
20    pub fn with_metadata(mut self, metadata: Metadata) -> Self {
21        self.metadata = metadata;
22        self
23    }
24}
25
26#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
27pub struct Chunk {
28    pub id: String,
29    pub document_id: String,
30    pub text: String,
31    pub metadata: Metadata,
32    pub start_char: usize,
33    pub end_char: usize,
34}
35
36impl Chunk {
37    pub fn new(
38        id: impl Into<String>,
39        document_id: impl Into<String>,
40        text: impl Into<String>,
41        start_char: usize,
42        end_char: usize,
43    ) -> Self {
44        Self {
45            id: id.into(),
46            document_id: document_id.into(),
47            text: text.into(),
48            metadata: Metadata::new(),
49            start_char,
50            end_char,
51        }
52    }
53
54    pub fn with_metadata(mut self, metadata: Metadata) -> Self {
55        self.metadata = metadata;
56        self
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::{Chunk, Document};
63    use crate::Metadata;
64
65    #[test]
66    fn document_constructor_sets_id_and_text() {
67        let doc = Document::new("doc-1", "hello");
68        assert_eq!(doc.id, "doc-1");
69        assert_eq!(doc.text, "hello");
70    }
71
72    #[test]
73    fn chunk_constructor_sets_offsets() {
74        let chunk = Chunk::new("chunk-1", "doc-1", "hello", 0, 5);
75        assert_eq!(chunk.document_id, "doc-1");
76        assert_eq!(chunk.start_char, 0);
77        assert_eq!(chunk.end_char, 5);
78    }
79
80    #[test]
81    fn document_with_metadata_replaces_default_metadata() {
82        let metadata = Metadata::new().with("source", "runbook.md");
83        let doc = Document::new("doc-1", "hello").with_metadata(metadata.clone());
84        assert_eq!(doc.metadata, metadata);
85    }
86
87    #[test]
88    fn chunk_with_metadata_replaces_default_metadata() {
89        let metadata = Metadata::new().with("score", 0.9);
90        let chunk = Chunk::new("chunk-1", "doc-1", "hello", 0, 5).with_metadata(metadata.clone());
91        assert_eq!(chunk.metadata, metadata);
92    }
93}