Skip to main content

rig_vectorize/client/
mod.rs

1//! Cloudflare Vectorize HTTP client.
2//!
3//! This module contains the HTTP client for interacting with the Cloudflare Vectorize API.
4//! It is designed to be potentially extracted into a standalone crate in the future.
5
6mod 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
21/// Base URL for the Cloudflare API.
22const CLOUDFLARE_API_BASE_URL: &str = "https://api.cloudflare.com/client/v4";
23
24/// HTTP client wrapper for Vectorize API operations.
25#[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    /// Creates a new Vectorize client.
35    ///
36    /// # Arguments
37    /// * `account_id` - Cloudflare account ID
38    /// * `index_name` - Name of the Vectorize index
39    /// * `api_token` - Cloudflare API token with Vectorize permissions
40    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    /// Returns the base URL for the index endpoints.
54    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    /// Performs a vector similarity query.
62    #[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    /// Upserts vectors (inserts or updates if ID already exists).
103    ///
104    /// This is the preferred method for inserting documents as it's idempotent.
105    /// Up to 5000 vectors can be upserted per request via the HTTP API.
106    #[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    /// Deletes vectors by their IDs.
150    ///
151    /// Up to 1000 vector IDs can be deleted per request.
152    #[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    /// Lists vector IDs in the index (paginated).
193    ///
194    /// Returns up to `limit` vector IDs (max 1000, default 100).
195    /// Use the `next_cursor` from the response to fetch the next page.
196    #[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}