vectorizer_rust_sdk/
client.rs

1//! Simplified client for Vectorizer
2
3use crate::error::{VectorizerError, Result};
4use crate::models::*;
5use reqwest::{Client, ClientBuilder, header::{HeaderMap, HeaderValue, CONTENT_TYPE}};
6use serde_json;
7
8/// Simplified client for Vectorizer
9pub struct VectorizerClient {
10    http_client: Client,
11    base_url: String,
12    api_key: Option<String>,
13}
14
15impl VectorizerClient {
16    /// Get the base URL
17    pub fn base_url(&self) -> &str {
18        &self.base_url
19    }
20
21    /// Create a new client
22    pub fn new_default() -> Result<Self> {
23        let mut headers = HeaderMap::new();
24        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
25
26        let client = ClientBuilder::new()
27            .timeout(std::time::Duration::from_secs(30))
28            .default_headers(headers)
29            .build()
30            .map_err(|e| VectorizerError::configuration(format!("Failed to create HTTP client: {}", e)))?;
31
32        Ok(Self {
33            http_client: client,
34            base_url: "http://localhost:15001".to_string(),
35            api_key: None,
36        })
37    }
38
39    /// Create client with custom URL
40    pub fn new_with_url(base_url: &str) -> Result<Self> {
41        let mut headers = HeaderMap::new();
42        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
43
44        let client = ClientBuilder::new()
45            .timeout(std::time::Duration::from_secs(30))
46            .default_headers(headers)
47            .build()
48            .map_err(|e| VectorizerError::configuration(format!("Failed to create HTTP client: {}", e)))?;
49
50        Ok(Self {
51            http_client: client,
52            base_url: base_url.to_string(),
53            api_key: None,
54        })
55    }
56
57    /// Create client with API key
58    pub fn new_with_api_key(base_url: &str, api_key: &str) -> Result<Self> {
59        let mut headers = HeaderMap::new();
60        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
61        headers.insert("Authorization", HeaderValue::from_str(&format!("Bearer {}", api_key))
62            .map_err(|e| VectorizerError::configuration(format!("Invalid API key: {}", e)))?);
63
64        let client = ClientBuilder::new()
65            .timeout(std::time::Duration::from_secs(30))
66            .default_headers(headers)
67            .build()
68            .map_err(|e| VectorizerError::configuration(format!("Failed to create HTTP client: {}", e)))?;
69
70        Ok(Self {
71            http_client: client,
72            base_url: base_url.to_string(),
73            api_key: Some(api_key.to_string()),
74        })
75    }
76
77    /// Health check
78    pub async fn health_check(&self) -> Result<HealthStatus> {
79        let response = self.make_request("GET", "/api/v1/health", None).await?;
80        let health: HealthStatus = serde_json::from_str(&response)
81            .map_err(|e| VectorizerError::server(format!("Failed to parse health check response: {}", e)))?;
82        Ok(health)
83    }
84
85    /// List collections
86    pub async fn list_collections(&self) -> Result<Vec<CollectionInfo>> {
87        let response = self.make_request("GET", "/api/v1/collections", None).await?;
88        let collections_response: CollectionsResponse = serde_json::from_str(&response)
89            .map_err(|e| VectorizerError::server(format!("Failed to parse collections response: {}", e)))?;
90        Ok(collections_response.collections)
91    }
92
93    /// Search vectors
94    pub async fn search_vectors(
95        &self,
96        collection: &str,
97        query: &str,
98        limit: Option<usize>,
99        score_threshold: Option<f32>,
100    ) -> Result<SearchResponse> {
101        let mut payload = serde_json::Map::new();
102        payload.insert("query".to_string(), serde_json::Value::String(query.to_string()));
103        payload.insert("limit".to_string(), serde_json::Value::Number(limit.unwrap_or(10).into()));
104
105        if let Some(threshold) = score_threshold {
106            payload.insert("score_threshold".to_string(), serde_json::Value::Number(serde_json::Number::from_f64(threshold as f64).unwrap()));
107        }
108
109        let response = self.make_request("POST", &format!("/api/v1/collections/{}/search/text", collection), Some(serde_json::Value::Object(payload))).await?;
110        let search_response: SearchResponse = serde_json::from_str(&response)
111            .map_err(|e| VectorizerError::server(format!("Failed to parse search response: {}", e)))?;
112        Ok(search_response)
113    }
114
115    /// Create collection
116    pub async fn create_collection(
117        &self,
118        name: &str,
119        dimension: usize,
120        metric: Option<SimilarityMetric>,
121    ) -> Result<CollectionInfo> {
122        let mut payload = serde_json::Map::new();
123        payload.insert("name".to_string(), serde_json::Value::String(name.to_string()));
124        payload.insert("dimension".to_string(), serde_json::Value::Number(dimension.into()));
125        payload.insert("metric".to_string(), serde_json::Value::String(format!("{:?}", metric.unwrap_or_default()).to_lowercase()));
126
127        let response = self.make_request("POST", "/api/v1/collections", Some(serde_json::Value::Object(payload))).await?;
128        let create_response: CreateCollectionResponse = serde_json::from_str(&response)
129            .map_err(|e| VectorizerError::server(format!("Failed to parse create collection response: {}", e)))?;
130
131        // Create a basic CollectionInfo from the response
132        let info = CollectionInfo {
133            name: create_response.collection,
134            dimension: dimension,
135            metric: format!("{:?}", metric.unwrap_or_default()).to_lowercase(),
136            vector_count: 0,
137            document_count: 0,
138            created_at: "".to_string(),
139            updated_at: "".to_string(),
140            indexing_status: crate::models::IndexingStatus {
141                status: "created".to_string(),
142                progress: 0.0,
143                total_documents: 0,
144                processed_documents: 0,
145                vector_count: 0,
146                estimated_time_remaining: None,
147                last_updated: "".to_string(),
148            },
149        };
150        Ok(info)
151    }
152
153    /// Insert texts
154    pub async fn insert_texts(
155        &self,
156        collection: &str,
157        texts: Vec<BatchTextRequest>,
158    ) -> Result<BatchResponse> {
159        let payload = serde_json::json!({
160            "texts": texts
161        });
162
163        let response = self.make_request("POST", &format!("/api/v1/collections/{}/documents", collection), Some(serde_json::to_value(payload)?)).await?;
164        let batch_response: BatchResponse = serde_json::from_str(&response)
165            .map_err(|e| VectorizerError::server(format!("Failed to parse insert texts response: {}", e)))?;
166        Ok(batch_response)
167    }
168
169    /// Delete collection
170    pub async fn delete_collection(&self, name: &str) -> Result<()> {
171        self.make_request("DELETE", &format!("/api/v1/collections/{}", name), None).await?;
172        Ok(())
173    }
174
175    /// Get vector
176    pub async fn get_vector(&self, collection: &str, vector_id: &str) -> Result<Vector> {
177        let response = self.make_request("GET", &format!("/api/v1/collections/{}/vectors/{}", collection, vector_id), None).await?;
178        let vector: Vector = serde_json::from_str(&response)
179            .map_err(|e| VectorizerError::server(format!("Failed to parse get vector response: {}", e)))?;
180        Ok(vector)
181    }
182
183    /// Get collection info
184    pub async fn get_collection_info(&self, collection: &str) -> Result<CollectionInfo> {
185        let response = self.make_request("GET", &format!("/api/v1/collections/{}", collection), None).await?;
186        let info: CollectionInfo = serde_json::from_str(&response)
187            .map_err(|e| VectorizerError::server(format!("Failed to parse collection info: {}", e)))?;
188        Ok(info)
189    }
190
191    /// Generate embeddings
192    pub async fn embed_text(&self, text: &str, model: Option<&str>) -> Result<EmbeddingResponse> {
193        let mut payload = serde_json::Map::new();
194        payload.insert("text".to_string(), serde_json::Value::String(text.to_string()));
195
196        if let Some(model) = model {
197            payload.insert("model".to_string(), serde_json::Value::String(model.to_string()));
198        }
199
200        let response = self.make_request("POST", "/api/v1/embed", Some(serde_json::Value::Object(payload))).await?;
201        let embedding_response: EmbeddingResponse = serde_json::from_str(&response)
202            .map_err(|e| VectorizerError::server(format!("Failed to parse embedding response: {}", e)))?;
203        Ok(embedding_response)
204    }
205
206    /// Make HTTP request
207    async fn make_request(
208        &self,
209        method: &str,
210        endpoint: &str,
211        payload: Option<serde_json::Value>,
212    ) -> Result<String> {
213        let url = format!("{}{}", self.base_url, endpoint);
214
215        let mut request = match method {
216            "GET" => self.http_client.get(&url),
217            "POST" => self.http_client.post(&url),
218            "PUT" => self.http_client.put(&url),
219            "DELETE" => self.http_client.delete(&url),
220            _ => return Err(VectorizerError::configuration(format!("Unsupported HTTP method: {}", method))),
221        };
222
223        if let Some(payload) = payload {
224            request = request.json(&payload);
225        }
226
227        let response = request
228            .send()
229            .await
230            .map_err(|e| VectorizerError::network(format!("Request failed: {}", e)))?;
231
232        let status = response.status();
233
234        if !status.is_success() {
235            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
236            return Err(VectorizerError::server(format!("HTTP {}: {}", status.as_u16(), error_text)));
237        }
238
239        let response_text = response
240            .text()
241            .await
242            .map_err(|e| VectorizerError::network(format!("Failed to read response: {}", e)))?;
243
244        Ok(response_text)
245    }
246}