lc_shared/
document_types.rs1use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct Document {
13 pub content: String,
15
16 pub metadata: HashMap<String, String>,
18
19 pub id: Option<String>,
21}
22
23impl Document {
24 pub fn new(content: impl Into<String>) -> Self {
26 Self {
27 content: content.into(),
28 metadata: HashMap::new(),
29 id: None,
30 }
31 }
32
33 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
35 self.metadata.insert(key.into(), value.into());
36 self
37 }
38
39 pub fn with_id(mut self, id: impl Into<String>) -> Self {
41 self.id = Some(id.into());
42 self
43 }
44
45 pub fn page_content(&self) -> &str {
47 &self.content
48 }
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct VectorDocument {
54 pub document: Document,
56
57 pub embedding: Vec<f32>,
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct SearchResult {
64 pub document: Document,
66
67 pub score: f32,
69}
70
71#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct ChunkDocument {
74 pub chunk_id: String,
76
77 pub parent_id: String,
79
80 pub content: String,
82
83 pub segment: usize,
85
86 pub metadata: HashMap<String, String>,
88}
89
90impl ChunkDocument {
91 pub fn new(chunk_id: String, parent_id: String, content: String, segment: usize) -> Self {
93 Self {
94 chunk_id,
95 parent_id,
96 content,
97 segment,
98 metadata: HashMap::new(),
99 }
100 }
101
102 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
104 self.metadata.insert(key.into(), value.into());
105 self
106 }
107
108 pub fn to_document(&self) -> Document {
110 Document {
111 content: self.content.clone(),
112 metadata: self.metadata.clone(),
113 id: Some(self.chunk_id.clone()),
114 }
115 }
116}
117
118#[cfg(test)]
119mod tests {
120 use super::*;
121
122 #[test]
123 fn test_document_creation() {
124 let doc = Document::new("Hello, world!")
125 .with_metadata("source", "test")
126 .with_id("doc-1");
127
128 assert_eq!(doc.content, "Hello, world!");
129 assert_eq!(doc.metadata.get("source"), Some(&"test".to_string()));
130 assert_eq!(doc.id, Some("doc-1".to_string()));
131 }
132
133 #[test]
134 fn test_document_page_content() {
135 let doc = Document::new("Test content");
136 assert_eq!(doc.page_content(), "Test content");
137 }
138
139 #[test]
140 fn test_chunk_document_to_document() {
141 let chunk = ChunkDocument::new("c1".to_string(), "p1".to_string(), "hello".to_string(), 0);
142 let doc = chunk.to_document();
143 assert_eq!(doc.content, "hello");
144 assert_eq!(doc.id, Some("c1".to_string()));
145 }
146}