rig_vectorize/client/
mod.rs1mod error;
7mod filter;
8mod types;
9
10pub use error::VectorizeError;
11pub use filter::VectorizeFilter;
12pub use types::{
13 DeleteByIdsRequest, DeleteResult, ListVectorsResult, QueryRequest, QueryResult, ReturnMetadata,
14 UpsertRequest, UpsertResult, VectorIdEntry, VectorInput, VectorMatch,
15};
16
17use reqwest::Client;
18use tracing::instrument;
19use types::ApiResponse;
20
21const CLOUDFLARE_API_BASE_URL: &str = "https://api.cloudflare.com/client/v4";
23
24#[derive(Debug, Clone)]
26pub struct VectorizeClient {
27 http_client: Client,
28 account_id: String,
29 index_name: String,
30 api_token: String,
31}
32
33impl VectorizeClient {
34 pub fn new(
41 account_id: impl Into<String>,
42 index_name: impl Into<String>,
43 api_token: impl Into<String>,
44 ) -> Self {
45 Self {
46 http_client: Client::new(),
47 account_id: account_id.into(),
48 index_name: index_name.into(),
49 api_token: api_token.into(),
50 }
51 }
52
53 fn index_url(&self) -> String {
55 format!(
56 "{}/accounts/{}/vectorize/v2/indexes/{}",
57 CLOUDFLARE_API_BASE_URL, self.account_id, self.index_name
58 )
59 }
60
61 #[instrument(skip(self, request), fields(index = %self.index_name, top_k = request.top_k))]
63 pub async fn query(&self, request: QueryRequest) -> Result<QueryResult, VectorizeError> {
64 let url = format!("{}/query", self.index_url());
65
66 tracing::debug!("Sending query request to Vectorize");
67
68 let response = self
69 .http_client
70 .post(&url)
71 .bearer_auth(&self.api_token)
72 .json(&request)
73 .send()
74 .await?;
75
76 let response_text = response.text().await?;
77 tracing::debug!("Raw Vectorize response: {}", response_text);
78
79 let api_response: ApiResponse<QueryResult> = serde_json::from_str(&response_text)?;
80
81 if !api_response.success {
82 let error = api_response
83 .errors
84 .first()
85 .map(|e| VectorizeError::ApiError {
86 code: e.code,
87 message: e.message.clone(),
88 })
89 .unwrap_or_else(|| VectorizeError::ApiError {
90 code: 0,
91 message: "Unknown error".to_string(),
92 });
93 return Err(error);
94 }
95
96 api_response.result.ok_or_else(|| VectorizeError::ApiError {
97 code: 0,
98 message: "No result in successful response".to_string(),
99 })
100 }
101
102 #[instrument(skip(self, request), fields(index = %self.index_name, count = request.vectors.len()))]
107 pub async fn upsert(&self, request: UpsertRequest) -> Result<UpsertResult, VectorizeError> {
108 let url = format!("{}/upsert", self.index_url());
109
110 tracing::debug!(
111 "Sending upsert request to Vectorize with {} vectors",
112 request.vectors.len()
113 );
114
115 let response = self
116 .http_client
117 .post(&url)
118 .bearer_auth(&self.api_token)
119 .json(&request)
120 .send()
121 .await?;
122
123 let response_text = response.text().await?;
124 tracing::debug!("Raw Vectorize upsert response: {}", response_text);
125
126 let api_response: ApiResponse<UpsertResult> = serde_json::from_str(&response_text)?;
127
128 if !api_response.success {
129 let error = api_response
130 .errors
131 .first()
132 .map(|e| VectorizeError::ApiError {
133 code: e.code,
134 message: e.message.clone(),
135 })
136 .unwrap_or_else(|| VectorizeError::ApiError {
137 code: 0,
138 message: "Unknown error".to_string(),
139 });
140 return Err(error);
141 }
142
143 api_response.result.ok_or_else(|| VectorizeError::ApiError {
144 code: 0,
145 message: "No result in successful upsert response".to_string(),
146 })
147 }
148
149 #[instrument(skip(self, ids), fields(index = %self.index_name, count = ids.len()))]
153 pub async fn delete_by_ids(&self, ids: Vec<String>) -> Result<DeleteResult, VectorizeError> {
154 let url = format!("{}/delete_by_ids", self.index_url());
155
156 let request = DeleteByIdsRequest { ids };
157
158 let response = self
159 .http_client
160 .post(&url)
161 .bearer_auth(&self.api_token)
162 .json(&request)
163 .send()
164 .await?;
165
166 let response_text = response.text().await?;
167 tracing::debug!("Raw Vectorize delete response: {}", response_text);
168
169 let api_response: ApiResponse<DeleteResult> = serde_json::from_str(&response_text)?;
170
171 if !api_response.success {
172 let error = api_response
173 .errors
174 .first()
175 .map(|e| VectorizeError::ApiError {
176 code: e.code,
177 message: e.message.clone(),
178 })
179 .unwrap_or_else(|| VectorizeError::ApiError {
180 code: 0,
181 message: "Unknown error".to_string(),
182 });
183 return Err(error);
184 }
185
186 api_response.result.ok_or_else(|| VectorizeError::ApiError {
187 code: 0,
188 message: "No result in successful delete response".to_string(),
189 })
190 }
191
192 #[instrument(skip(self), fields(index = %self.index_name))]
197 pub async fn list_vectors(
198 &self,
199 limit: Option<u32>,
200 cursor: Option<&str>,
201 ) -> Result<ListVectorsResult, VectorizeError> {
202 let mut url = format!("{}/list", self.index_url());
203
204 let mut query_params = Vec::new();
205 if let Some(limit) = limit {
206 query_params.push(format!("count={}", limit));
207 }
208 if let Some(cursor) = cursor {
209 query_params.push(format!("cursor={}", cursor));
210 }
211 if !query_params.is_empty() {
212 url = format!("{}?{}", url, query_params.join("&"));
213 }
214
215 let response = self
216 .http_client
217 .get(&url)
218 .bearer_auth(&self.api_token)
219 .send()
220 .await?;
221
222 let response_text = response.text().await?;
223 tracing::debug!("Raw Vectorize list response: {}", response_text);
224
225 let api_response: ApiResponse<ListVectorsResult> = serde_json::from_str(&response_text)?;
226
227 if !api_response.success {
228 let error = api_response
229 .errors
230 .first()
231 .map(|e| VectorizeError::ApiError {
232 code: e.code,
233 message: e.message.clone(),
234 })
235 .unwrap_or_else(|| VectorizeError::ApiError {
236 code: 0,
237 message: "Unknown error".to_string(),
238 });
239 return Err(error);
240 }
241
242 api_response.result.ok_or_else(|| VectorizeError::ApiError {
243 code: 0,
244 message: "No result in successful list response".to_string(),
245 })
246 }
247}