Skip to main content

lc_vector_stores/
provider.rs

1// lc-vector-stores/src/provider.rs
2//! 向量存储提供者
3//!
4//! 提供多种向量存储引擎的选择:内存、持久化、Qdrant等
5
6use crate::{VectorStore, VectorStoreError};
7use std::sync::Arc;
8
9/// 向量存储类型枚举
10#[derive(Debug, Clone)]
11pub enum VectorStoreType {
12    /// 内存存储,适用于测试和小型应用
13    InMemory,
14
15    /// 文件持久化存储,适用于个人知识库
16    FileBacked {
17        /// 存储文件路径
18        path: String,
19        /// 向量维度
20        dimension: usize,
21    },
22
23    /// Qdrant 向量数据库,适用于生产环境
24    Qdrant { url: String, collection: String },
25}
26
27/// 向量存储提供者
28pub struct VectorStoreProvider;
29
30impl VectorStoreProvider {
31    /// 创建向量存储实例
32    pub async fn create(
33        store_type: VectorStoreType,
34    ) -> Result<Arc<dyn VectorStore>, VectorStoreError> {
35        match store_type {
36            VectorStoreType::InMemory => {
37                use crate::InMemoryVectorStore;
38                Ok(Arc::new(InMemoryVectorStore::new()))
39            }
40            VectorStoreType::FileBacked { path, dimension } => {
41                use crate::FileVectorStore;
42                let store = FileVectorStore::new(std::path::PathBuf::from(path), dimension)
43                    .await
44                    .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
45                Ok(Arc::new(store))
46            }
47            VectorStoreType::Qdrant { url, collection } => {
48                Self::create_qdrant_store(url, collection).await
49            }
50        }
51    }
52
53    /// 创建 Qdrant 向量存储
54    async fn create_qdrant_store(
55        url: String,
56        collection: String,
57    ) -> Result<Arc<dyn VectorStore>, VectorStoreError> {
58        #[cfg(feature = "qdrant-integration")]
59        {
60            use crate::{QdrantConfig, QdrantVectorStore};
61            let config = QdrantConfig::new(url, collection);
62            let store = QdrantVectorStore::new(config).await?;
63            Ok(Arc::new(store))
64        }
65
66        #[cfg(not(feature = "qdrant-integration"))]
67        {
68            let _ = (url, collection);
69            eprintln!("Warning: Qdrant requested but feature 'qdrant-integration' not enabled. Falling back to InMemory store.");
70            use crate::InMemoryVectorStore;
71            Ok(Arc::new(InMemoryVectorStore::new()))
72        }
73    }
74}
75
76/// 向量存储构建器,提供便利的创建方法
77pub struct VectorStoreBuilder {
78    store_type: VectorStoreType,
79}
80
81impl Default for VectorStoreBuilder {
82    fn default() -> Self {
83        Self::new()
84    }
85}
86
87impl VectorStoreBuilder {
88    pub fn new() -> Self {
89        Self {
90            store_type: VectorStoreType::InMemory,
91        }
92    }
93
94    pub fn in_memory() -> Self {
95        Self {
96            store_type: VectorStoreType::InMemory,
97        }
98    }
99
100    pub fn file_backed(path: impl Into<String>, dimension: usize) -> Self {
101        Self {
102            store_type: VectorStoreType::FileBacked {
103                path: path.into(),
104                dimension,
105            },
106        }
107    }
108
109    pub fn qdrant(url: impl Into<String>, collection: impl Into<String>) -> Self {
110        Self {
111            store_type: VectorStoreType::Qdrant {
112                url: url.into(),
113                collection: collection.into(),
114            },
115        }
116    }
117
118    pub async fn build(self) -> Result<Arc<dyn VectorStore>, VectorStoreError> {
119        VectorStoreProvider::create(self.store_type).await
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126
127    #[tokio::test]
128    async fn test_create_in_memory() {
129        let result = VectorStoreProvider::create(VectorStoreType::InMemory).await;
130        assert!(result.is_ok());
131    }
132
133    #[tokio::test]
134    async fn test_builder_in_memory() {
135        let builder = VectorStoreBuilder::in_memory();
136        let store = builder.build().await;
137        assert!(store.is_ok());
138    }
139
140    #[tokio::test]
141    async fn test_builder_qdrant_fallback() {
142        // 没有 feature 时,应该回退到内存存储
143        let builder = VectorStoreBuilder::qdrant("http://localhost:6334", "test_collection");
144        let store = builder.build().await;
145        assert!(store.is_ok());
146    }
147}