lc_vector_stores/
memory.rs

1// lc-vector-stores/src/memory.rs
2//! 内存向量存储
3//!
4//! 将文档和向量存储在内存中，适用于小规模数据和测试。
5
6use crate::{
7    cosine_similarity, Document, SearchResult, VectorDocument, VectorStore, VectorStoreError,
8};
9use async_trait::async_trait;
10use std::collections::HashMap;
11use std::sync::Arc;
12use tokio::sync::RwLock;
13use uuid::Uuid;
14
15/// 内存向量存储
16pub struct InMemoryVectorStore {
17    /// 文档存储
18    documents: Arc<RwLock<HashMap<String, VectorDocument>>>,
19}
20
21impl InMemoryVectorStore {
22    /// 创建新的内存向量存储
23    pub fn new() -> Self {
24        Self {
25            documents: Arc::new(RwLock::new(HashMap::new())),
26        }
27    }
28}
29
30impl Default for InMemoryVectorStore {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36#[async_trait]
37impl VectorStore for InMemoryVectorStore {
38    async fn add_documents(
39        &self,
40        documents: Vec<Document>,
41        embeddings: Vec<Vec<f32>>,
42    ) -> Result<Vec<String>, VectorStoreError> {
43        if documents.len() != embeddings.len() {
44            return Err(VectorStoreError::StorageError(
45                "文档数量和嵌入向量数量不匹配".to_string(),
46            ));
47        }
48
49        let mut store = self.documents.write().await;
50        let mut ids = Vec::new();
51
52        for (doc, embedding) in documents.into_iter().zip(embeddings.into_iter()) {
53            let id = doc.id.clone().unwrap_or_else(|| Uuid::new_v4().to_string());
54
55            let vector_doc = VectorDocument {
56                document: Document {
57                    id: Some(id.clone()),
58                    content: doc.content,
59                    metadata: doc.metadata,
60                },
61                embedding,
62            };
63
64            store.insert(id.clone(), vector_doc);
65            ids.push(id);
66        }
67
68        Ok(ids)
69    }
70
71    async fn similarity_search(
72        &self,
73        query_embedding: &[f32],
74        k: usize,
75    ) -> Result<Vec<SearchResult>, VectorStoreError> {
76        let store = self.documents.read().await;
77
78        // 计算所有文档的相似度并过滤负分 (H27)
79        let mut results: Vec<SearchResult> = store
80            .values()
81            .filter_map(|vd| {
82                let score = cosine_similarity(query_embedding, &vd.embedding).unwrap_or(0.0);
83                if score > 0.0 {
84                    Some(SearchResult {
85                        document: vd.document.clone(),
86                        score,
87                    })
88                } else {
89                    None
90                }
91            })
92            .collect();
93
94        // 按相似度降序排序
95        results.sort_by(|a, b| {
96            b.score
97                .partial_cmp(&a.score)
98                .unwrap_or(std::cmp::Ordering::Equal)
99        });
100
101        // 返回前 k 个结果
102        Ok(results.into_iter().take(k).collect())
103    }
104
105    async fn get_document(&self, id: &str) -> Result<Option<Document>, VectorStoreError> {
106        let store = self.documents.read().await;
107        Ok(store.get(id).map(|vd| vd.document.clone()))
108    }
109
110    async fn get_embedding(&self, id: &str) -> Result<Option<Vec<f32>>, VectorStoreError> {
111        let store = self.documents.read().await;
112        Ok(store.get(id).map(|vd| vd.embedding.clone()))
113    }
114
115    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
116        let mut store = self.documents.write().await;
117        store.remove(id);
118        Ok(())
119    }
120
121    async fn count(&self) -> usize {
122        let store = self.documents.read().await;
123        store.len()
124    }
125
126    async fn clear(&self) -> Result<(), VectorStoreError> {
127        let mut store = self.documents.write().await;
128        store.clear();
129        Ok(())
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136
137    #[tokio::test]
138    async fn test_add_and_search() {
139        let store = InMemoryVectorStore::new();
140
141        // 添加文档
142        let docs = vec![
143            Document::new("Rust is a systems programming language"),
144            Document::new("Python is a scripting language"),
145            Document::new("JavaScript is used for web development"),
146        ];
147
148        // 创建简单的模拟嵌入向量
149        let embeddings = vec![
150            vec![1.0, 0.0, 0.0], // Rust 相关
151            vec![0.0, 1.0, 0.0], // Python 相关
152            vec![0.0, 0.0, 1.0], // JavaScript 相关
153        ];
154
155        let ids = store.add_documents(docs, embeddings).await.unwrap();
156        assert_eq!(ids.len(), 3);
157        assert_eq!(store.count().await, 3);
158
159        // 搜索相似文档
160        let query = vec![0.9, 0.1, 0.0]; // 更接近 Rust
161        let results = store.similarity_search(&query, 2).await.unwrap();
162
163        assert_eq!(results.len(), 2);
164        assert!(results[0].document.content.contains("Rust"));
165        assert!(results[0].score > results[1].score);
166    }
167
168    #[tokio::test]
169    async fn test_get_and_delete() {
170        let store = InMemoryVectorStore::new();
171
172        let doc = Document::new("Test document").with_id("test-id");
173        let embeddings = vec![vec![1.0, 0.0, 0.0]];
174
175        store.add_documents(vec![doc], embeddings).await.unwrap();
176
177        // 获取文档
178        let retrieved = store.get_document("test-id").await.unwrap();
179        assert!(retrieved.is_some());
180        assert_eq!(retrieved.unwrap().content, "Test document");
181
182        // 删除文档
183        store.delete_document("test-id").await.unwrap();
184        assert_eq!(store.count().await, 0);
185
186        // 再次获取应该返回 None
187        let retrieved = store.get_document("test-id").await.unwrap();
188        assert!(retrieved.is_none());
189    }
190
191    #[tokio::test]
192    async fn test_clear() {
193        let store = InMemoryVectorStore::new();
194
195        let docs = vec![Document::new("Doc 1"), Document::new("Doc 2")];
196        let embeddings = vec![vec![1.0, 0.0], vec![0.0, 1.0]];
197
198        store.add_documents(docs, embeddings).await.unwrap();
199        assert_eq!(store.count().await, 2);
200
201        store.clear().await.unwrap();
202        assert_eq!(store.count().await, 0);
203    }
204
205    #[test]
206    fn test_cosine_similarity() {
207        // Identical vectors
208        let a = vec![1.0, 0.0, 0.0];
209        let b = vec![1.0, 0.0, 0.0];
210        assert!((cosine_similarity(&a, &b).unwrap() - 1.0).abs() < 0.0001);
211
212        // Orthogonal vectors
213        let a = vec![1.0, 0.0, 0.0];
214        let b = vec![0.0, 1.0, 0.0];
215        assert!((cosine_similarity(&a, &b).unwrap() - 0.0).abs() < 0.0001);
216    }
217}
lc_vector_stores/memory.rs

lc_vector_stores/
memory.rs
