Skip to main content

lc_vector_stores/
lib.rs

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