litellm_rs/storage/vector/
qdrant.rs

1//! Qdrant vector store implementation
2
3use crate::config::models::file_storage::VectorDbConfig;
4use crate::core::http::outbound::default_outbound_client;
5use crate::utils::error::gateway_error::{GatewayError, Result};
6use tracing::{debug, info};
7
8use super::types::{SearchResult, VectorPoint};
9
10/// Qdrant vector store
11#[derive(Debug, Clone)]
12pub struct QdrantStore {
13    url: String,
14    api_key: Option<String>,
15    collection: String,
16    client: reqwest::Client,
17}
18
19impl QdrantStore {
20    /// Create a new Qdrant store
21    pub async fn new(config: &VectorDbConfig) -> Result<Self> {
22        let client = default_outbound_client().clone();
23
24        let store = Self {
25            url: config.url.clone(),
26            api_key: Some(config.api_key.clone()),
27            collection: config.index_name.clone(),
28            client,
29        };
30
31        // Ensure collection exists
32        store.ensure_collection().await?;
33
34        info!("Qdrant vector store initialized");
35        Ok(store)
36    }
37
38    /// Ensure collection exists
39    async fn ensure_collection(&self) -> Result<()> {
40        let url = format!("{}/collections/{}", self.url, self.collection);
41        let mut request = self.client.get(&url);
42
43        if let Some(api_key) = &self.api_key {
44            request = request.header("api-key", api_key);
45        }
46
47        let response = request
48            .send()
49            .await
50            .map_err(|e| GatewayError::Storage(format!("Failed to check collection: {}", e)))?;
51
52        if response.status() == 404 {
53            // Collection doesn't exist, create it
54            self.create_collection().await?;
55        } else if !response.status().is_success() {
56            return Err(GatewayError::Storage(format!(
57                "Failed to check collection: {}",
58                response.status()
59            )));
60        }
61
62        Ok(())
63    }
64
65    /// Create collection
66    async fn create_collection(&self) -> Result<()> {
67        let url = format!("{}/collections/{}", self.url, self.collection);
68        let payload = serde_json::json!({
69            "vectors": {
70                "size": 1536, // Default OpenAI embedding size
71                "distance": "Cosine"
72            }
73        });
74
75        let mut request = self.client.put(&url).json(&payload);
76
77        if let Some(api_key) = &self.api_key {
78            request = request.header("api-key", api_key);
79        }
80
81        let response = request
82            .send()
83            .await
84            .map_err(|e| GatewayError::Storage(format!("Failed to create collection: {}", e)))?;
85
86        if !response.status().is_success() {
87            return Err(GatewayError::Storage(format!(
88                "Failed to create collection: {}",
89                response.status()
90            )));
91        }
92
93        info!("Created Qdrant collection: {}", self.collection);
94        Ok(())
95    }
96
97    /// Store a vector
98    pub async fn store(
99        &self,
100        id: &str,
101        vector: &[f32],
102        metadata: Option<serde_json::Value>,
103    ) -> Result<()> {
104        let url = format!("{}/collections/{}/points", self.url, self.collection);
105        let payload = serde_json::json!({
106            "points": [{
107                "id": id,
108                "vector": vector,
109                "payload": metadata.unwrap_or_default()
110            }]
111        });
112
113        let mut request = self.client.put(&url).json(&payload);
114
115        if let Some(api_key) = &self.api_key {
116            request = request.header("api-key", api_key);
117        }
118
119        let response = request
120            .send()
121            .await
122            .map_err(|e| GatewayError::Storage(format!("Failed to store vector: {}", e)))?;
123
124        if !response.status().is_success() {
125            return Err(GatewayError::Storage(format!(
126                "Failed to store vector: {}",
127                response.status()
128            )));
129        }
130
131        debug!("Stored vector: {}", id);
132        Ok(())
133    }
134
135    /// Search for similar vectors
136    pub async fn search(
137        &self,
138        query_vector: &[f32],
139        limit: usize,
140        threshold: Option<f32>,
141    ) -> Result<Vec<SearchResult>> {
142        let url = format!("{}/collections/{}/points/search", self.url, self.collection);
143        let mut payload = serde_json::json!({
144            "vector": query_vector,
145            "limit": limit,
146            "with_payload": true,
147            "with_vector": false
148        });
149
150        if let Some(threshold) = threshold {
151            payload["score_threshold"] = serde_json::Number::from_f64(threshold as f64)
152                .map(serde_json::Value::Number)
153                .ok_or_else(|| {
154                    GatewayError::Storage(format!(
155                        "Invalid score threshold: {} (NaN/Infinity not allowed)",
156                        threshold
157                    ))
158                })?;
159        }
160
161        let mut request = self.client.post(&url).json(&payload);
162
163        if let Some(api_key) = &self.api_key {
164            request = request.header("api-key", api_key);
165        }
166
167        let response = request
168            .send()
169            .await
170            .map_err(|e| GatewayError::Storage(format!("Failed to search vectors: {}", e)))?;
171
172        if !response.status().is_success() {
173            return Err(GatewayError::Storage(format!(
174                "Failed to search vectors: {}",
175                response.status()
176            )));
177        }
178
179        let result: serde_json::Value = response.json().await.map_err(|e| {
180            GatewayError::Storage(format!("Failed to parse search response: {}", e))
181        })?;
182
183        let mut search_results = Vec::new();
184        if let Some(points) = result["result"].as_array() {
185            for point in points {
186                if let (Some(id), Some(score)) = (point["id"].as_str(), point["score"].as_f64()) {
187                    search_results.push(SearchResult {
188                        id: id.to_string(),
189                        score: score as f32,
190                        metadata: point["payload"].clone().into(),
191                        vector: None,
192                    });
193                }
194            }
195        }
196
197        Ok(search_results)
198    }
199
200    /// Delete a vector
201    pub async fn delete(&self, id: &str) -> Result<()> {
202        let url = format!("{}/collections/{}/points/delete", self.url, self.collection);
203        let payload = serde_json::json!({
204            "points": [id]
205        });
206
207        let mut request = self.client.post(&url).json(&payload);
208
209        if let Some(api_key) = &self.api_key {
210            request = request.header("api-key", api_key);
211        }
212
213        let response = request
214            .send()
215            .await
216            .map_err(|e| GatewayError::Storage(format!("Failed to delete vector: {}", e)))?;
217
218        if !response.status().is_success() {
219            return Err(GatewayError::Storage(format!(
220                "Failed to delete vector: {}",
221                response.status()
222            )));
223        }
224
225        debug!("Deleted vector: {}", id);
226        Ok(())
227    }
228
229    /// Get a vector by ID
230    pub async fn get(&self, id: &str) -> Result<Option<VectorPoint>> {
231        let url = format!("{}/collections/{}/points/{}", self.url, self.collection, id);
232        let mut request = self.client.get(&url);
233
234        if let Some(api_key) = &self.api_key {
235            request = request.header("api-key", api_key);
236        }
237
238        let response = request
239            .send()
240            .await
241            .map_err(|e| GatewayError::Storage(format!("Failed to get vector: {}", e)))?;
242
243        if response.status() == 404 {
244            return Ok(None);
245        }
246
247        if !response.status().is_success() {
248            return Err(GatewayError::Storage(format!(
249                "Failed to get vector: {}",
250                response.status()
251            )));
252        }
253
254        let result: serde_json::Value = response
255            .json()
256            .await
257            .map_err(|e| GatewayError::Storage(format!("Failed to parse get response: {}", e)))?;
258
259        if let Some(point) = result["result"].as_object()
260            && let (Some(id), Some(vector)) = (point["id"].as_str(), point["vector"].as_array())
261        {
262            let vector_data: Vec<f32> = vector
263                .iter()
264                .filter_map(|v| v.as_f64().map(|f| f as f32))
265                .collect();
266
267            return Ok(Some(VectorPoint {
268                id: id.to_string(),
269                vector: vector_data,
270                metadata: point["payload"].clone().into(),
271            }));
272        }
273
274        Ok(None)
275    }
276
277    /// Health check
278    pub async fn health_check(&self) -> Result<()> {
279        let url = format!("{}/", self.url);
280        let mut request = self.client.get(&url);
281
282        if let Some(api_key) = &self.api_key {
283            request = request.header("api-key", api_key);
284        }
285
286        let response = request
287            .send()
288            .await
289            .map_err(|e| GatewayError::Storage(format!("Qdrant health check failed: {}", e)))?;
290
291        if !response.status().is_success() {
292            return Err(GatewayError::Storage(format!(
293                "Qdrant health check failed: {}",
294                response.status()
295            )));
296        }
297
298        Ok(())
299    }
300
301    /// Close connections
302    pub async fn close(&self) -> Result<()> {
303        // HTTP client doesn't need explicit closing
304        Ok(())
305    }
306
307    /// Batch store vectors
308    pub async fn batch_store(&self, points: &[VectorPoint]) -> Result<()> {
309        let url = format!("{}/collections/{}/points", self.url, self.collection);
310        let qdrant_points: Vec<serde_json::Value> = points
311            .iter()
312            .map(|point| {
313                serde_json::json!({
314                    "id": point.id,
315                    "vector": point.vector,
316                    "payload": point.metadata.clone().unwrap_or_default()
317                })
318            })
319            .collect();
320
321        let payload = serde_json::json!({
322            "points": qdrant_points
323        });
324
325        let mut request = self.client.put(&url).json(&payload);
326
327        if let Some(api_key) = &self.api_key {
328            request = request.header("api-key", api_key);
329        }
330
331        let response = request
332            .send()
333            .await
334            .map_err(|e| GatewayError::Storage(format!("Failed to batch store vectors: {}", e)))?;
335
336        if !response.status().is_success() {
337            return Err(GatewayError::Storage(format!(
338                "Failed to batch store vectors: {}",
339                response.status()
340            )));
341        }
342
343        debug!("Batch stored {} vectors", points.len());
344        Ok(())
345    }
346
347    /// Count vectors in collection
348    pub async fn count(&self) -> Result<u64> {
349        let url = format!("{}/collections/{}", self.url, self.collection);
350        let mut request = self.client.get(&url);
351
352        if let Some(api_key) = &self.api_key {
353            request = request.header("api-key", api_key);
354        }
355
356        let response = request
357            .send()
358            .await
359            .map_err(|e| GatewayError::Storage(format!("Failed to get collection info: {}", e)))?;
360
361        if !response.status().is_success() {
362            return Err(GatewayError::Storage(format!(
363                "Failed to get collection info: {}",
364                response.status()
365            )));
366        }
367
368        let result: serde_json::Value = response.json().await.map_err(|e| {
369            GatewayError::Storage(format!("Failed to parse collection info: {}", e))
370        })?;
371
372        if let Some(count) = result["result"]["points_count"].as_u64() {
373            Ok(count)
374        } else {
375            Ok(0)
376        }
377    }
378}
litellm_rs/storage/vector/qdrant.rs

litellm_rs/storage/vector/
qdrant.rs
