lumosai_vector_milvus/
lib.rs

1//! # LumosAI Milvus Integration
2//!
3//! This crate provides Milvus integration for LumosAI vector storage,
4//! offering high-performance vector database capabilities with cloud-native features.
5//!
6//! ## Features
7//!
8//! - **High Performance**: Distributed vector database optimized for large-scale applications
9//! - **Cloud Native**: Kubernetes-ready with horizontal scaling
10//! - **Rich Indexing**: Multiple index types (IVF, HNSW, ANNOY, etc.)
11//! - **Metadata Filtering**: Complex filtering with boolean expressions
12//! - **Multi-tenancy**: Collection-based isolation and resource management
13//! - **ACID Transactions**: Consistency guarantees for critical operations
14//! - **Real-time**: Support for real-time data ingestion and querying
15//!
16//! ## Quick Start
17//!
18//! ```rust
19//! use lumosai_vector_milvus::{MilvusStorage, MilvusConfig};
20//! use lumosai_vector_core::traits::VectorStorage;
21//!
22//! #[tokio::main]
23//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
24//!     // Create Milvus storage
25//!     let config = MilvusConfig::new("http://localhost:19530")
26//!         .with_database("default")
27//!         .with_auth("username", "password");
28//!     let storage = MilvusStorage::new(config).await?;
29//!     
30//!     // Create a collection
31//!     let index_config = IndexConfig::new("documents", 384)
32//!         .with_metric(SimilarityMetric::Cosine);
33//!     storage.create_index(index_config).await?;
34//!     
35//!     // Insert documents
36//!     let docs = vec![
37//!         Document::new("doc1", "Hello world")
38//!             .with_embedding(vec![0.1; 384])
39//!             .with_metadata("category", "greeting"),
40//!     ];
41//!     storage.upsert_documents("documents", docs).await?;
42//!     
43//!     Ok(())
44//! }
45//! ```
46
47// Removed unused imports
48
49pub mod storage;
50pub mod config;
51pub mod error;
52pub mod client;
53pub mod types;
54
55pub use storage::MilvusStorage;
56pub use config::{MilvusConfig, MilvusConfigBuilder};
57pub use error::{MilvusError, MilvusResult};
58pub use client::MilvusClient;
59pub use types::{MilvusEntity, AuthRequest, AuthResponse, CollectionInfo, CollectionSchema};
60
61// Re-export core types for convenience (avoiding conflicts)
62pub use lumosai_vector_core::types::{
63    Document, SimilarityMetric, MetadataValue
64};
65pub use lumosai_vector_core::PerformanceMetrics;
66pub use lumosai_vector_core::traits::VectorStorage;
67
68/// Milvus client for managing connections and databases
69#[derive(Clone)]
70pub struct MilvusConnection {
71    /// HTTP client
72    client: reqwest::Client,
73    
74    /// Configuration
75    config: MilvusConfig,
76    
77    /// Authentication token (if using token-based auth)
78    auth_token: Option<String>,
79}
80
81impl MilvusConnection {
82    /// Create a new Milvus connection
83    pub async fn new(config: MilvusConfig) -> MilvusResult<Self> {
84        let client = reqwest::Client::builder()
85            .timeout(config.timeout)
86            .build()
87            .map_err(|e| MilvusError::Connection(e.to_string()))?;
88        
89        let mut connection = Self {
90            client,
91            config,
92            auth_token: None,
93        };
94        
95        // Authenticate if credentials are provided
96        if let Some(ref _auth) = connection.config.auth {
97            connection.authenticate().await?;
98        }
99        
100        Ok(connection)
101    }
102    
103    /// Authenticate with Milvus
104    async fn authenticate(&mut self) -> MilvusResult<()> {
105        if let Some(ref auth) = self.config.auth {
106            let auth_request = AuthRequest {
107                username: auth.username.clone(),
108                password: auth.password.clone(),
109            };
110            
111            let url = format!("{}/v1/auth/login", self.config.endpoint);
112            let response = self.client
113                .post(&url)
114                .json(&auth_request)
115                .send()
116                .await
117                .map_err(|e| MilvusError::Connection(e.to_string()))?;
118            
119            if response.status().is_success() {
120                let auth_response: AuthResponse = response
121                    .json()
122                    .await
123                    .map_err(|e| MilvusError::Serialization(e.to_string()))?;
124                
125                self.auth_token = Some(auth_response.token);
126                tracing::info!("Successfully authenticated with Milvus");
127            } else {
128                let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
129                return Err(MilvusError::Authentication(format!("Authentication failed: {}", error_text)));
130            }
131        }
132        
133        Ok(())
134    }
135    
136    /// Get the HTTP client with authentication headers
137    pub fn authenticated_client(&self) -> reqwest::RequestBuilder {
138        let mut builder = self.client.get(&self.config.endpoint);
139        
140        if let Some(ref token) = self.auth_token {
141            builder = builder.header("Authorization", format!("Bearer {}", token));
142        }
143        
144        builder
145    }
146    
147    /// Make an authenticated POST request
148    pub fn post(&self, url: &str) -> reqwest::RequestBuilder {
149        let mut builder = self.client.post(url);
150        
151        if let Some(ref token) = self.auth_token {
152            builder = builder.header("Authorization", format!("Bearer {}", token));
153        }
154        
155        builder
156    }
157    
158    /// Make an authenticated GET request
159    pub fn get(&self, url: &str) -> reqwest::RequestBuilder {
160        let mut builder = self.client.get(url);
161        
162        if let Some(ref token) = self.auth_token {
163            builder = builder.header("Authorization", format!("Bearer {}", token));
164        }
165        
166        builder
167    }
168    
169    /// Make an authenticated DELETE request
170    pub fn delete(&self, url: &str) -> reqwest::RequestBuilder {
171        let mut builder = self.client.delete(url);
172        
173        if let Some(ref token) = self.auth_token {
174            builder = builder.header("Authorization", format!("Bearer {}", token));
175        }
176        
177        builder
178    }
179    
180    /// Get the configuration
181    pub fn config(&self) -> &MilvusConfig {
182        &self.config
183    }
184    
185    /// Check connection health
186    pub async fn health_check(&self) -> MilvusResult<()> {
187        let url = format!("{}/health", self.config.endpoint);
188        let response = self.get(&url)
189            .send()
190            .await
191            .map_err(|e| MilvusError::Connection(e.to_string()))?;
192        
193        if response.status().is_success() {
194            Ok(())
195        } else {
196            Err(MilvusError::Connection("Health check failed".to_string()))
197        }
198    }
199}
200
201/// Create a new Milvus storage instance
202pub async fn create_milvus_storage(endpoint: &str) -> MilvusResult<MilvusStorage> {
203    let config = MilvusConfig::new(endpoint);
204    MilvusStorage::new(config).await
205}
206
207/// Create a new Milvus storage instance with configuration
208pub async fn create_milvus_storage_with_config(config: MilvusConfig) -> MilvusResult<MilvusStorage> {
209    MilvusStorage::new(config).await
210}
211
212/// Utility functions for Milvus operations
213pub mod utils {
214    use super::*;
215    use lumosai_vector_core::types::{Document, MetadataValue};
216    
217    /// Convert LumosAI documents to Milvus entities
218    pub fn documents_to_entities(documents: &[Document]) -> MilvusResult<Vec<MilvusEntity>> {
219        let mut entities = Vec::new();
220        
221        for doc in documents {
222            let embedding = doc.embedding.as_ref()
223                .ok_or_else(|| MilvusError::InvalidData("Document missing embedding".to_string()))?;
224            
225            let entity = MilvusEntity {
226                id: doc.id.clone(),
227                vector: embedding.clone(),
228                content: doc.content.clone(),
229                metadata: doc.metadata.clone(),
230            };
231            
232            entities.push(entity);
233        }
234        
235        Ok(entities)
236    }
237    
238    /// Convert Milvus entities to LumosAI documents
239    pub fn entities_to_documents(entities: &[MilvusEntity]) -> Vec<Document> {
240        entities.iter().map(|entity| {
241            let mut document = Document::new(&entity.id, &entity.content);
242            document.embedding = Some(entity.vector.clone());
243            document.metadata = entity.metadata.clone();
244            document
245        }).collect()
246    }
247    
248    /// Convert metadata value to Milvus-compatible value
249    pub fn metadata_value_to_milvus(value: &MetadataValue) -> serde_json::Value {
250        match value {
251            MetadataValue::String(s) => serde_json::Value::String(s.clone()),
252            MetadataValue::Integer(i) => serde_json::Value::Number(serde_json::Number::from(*i)),
253            MetadataValue::Float(f) => {
254                serde_json::Number::from_f64(*f)
255                    .map(serde_json::Value::Number)
256                    .unwrap_or(serde_json::Value::Null)
257            }
258            MetadataValue::Boolean(b) => serde_json::Value::Bool(*b),
259            MetadataValue::Null => serde_json::Value::Null,
260            MetadataValue::Array(arr) => {
261                let json_arr: Vec<serde_json::Value> = arr
262                    .iter()
263                    .map(metadata_value_to_milvus)
264                    .collect();
265                serde_json::Value::Array(json_arr)
266            }
267            MetadataValue::Object(obj) => {
268                let json_obj: serde_json::Map<String, serde_json::Value> = obj
269                    .iter()
270                    .map(|(k, v)| (k.clone(), metadata_value_to_milvus(v)))
271                    .collect();
272                serde_json::Value::Object(json_obj)
273            }
274        }
275    }
276    
277    /// Validate vector dimensions
278    pub fn validate_vector_dimension(vectors: &[Vec<f32>]) -> MilvusResult<usize> {
279        if vectors.is_empty() {
280            return Err(MilvusError::InvalidData("No vectors provided".to_string()));
281        }
282        
283        let expected_dim = vectors[0].len();
284        if expected_dim == 0 {
285            return Err(MilvusError::InvalidData("Vector dimension cannot be zero".to_string()));
286        }
287        
288        for (i, vector) in vectors.iter().enumerate() {
289            if vector.len() != expected_dim {
290                return Err(MilvusError::InvalidData(
291                    format!("Vector {} has dimension {} but expected {}", i, vector.len(), expected_dim)
292                ));
293            }
294        }
295        
296        Ok(expected_dim)
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use super::*;
303    
304    #[test]
305    fn test_config_creation() {
306        let config = MilvusConfig::new("http://localhost:19530");
307        assert_eq!(config.endpoint, "http://localhost:19530");
308        assert_eq!(config.database, "default");
309    }
310    
311    #[test]
312    fn test_utils_document_conversion() {
313        let doc = Document::new("test", "content")
314            .with_embedding(vec![1.0, 2.0, 3.0])
315            .with_metadata("key", "value");
316        
317        let entities = utils::documents_to_entities(&[doc]).unwrap();
318        assert_eq!(entities.len(), 1);
319        assert_eq!(entities[0].id, "test");
320        assert_eq!(entities[0].vector, vec![1.0, 2.0, 3.0]);
321        
322        let docs = utils::entities_to_documents(&entities);
323        assert_eq!(docs.len(), 1);
324        assert_eq!(docs[0].id, "test");
325    }
326    
327    #[test]
328    fn test_vector_validation() {
329        let vectors = vec![
330            vec![1.0, 2.0, 3.0],
331            vec![4.0, 5.0, 6.0],
332        ];
333        
334        let dim = utils::validate_vector_dimension(&vectors).unwrap();
335        assert_eq!(dim, 3);
336        
337        let invalid_vectors = vec![
338            vec![1.0, 2.0, 3.0],
339            vec![4.0, 5.0], // Wrong dimension
340        ];
341        
342        assert!(utils::validate_vector_dimension(&invalid_vectors).is_err());
343    }
344}
lumosai_vector_milvus/lib.rs

lumosai_vector_milvus/
lib.rs
