Skip to main content

lc_vector_stores/
lib.rs

1// lc-vector-stores/src/lib.rs
2//! Vector store implementations.
3//!
4//! Provides document vector storage and retrieval functionality.
5
6pub mod chromadb;
7pub mod chunked_vector_store;
8pub mod document_store;
9mod file_store;
10pub mod lancedb;
11mod memory;
12pub mod neo4j;
13mod provider;
14
15#[cfg(feature = "mongodb-persistence")]
16mod mongo_document_store;
17
18#[cfg(feature = "qdrant-integration")]
19mod qdrant;
20
21#[cfg(feature = "redis-storage")]
22pub mod redis_store;
23
24#[cfg(feature = "sqlite-storage")]
25pub mod sqlite_store;
26
27#[cfg(feature = "pgvector-storage")]
28pub mod pgvector;
29
30pub mod pinecone;
31
32pub use chunked_vector_store::ChunkedVectorStore;
33pub use document_store::{
34    ChunkedDocumentStore, ChunkedDocumentStoreTrait, DocumentStore, InMemoryChunkedDocumentStore,
35    InMemoryDocumentStore,
36};
37pub use file_store::FileVectorStore;
38pub use lancedb::{LanceDBConfig, LanceDBVectorStore};
39pub use memory::InMemoryVectorStore;
40pub use neo4j::{Neo4jConfig, Neo4jVectorStore};
41pub use pinecone::PineconeStore;
42pub use provider::{VectorStoreBuilder, VectorStoreProvider, VectorStoreType};
43
44#[cfg(feature = "mongodb-persistence")]
45pub use mongo_document_store::{MongoChunkedDocumentStore, MongoStoreConfig};
46
47#[cfg(feature = "qdrant-integration")]
48pub use qdrant::{QdrantConfig, QdrantDistance, QdrantVectorStore};
49
50pub use chromadb::{ChromaDBConfig, ChromaDBVectorStore};
51
52#[cfg(feature = "redis-storage")]
53pub use redis_store::{RedisDocumentStore, RedisStoreConfig};
54
55#[cfg(feature = "sqlite-storage")]
56pub use sqlite_store::{SQLiteDocumentStore, SQLiteStoreConfig};
57
58use async_trait::async_trait;
59use std::error::Error;
60
61// Re-export shared document types from lc-shared
62pub use lc_shared::document::{ChunkDocument, Document, SearchResult, VectorDocument};
63
64// Re-export cosine_similarity from lc-core
65pub use lc_core::math::cosine_similarity;
66
67/// Vector store error types.
68#[derive(Debug)]
69pub enum VectorStoreError {
70    /// Document not found.
71    DocumentNotFound(String),
72
73    /// Embedding error.
74    EmbeddingError(String),
75
76    /// Storage error.
77    StorageError(String),
78
79    /// Connection error (for remote vector databases).
80    ConnectionError(String),
81}
82
83impl std::fmt::Display for VectorStoreError {
84    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85        match self {
86            VectorStoreError::DocumentNotFound(id) => write!(f, "Document not found: {}", id),
87            VectorStoreError::EmbeddingError(msg) => write!(f, "Embedding error: {}", msg),
88            VectorStoreError::StorageError(msg) => write!(f, "Storage error: {}", msg),
89            VectorStoreError::ConnectionError(msg) => write!(f, "Connection error: {}", msg),
90        }
91    }
92}
93
94impl Error for VectorStoreError {}
95
96/// Vector store trait.
97#[async_trait]
98pub trait VectorStore: Send + Sync {
99    /// Adds documents.
100    ///
101    /// # Arguments
102    /// * `documents` - Document list.
103    /// * `embeddings` - Embedding vectors for documents.
104    ///
105    /// # Returns
106    /// Document ID list.
107    async fn add_documents(
108        &self,
109        documents: Vec<Document>,
110        embeddings: Vec<Vec<f32>>,
111    ) -> Result<Vec<String>, VectorStoreError>;
112
113    /// Searches similar documents.
114    ///
115    /// # Arguments
116    /// * `query_embedding` - Query vector.
117    /// * `k` - Number of documents to return.
118    ///
119    /// # Returns
120    /// Similar document list (sorted by similarity descending).
121    async fn similarity_search(
122        &self,
123        query_embedding: &[f32],
124        k: usize,
125    ) -> Result<Vec<SearchResult>, VectorStoreError>;
126
127    /// 返回该向量存储自带的文本嵌入器(若有)。
128    ///
129    /// Q1: 此前 trait 只接收 `query_embedding: &[f32]`,调用方必须自己嵌入查询
130    /// 文本,却没有契约告诉它"该用哪个嵌入器"。有了该 getter,内嵌嵌入器的实现
131    /// 可以直接用 [`similarity_search_text`](Self::similarity_search_text) 传文本;
132    /// 没有的返回 `None`,调用方会收到显式错误而不是静默地用错模型。
133    fn embed_query(&self) -> Option<&dyn Embeddings> {
134        None
135    }
136
137    /// 文本相似度检索:用 [`embed_query`](Self::embed_query) 返回的嵌入器把
138    /// `query` 向量化后再检索。
139    ///
140    /// 未配置嵌入器时返回 [`VectorStoreError::EmbeddingError`],提示改用
141    /// [`similarity_search`](Self::similarity_search) 直接传入查询向量。
142    async fn similarity_search_text(
143        &self,
144        query: &str,
145        k: usize,
146    ) -> Result<Vec<SearchResult>, VectorStoreError> {
147        let Some(embeddings) = self.embed_query() else {
148            return Err(VectorStoreError::EmbeddingError(
149                "该向量存储未配置嵌入器,无法自动向量化查询文本;请改用 similarity_search 直接传入查询向量"
150                    .to_string(),
151            ));
152        };
153        let query_embedding = embeddings
154            .embed_query(query)
155            .await
156            .map_err(|e| VectorStoreError::EmbeddingError(e.to_string()))?;
157        self.similarity_search(&query_embedding, k).await
158    }
159
160    /// 带最低分数阈值的相似度检索。
161    ///
162    /// - `min_score: None` —— 不过滤,返回全库 top-k(即使分数为负)。
163    /// - `min_score: Some(t)` —— 只返回 `score >= t` 的结果,最多 `k` 条。
164    ///
165    /// Q2: 默认实现基于 [`similarity_search`](Self::similarity_search) 的结果做二次过滤
166    /// (对检索期无法直接按阈值过滤的后端是最佳近似)。本地计算相似度的实现应覆盖此
167    /// 方法,以获得"先过滤再取 top-k"的精确语义。
168    async fn similarity_search_with_min_score(
169        &self,
170        query_embedding: &[f32],
171        k: usize,
172        min_score: Option<f32>,
173    ) -> Result<Vec<SearchResult>, VectorStoreError> {
174        let results = self.similarity_search(query_embedding, k).await?;
175        match min_score {
176            Some(threshold) => Ok(results
177                .into_iter()
178                .filter(|r| r.score >= threshold)
179                .collect()),
180            None => Ok(results),
181        }
182    }
183
184    /// Gets document by ID.
185    async fn get_document(&self, id: &str) -> Result<Option<Document>, VectorStoreError>;
186
187    /// Gets document embedding by ID.
188    async fn get_embedding(&self, id: &str) -> Result<Option<Vec<f32>>, VectorStoreError>;
189
190    /// Deletes document.
191    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError>;
192
193    /// Returns document count.
194    async fn count(&self) -> usize;
195
196    /// Clears store.
197    async fn clear(&self) -> Result<(), VectorStoreError>;
198}
199
200/// Embedding model trait — re-exported from lc-embeddings.
201///
202/// Used by vector store implementations that need to embed documents on the fly
203/// (e.g., Pinecone's `upsert` method).
204pub use lc_embeddings::{EmbeddingError, Embeddings};
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209
210    #[test]
211    fn test_document_creation() {
212        let doc = Document::new("Hello, world!")
213            .with_metadata("source", "test")
214            .with_id("doc-1");
215
216        assert_eq!(doc.content, "Hello, world!");
217        assert_eq!(doc.metadata.get("source"), Some(&"test".to_string()));
218        assert_eq!(doc.id, Some("doc-1".to_string()));
219    }
220
221    #[test]
222    fn test_document_page_content() {
223        let doc = Document::new("Test content");
224        assert_eq!(doc.page_content(), "Test content");
225    }
226}