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}