lc_shared/
document_types.rs1use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct Document {
14 pub content: String,
16
17 pub metadata: HashMap<String, Value>,
19
20 pub id: Option<String>,
22}
23
24impl Document {
25 pub fn new(content: impl Into<String>) -> Self {
27 Self {
28 content: content.into(),
29 metadata: HashMap::new(),
30 id: None,
31 }
32 }
33
34 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
36 self.metadata.insert(key.into(), value.into());
37 self
38 }
39
40 pub fn with_id(mut self, id: impl Into<String>) -> Self {
42 self.id = Some(id.into());
43 self
44 }
45
46 pub fn page_content(&self) -> &str {
48 &self.content
49 }
50}
51
52#[derive(Debug, Clone, Serialize, Deserialize)]
54pub struct VectorDocument {
55 pub document: Document,
57
58 pub embedding: Vec<f32>,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct SearchResult {
65 pub document: Document,
67
68 pub score: f32,
70}
71
72#[derive(Debug, Clone, Serialize, Deserialize)]
74pub struct ChunkDocument {
75 pub chunk_id: String,
77
78 pub parent_id: String,
80
81 pub content: String,
83
84 pub segment: usize,
86
87 pub metadata: HashMap<String, Value>,
89}
90
91impl ChunkDocument {
92 pub fn new(
94 chunk_id: impl Into<String>,
95 parent_id: impl Into<String>,
96 content: impl Into<String>,
97 segment: usize,
98 ) -> Self {
99 Self {
100 chunk_id: chunk_id.into(),
101 parent_id: parent_id.into(),
102 content: content.into(),
103 segment,
104 metadata: HashMap::new(),
105 }
106 }
107
108 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
110 self.metadata.insert(key.into(), value.into());
111 self
112 }
113
114 pub fn to_document(&self) -> Document {
116 Document {
117 content: self.content.clone(),
118 metadata: self.metadata.clone(),
119 id: Some(self.chunk_id.clone()),
120 }
121 }
122}
123
124#[cfg(test)]
125mod tests {
126 use super::*;
127
128 #[test]
129 fn test_document_creation() {
130 let doc = Document::new("Hello, world!")
131 .with_metadata("source", "test")
132 .with_id("doc-1");
133
134 assert_eq!(doc.content, "Hello, world!");
135 assert_eq!(
136 doc.metadata.get("source").and_then(|v| v.as_str()),
137 Some("test")
138 );
139 assert_eq!(doc.id, Some("doc-1".to_string()));
140 }
141
142 #[test]
143 fn test_document_page_content() {
144 let doc = Document::new("Test content");
145 assert_eq!(doc.page_content(), "Test content");
146 }
147
148 #[test]
149 fn test_chunk_document_to_document() {
150 let chunk = ChunkDocument::new("c1".to_string(), "p1".to_string(), "hello".to_string(), 0);
151 let doc = chunk.to_document();
152 assert_eq!(doc.content, "hello");
153 assert_eq!(doc.id, Some("c1".to_string()));
154 }
155}