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 std::collections::HashMap;
9
10/// Document structure.
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct Document {
13    /// Document content.
14    pub content: String,
15
16    /// Document metadata.
17    pub metadata: HashMap<String, String>,
18
19    /// Document ID (optional).
20    pub id: Option<String>,
21}
22
23impl Document {
24    /// Creates a new document.
25    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    /// Adds metadata.
34    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    /// Sets ID.
40    pub fn with_id(mut self, id: impl Into<String>) -> Self {
41        self.id = Some(id.into());
42        self
43    }
44
45    /// Returns page content (alias).
46    pub fn page_content(&self) -> &str {
47        &self.content
48    }
49}
50
51/// Vector document with embedding.
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct VectorDocument {
54    /// Document.
55    pub document: Document,
56
57    /// Embedding vector.
58    pub embedding: Vec<f32>,
59}
60
61/// Search result.
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct SearchResult {
64    /// Document.
65    pub document: Document,
66
67    /// Similarity score.
68    pub score: f32,
69}
70
71/// Chunk document (split document fragment).
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct ChunkDocument {
74    /// Chunk ID
75    pub chunk_id: String,
76
77    /// Original document ID (Parent ID)
78    pub parent_id: String,
79
80    /// Chunk content
81    pub content: String,
82
83    /// Chunk sequence number
84    pub segment: usize,
85
86    /// Chunk metadata
87    pub metadata: HashMap<String, String>,
88}
89
90impl ChunkDocument {
91    /// Create a new chunk document
92    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    /// Add metadata
103    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    /// Convert to Document
109    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}