lc_vector_stores/
file_store.rs

1// lc-vector-stores/src/file_store.rs
2//! File-persistent vector store
3//!
4//! Persists documents and vectors to a local file (JSON serialization), suitable for personal
5//! knowledge bases and offline scenarios. Fills the gap between InMemory (not persistent) and
6//! external databases (too heavy), similar to SQLite vs MySQL.
7//!
8//! # Usage
9//! ```ignore
10//! use lc_vector_stores::FileVectorStore;
11//! use std::path::PathBuf;
12//!
13//! let store = FileVectorStore::new(PathBuf::from("./my_vectors.json"), 1536).unwrap();
14//! // add_documents / similarity_search share the same interface as InMemoryVectorStore
15//! // every insert/delete is automatically persisted to disk
16//! ```
17
18use crate::{
19    cosine_similarity, Document, MetadataFilter, SearchResult, VectorDocument, VectorStore,
20    VectorStoreError,
21};
22use async_trait::async_trait;
23use std::collections::HashMap;
24use std::path::PathBuf;
25use tokio::sync::RwLock;
26use uuid::Uuid;
27
28/// Persistent data structure (serialized as JSON)
29#[derive(serde::Serialize, serde::Deserialize)]
30struct FileStoreData {
31    /// Vector dimension (for validation)
32    dimension: usize,
33    /// Documents + vectors
34    documents: HashMap<String, VectorDocument>,
35}
36
37/// File-persistent vector store
38///
39/// Vectors + metadata are serialized to JSON on disk, loaded at startup, and written back
40/// automatically after inserts/deletes. Uses `RwLock` for concurrency safety: read operations
41/// take a read lock, write operations take a write lock plus persistence.
42pub struct FileVectorStore {
43    /// Storage file path
44    path: PathBuf,
45    /// Vector dimension
46    dimension: usize,
47    /// In-memory data (kept in sync with the file)
48    data: RwLock<FileStoreData>,
49}
50
51impl FileVectorStore {
52    /// Creates or loads a file vector store
53    ///
54    /// If the file already exists, loads its data; otherwise creates an empty store.
55    ///
56    /// # Arguments
57    /// * `path` - storage file path (a `.json` extension is recommended)
58    /// * `dimension` - vector dimension (for validation; when a file exists the file wins)
59    pub async fn new(path: PathBuf, dimension: usize) -> Result<Self, VectorStoreError> {
60        let data = if path.exists() {
61            let content = tokio::fs::read_to_string(&path).await.map_err(|e| {
62                VectorStoreError::StorageError(format!("failed to read file: {}", e))
63            })?;
64            serde_json::from_str::<FileStoreData>(&content).map_err(|e| {
65                VectorStoreError::StorageError(format!("failed to parse file: {}", e))
66            })?
67        } else {
68            // ensure the parent directory exists
69            if let Some(parent) = path.parent() {
70                if !parent.as_os_str().is_empty() {
71                    tokio::fs::create_dir_all(parent).await.map_err(|e| {
72                        VectorStoreError::StorageError(format!("failed to create directory: {}", e))
73                    })?;
74                }
75            }
76            let data = FileStoreData {
77                dimension,
78                documents: HashMap::new(),
79            };
80            // persist an empty file on first creation so path.exists() is true
81            Self::persist(&data, &path).await?;
82            data
83        };
84
85        Ok(Self {
86            path,
87            dimension: data.dimension,
88            data: RwLock::new(data),
89        })
90    }
91
92    /// Persists the current data to disk
93    async fn persist(data: &FileStoreData, path: &PathBuf) -> Result<(), VectorStoreError> {
94        let json = serde_json::to_string(data)
95            .map_err(|e| VectorStoreError::StorageError(format!("failed to serialize: {}", e)))?;
96        // write to a temporary file first, then rename, to avoid corruption from a mid-write power loss
97        let tmp_path = path.with_extension("json.tmp");
98        tokio::fs::write(&tmp_path, &json).await.map_err(|e| {
99            VectorStoreError::StorageError(format!("failed to write temporary file: {}", e))
100        })?;
101        tokio::fs::rename(&tmp_path, path)
102            .await
103            .map_err(|e| VectorStoreError::StorageError(format!("failed to rename file: {}", e)))?;
104        Ok(())
105    }
106
107    /// Returns the vector dimension
108    pub fn dimension(&self) -> usize {
109        self.dimension
110    }
111
112    /// Returns the storage file path
113    pub fn path(&self) -> &PathBuf {
114        &self.path
115    }
116}
117
118#[async_trait]
119impl VectorStore for FileVectorStore {
120    async fn add_documents(
121        &self,
122        documents: Vec<Document>,
123        embeddings: Vec<Vec<f32>>,
124    ) -> Result<Vec<String>, VectorStoreError> {
125        if documents.len() != embeddings.len() {
126            return Err(VectorStoreError::StorageError(
127                "document count and embedding count mismatch".to_string(),
128            ));
129        }
130
131        let mut data = self.data.write().await;
132        let mut ids = Vec::new();
133
134        for (doc, embedding) in documents.into_iter().zip(embeddings.into_iter()) {
135            // validate the dimension
136            if !embedding.is_empty() && embedding.len() != data.dimension {
137                return Err(VectorStoreError::StorageError(format!(
138                    "embedding dimension {} does not match storage dimension {}",
139                    embedding.len(),
140                    data.dimension
141                )));
142            }
143
144            let id = doc.id.clone().unwrap_or_else(|| Uuid::new_v4().to_string());
145            let vector_doc = VectorDocument {
146                document: Document {
147                    id: Some(id.clone()),
148                    content: doc.content,
149                    metadata: doc.metadata,
150                },
151                embedding,
152            };
153            data.documents.insert(id.clone(), vector_doc);
154            ids.push(id);
155        }
156
157        Self::persist(&data, &self.path).await?;
158        Ok(ids)
159    }
160
161    async fn similarity_search(
162        &self,
163        query_embedding: &[f32],
164        k: usize,
165    ) -> Result<Vec<SearchResult>, VectorStoreError> {
166        let data = self.data.read().await;
167
168        let mut results: Vec<SearchResult> = data
169            .documents
170            .values()
171            .map(|vd| {
172                let score = cosine_similarity(query_embedding, &vd.embedding).unwrap_or(0.0);
173                SearchResult {
174                    document: vd.document.clone(),
175                    score,
176                }
177            })
178            .collect();
179
180        results.sort_by(|a, b| {
181            b.score
182                .partial_cmp(&a.score)
183                .unwrap_or(std::cmp::Ordering::Equal)
184        });
185
186        Ok(results.into_iter().take(k).collect())
187    }
188
189    /// S3: file-store metadata filtering — the same "filter first, then compute similarity" semantics as the in-memory store.
190    async fn similarity_search_with_filter(
191        &self,
192        query_embedding: &[f32],
193        k: usize,
194        filter: Option<&MetadataFilter>,
195    ) -> Result<Vec<SearchResult>, VectorStoreError> {
196        let data = self.data.read().await;
197
198        let mut results: Vec<SearchResult> = data
199            .documents
200            .values()
201            .filter(|vd| filter.is_none_or(|f| f.matches(&vd.document.metadata)))
202            .map(|vd| {
203                let score = cosine_similarity(query_embedding, &vd.embedding).unwrap_or(0.0);
204                SearchResult {
205                    document: vd.document.clone(),
206                    score,
207                }
208            })
209            .collect();
210
211        results.sort_by(|a, b| {
212            b.score
213                .partial_cmp(&a.score)
214                .unwrap_or(std::cmp::Ordering::Equal)
215        });
216
217        Ok(results.into_iter().take(k).collect())
218    }
219
220    async fn get_document(&self, id: &str) -> Result<Option<Document>, VectorStoreError> {
221        let data = self.data.read().await;
222        Ok(data.documents.get(id).map(|vd| vd.document.clone()))
223    }
224
225    async fn get_embedding(&self, id: &str) -> Result<Option<Vec<f32>>, VectorStoreError> {
226        let data = self.data.read().await;
227        Ok(data.documents.get(id).map(|vd| vd.embedding.clone()))
228    }
229
230    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
231        let mut data = self.data.write().await;
232        data.documents
233            .remove(id)
234            .ok_or_else(|| VectorStoreError::DocumentNotFound(id.to_string()))?;
235        Self::persist(&data, &self.path).await?;
236        Ok(())
237    }
238
239    async fn count(&self) -> usize {
240        let data = self.data.read().await;
241        data.documents.len()
242    }
243
244    async fn clear(&self) -> Result<(), VectorStoreError> {
245        let mut data = self.data.write().await;
246        data.documents.clear();
247        Self::persist(&data, &self.path).await?;
248        Ok(())
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255    use tempfile::TempDir;
256
257    fn test_store_path(dir: &TempDir) -> PathBuf {
258        dir.path().join("test_vectors.json")
259    }
260
261    #[tokio::test]
262    async fn test_new_creates_empty_store() {
263        let dir = TempDir::new().unwrap();
264        let path = test_store_path(&dir);
265        let store = FileVectorStore::new(path.clone(), 3).await.unwrap();
266        assert_eq!(store.count().await, 0);
267        assert_eq!(store.dimension(), 3);
268        assert!(path.exists());
269    }
270
271    #[tokio::test]
272    async fn test_add_and_search() {
273        let dir = TempDir::new().unwrap();
274        let path = test_store_path(&dir);
275        let store = FileVectorStore::new(path, 3).await.unwrap();
276
277        let docs = vec![
278            Document::new("Rust is a systems programming language"),
279            Document::new("Python is a scripting language"),
280            Document::new("JavaScript is used for web development"),
281        ];
282        let embeddings = vec![
283            vec![1.0, 0.0, 0.0],
284            vec![0.0, 1.0, 0.0],
285            vec![0.0, 0.0, 1.0],
286        ];
287
288        let ids = store.add_documents(docs, embeddings).await.unwrap();
289        assert_eq!(ids.len(), 3);
290        assert_eq!(store.count().await, 3);
291
292        let query = vec![0.9, 0.1, 0.0];
293        let results = store.similarity_search(&query, 2).await.unwrap();
294        assert_eq!(results.len(), 2);
295        assert!(results[0].document.content.contains("Rust"));
296        assert!(results[0].score > results[1].score);
297    }
298
299    #[tokio::test]
300    async fn test_persistence_across_instances() {
301        let dir = TempDir::new().unwrap();
302        let path = test_store_path(&dir);
303
304        // first instance: write
305        {
306            let store = FileVectorStore::new(path.clone(), 3).await.unwrap();
307            let doc = Document::new("persistent doc").with_id("p1");
308            store
309                .add_documents(vec![doc], vec![vec![1.0, 0.0, 0.0]])
310                .await
311                .unwrap();
312        }
313
314        // second instance: load and verify
315        {
316            let store = FileVectorStore::new(path.clone(), 3).await.unwrap();
317            assert_eq!(store.count().await, 1);
318            let doc = store.get_document("p1").await.unwrap().unwrap();
319            assert_eq!(doc.content, "persistent doc");
320        }
321    }
322
323    #[tokio::test]
324    async fn test_delete_persists() {
325        let dir = TempDir::new().unwrap();
326        let path = test_store_path(&dir);
327
328        {
329            let store = FileVectorStore::new(path.clone(), 3).await.unwrap();
330            let doc = Document::new("to delete").with_id("d1");
331            store
332                .add_documents(vec![doc], vec![vec![1.0, 0.0, 0.0]])
333                .await
334                .unwrap();
335            store.delete_document("d1").await.unwrap();
336        }
337
338        let store = FileVectorStore::new(path.clone(), 3).await.unwrap();
339        assert_eq!(store.count().await, 0);
340    }
341
342    #[tokio::test]
343    async fn test_clear_persists() {
344        let dir = TempDir::new().unwrap();
345        let path = test_store_path(&dir);
346
347        {
348            let store = FileVectorStore::new(path.clone(), 3).await.unwrap();
349            let docs = vec![Document::new("a"), Document::new("b")];
350            let embeddings = vec![vec![1.0, 0.0, 0.0], vec![0.0, 1.0, 0.0]];
351            store.add_documents(docs, embeddings).await.unwrap();
352            store.clear().await.unwrap();
353        }
354
355        let store = FileVectorStore::new(path.clone(), 3).await.unwrap();
356        assert_eq!(store.count().await, 0);
357    }
358
359    #[tokio::test]
360    async fn test_dimension_mismatch() {
361        let dir = TempDir::new().unwrap();
362        let path = test_store_path(&dir);
363        let store = FileVectorStore::new(path, 3).await.unwrap();
364
365        let doc = Document::new("wrong dim");
366        let wrong_embedding = vec![1.0, 0.0]; // dimension 2, storage dimension 3
367        let result = store.add_documents(vec![doc], vec![wrong_embedding]).await;
368        assert!(result.is_err());
369    }
370
371    #[tokio::test]
372    async fn test_get_embedding() {
373        let dir = TempDir::new().unwrap();
374        let path = test_store_path(&dir);
375        let store = FileVectorStore::new(path, 3).await.unwrap();
376
377        let doc = Document::new("embed test").with_id("e1");
378        store
379            .add_documents(vec![doc], vec![vec![0.5, 0.5, 0.0]])
380            .await
381            .unwrap();
382
383        let emb = store.get_embedding("e1").await.unwrap().unwrap();
384        assert_eq!(emb, vec![0.5, 0.5, 0.0]);
385    }
386
387    #[tokio::test]
388    async fn test_delete_nonexistent() {
389        let dir = TempDir::new().unwrap();
390        let path = test_store_path(&dir);
391        let store = FileVectorStore::new(path, 3).await.unwrap();
392
393        let result = store.delete_document("no-such-id").await;
394        assert!(result.is_err());
395    }
396
397    #[tokio::test]
398    async fn test_cosine_similarity() {
399        let a = vec![1.0, 0.0, 0.0];
400        let b = vec![1.0, 0.0, 0.0];
401        assert!((cosine_similarity(&a, &b).unwrap() - 1.0).abs() < 0.0001);
402
403        let a = vec![1.0, 0.0, 0.0];
404        let b = vec![0.0, 1.0, 0.0];
405        assert!((cosine_similarity(&a, &b).unwrap() - 0.0).abs() < 0.0001);
406    }
407
408    /// S3: file-store metadata filtering — single condition + AND combination.
409    #[tokio::test]
410    async fn test_metadata_filter() {
411        use crate::FilterOp;
412
413        let dir = TempDir::new().unwrap();
414        let path = test_store_path(&dir);
415        let store = FileVectorStore::new(path, 3).await.unwrap();
416
417        store
418            .add_documents(
419                vec![
420                    Document::new("rust doc")
421                        .with_metadata("lang", "rust")
422                        .with_metadata("year", 2024),
423                    Document::new("python doc")
424                        .with_metadata("lang", "python")
425                        .with_metadata("year", 2023),
426                ],
427                vec![vec![1.0, 0.0, 0.0], vec![0.0, 1.0, 0.0]],
428            )
429            .await
430            .unwrap();
431
432        let query = vec![1.0, 0.0, 0.0];
433
434        let eq = MetadataFilter::field("lang", FilterOp::Eq, "rust");
435        let r = store
436            .similarity_search_with_filter(&query, 5, Some(&eq))
437            .await
438            .unwrap();
439        assert_eq!(r.len(), 1);
440        assert_eq!(r[0].document.content, "rust doc");
441
442        let and = MetadataFilter::and(vec![
443            MetadataFilter::field("lang", FilterOp::Eq, "rust"),
444            MetadataFilter::field("year", FilterOp::Gt, 2020),
445        ]);
446        let r = store
447            .similarity_search_with_filter(&query, 5, Some(&and))
448            .await
449            .unwrap();
450        assert_eq!(r.len(), 1);
451
452        // filter: None behaves identically to similarity_search
453        let none = store
454            .similarity_search_with_filter(&query, 5, None)
455            .await
456            .unwrap();
457        let base = store.similarity_search(&query, 5).await.unwrap();
458        assert_eq!(none.len(), base.len());
459    }
460}
lc_vector_stores/file_store.rs

lc_vector_stores/
file_store.rs
