1use crate::error::{VectorizerError, Result};
4use crate::models::*;
5use reqwest::{Client, ClientBuilder, header::{HeaderMap, HeaderValue, CONTENT_TYPE}};
6use serde_json;
7
8pub struct VectorizerClient {
10 http_client: Client,
11 base_url: String,
12 api_key: Option<String>,
13}
14
15impl VectorizerClient {
16 pub fn base_url(&self) -> &str {
18 &self.base_url
19 }
20
21 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 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 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 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 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 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 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 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 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 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 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 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 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 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}