lc_shared/
document_types.rs

1// lc-shared/src/document_types.rs
2//! Document types shared across crates.
3//!
4//! These types are needed by both `lc-vector-stores` and `lc-rag`,
5//! so they live here to break the circular dependency.
6
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10
11/// Document structure.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct Document {
14    /// Document content.
15    pub content: String,
16
17    /// Document metadata.
18    pub metadata: HashMap<String, Value>,
19
20    /// Document ID (optional).
21    pub id: Option<String>,
22}
23
24impl Document {
25    /// Creates a new document.
26    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    /// Adds metadata.
35    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    /// Sets ID.
41    pub fn with_id(mut self, id: impl Into<String>) -> Self {
42        self.id = Some(id.into());
43        self
44    }
45
46    /// Returns page content (alias).
47    pub fn page_content(&self) -> &str {
48        &self.content
49    }
50}
51
52/// Vector document with embedding.
53#[derive(Debug, Clone, Serialize, Deserialize)]
54pub struct VectorDocument {
55    /// Document.
56    pub document: Document,
57
58    /// Embedding vector.
59    pub embedding: Vec<f32>,
60}
61
62/// Search result.
63#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct SearchResult {
65    /// Document.
66    pub document: Document,
67
68    /// Similarity score.
69    pub score: f32,
70}
71
72/// Chunk document (split document fragment).
73#[derive(Debug, Clone, Serialize, Deserialize)]
74pub struct ChunkDocument {
75    /// Chunk ID
76    pub chunk_id: String,
77
78    /// Original document ID (Parent ID)
79    pub parent_id: String,
80
81    /// Chunk content
82    pub content: String,
83
84    /// Chunk sequence number
85    pub segment: usize,
86
87    /// Chunk metadata
88    pub metadata: HashMap<String, Value>,
89}
90
91impl ChunkDocument {
92    /// Create a new chunk document
93    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    /// Add metadata
109    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    /// 整体替换 metadata(用于 chunk 继承父文档元数据等场景)。
115    pub fn with_metadata_map(mut self, metadata: HashMap<String, Value>) -> Self {
116        self.metadata = metadata;
117        self
118    }
119
120    /// Convert to Document
121    pub fn to_document(&self) -> Document {
122        Document {
123            content: self.content.clone(),
124            metadata: self.metadata.clone(),
125            id: Some(self.chunk_id.clone()),
126        }
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[test]
135    fn test_document_creation() {
136        let doc = Document::new("Hello, world!")
137            .with_metadata("source", "test")
138            .with_id("doc-1");
139
140        assert_eq!(doc.content, "Hello, world!");
141        assert_eq!(
142            doc.metadata.get("source").and_then(|v| v.as_str()),
143            Some("test")
144        );
145        assert_eq!(doc.id, Some("doc-1".to_string()));
146    }
147
148    #[test]
149    fn test_document_page_content() {
150        let doc = Document::new("Test content");
151        assert_eq!(doc.page_content(), "Test content");
152    }
153
154    #[test]
155    fn test_chunk_document_to_document() {
156        let chunk = ChunkDocument::new("c1".to_string(), "p1".to_string(), "hello".to_string(), 0);
157        let doc = chunk.to_document();
158        assert_eq!(doc.content, "hello");
159        assert_eq!(doc.id, Some("c1".to_string()));
160    }
161}
lc_shared/document_types.rs

lc_shared/
document_types.rs
