Skip to main content

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    /// Convert to Document
115    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}