1use std::collections::HashMap;
4use std::time::Duration;
5
6#[cfg(feature = "serde")]
7use serde::{Deserialize, Serialize};
8
9use crate::types::{MetadataValue, SimilarityMetric};
10
11#[derive(Debug, Clone)]
13#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
14pub enum StorageConfig {
15 Memory {
17 initial_capacity: Option<usize>,
19 max_vectors: Option<usize>,
21 },
22
23 Sqlite {
25 database_path: String,
27 pool_size: Option<usize>,
29 wal_mode: bool,
31 synchronous: SqliteSynchronous,
33 },
34
35 Qdrant {
37 url: String,
39 api_key: Option<String>,
41 collection_name: String,
43 timeout: Option<Duration>,
45 tls: bool,
47 },
48
49 MongoDB {
51 connection_string: String,
53 database: String,
55 collection: String,
57 timeout: Option<Duration>,
59 },
60
61 Custom {
63 backend: String,
65 options: HashMap<String, MetadataValue>,
67 },
68}
69
70#[derive(Debug, Clone, Copy)]
72#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
73pub enum SqliteSynchronous {
74 Off,
75 Normal,
76 Full,
77 Extra,
78}
79
80impl Default for SqliteSynchronous {
81 fn default() -> Self {
82 SqliteSynchronous::Normal
83 }
84}
85
86#[derive(Debug, Clone)]
88#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
89pub enum EmbeddingConfig {
90 OpenAI {
92 api_key: String,
94 model: String,
96 base_url: Option<String>,
98 organization: Option<String>,
100 timeout: Option<Duration>,
102 max_retries: u32,
104 },
105
106 Ollama {
108 url: String,
110 model: String,
112 timeout: Option<Duration>,
114 },
115
116 Local {
118 model_path: String,
120 device: String,
122 options: HashMap<String, MetadataValue>,
124 },
125
126 Custom {
128 provider: String,
130 options: HashMap<String, MetadataValue>,
132 },
133}
134
135#[derive(Debug, Clone)]
137#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
138pub struct IndexCreateConfig {
139 pub name: String,
141 pub dimension: usize,
143 pub metric: SimilarityMetric,
145 pub options: IndexOptions,
147}
148
149#[derive(Debug, Clone)]
151#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
152pub struct IndexOptions {
153 pub approximate: bool,
155 pub num_trees: Option<usize>,
157 pub accuracy: Option<f32>,
159 pub max_vectors: Option<usize>,
161 pub compression: bool,
163 pub custom: HashMap<String, MetadataValue>,
165}
166
167impl Default for IndexOptions {
168 fn default() -> Self {
169 Self {
170 approximate: false,
171 num_trees: None,
172 accuracy: None,
173 max_vectors: None,
174 compression: false,
175 custom: HashMap::new(),
176 }
177 }
178}
179
180#[derive(Debug, Clone)]
182#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
183pub struct SearchConfig {
184 pub timeout: Option<Duration>,
186 pub approximate: bool,
188 pub accuracy: Option<f32>,
190 pub cache_results: bool,
192 pub cache_ttl: Option<Duration>,
194 pub custom: HashMap<String, MetadataValue>,
196}
197
198impl Default for SearchConfig {
199 fn default() -> Self {
200 Self {
201 timeout: Some(Duration::from_secs(30)),
202 approximate: false,
203 accuracy: None,
204 cache_results: false,
205 cache_ttl: None,
206 custom: HashMap::new(),
207 }
208 }
209}
210
211#[derive(Debug, Clone)]
213#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
214pub struct VectorSystemConfig {
215 pub storage: StorageConfig,
217 pub embedding: Option<EmbeddingConfig>,
219 pub search: SearchConfig,
221 pub options: HashMap<String, MetadataValue>,
223}
224
225impl VectorSystemConfig {
226 pub fn memory() -> Self {
228 Self {
229 storage: StorageConfig::Memory {
230 initial_capacity: None,
231 max_vectors: None,
232 },
233 embedding: None,
234 search: SearchConfig::default(),
235 options: HashMap::new(),
236 }
237 }
238
239 pub fn sqlite(database_path: impl Into<String>) -> Self {
241 Self {
242 storage: StorageConfig::Sqlite {
243 database_path: database_path.into(),
244 pool_size: None,
245 wal_mode: true,
246 synchronous: SqliteSynchronous::default(),
247 },
248 embedding: None,
249 search: SearchConfig::default(),
250 options: HashMap::new(),
251 }
252 }
253
254 pub fn qdrant(url: impl Into<String>, collection: impl Into<String>) -> Self {
256 Self {
257 storage: StorageConfig::Qdrant {
258 url: url.into(),
259 api_key: None,
260 collection_name: collection.into(),
261 timeout: None,
262 tls: false,
263 },
264 embedding: None,
265 search: SearchConfig::default(),
266 options: HashMap::new(),
267 }
268 }
269
270 pub fn with_embedding(mut self, embedding: EmbeddingConfig) -> Self {
272 self.embedding = Some(embedding);
273 self
274 }
275
276 pub fn with_search_config(mut self, search: SearchConfig) -> Self {
278 self.search = search;
279 self
280 }
281
282 pub fn with_option(mut self, key: impl Into<String>, value: impl Into<MetadataValue>) -> Self {
284 self.options.insert(key.into(), value.into());
285 self
286 }
287}
288
289pub struct StorageConfigBuilder {
291 config: StorageConfig,
292}
293
294impl StorageConfigBuilder {
295 pub fn memory() -> Self {
297 Self {
298 config: StorageConfig::Memory {
299 initial_capacity: None,
300 max_vectors: None,
301 },
302 }
303 }
304
305 pub fn sqlite(database_path: impl Into<String>) -> Self {
307 Self {
308 config: StorageConfig::Sqlite {
309 database_path: database_path.into(),
310 pool_size: None,
311 wal_mode: true,
312 synchronous: SqliteSynchronous::default(),
313 },
314 }
315 }
316
317 pub fn qdrant(url: impl Into<String>, collection: impl Into<String>) -> Self {
319 Self {
320 config: StorageConfig::Qdrant {
321 url: url.into(),
322 api_key: None,
323 collection_name: collection.into(),
324 timeout: None,
325 tls: false,
326 },
327 }
328 }
329
330 pub fn with_initial_capacity(mut self, capacity: usize) -> Self {
332 if let StorageConfig::Memory { ref mut initial_capacity, .. } = self.config {
333 *initial_capacity = Some(capacity);
334 }
335 self
336 }
337
338 pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
340 if let StorageConfig::Qdrant { api_key: ref mut key, .. } = self.config {
341 *key = Some(api_key.into());
342 }
343 self
344 }
345
346 pub fn with_tls(mut self, enable: bool) -> Self {
348 if let StorageConfig::Qdrant { ref mut tls, .. } = self.config {
349 *tls = enable;
350 }
351 self
352 }
353
354 pub fn build(self) -> StorageConfig {
356 self.config
357 }
358}