1use reqwest::Method;
9use serde::de::DeserializeOwned;
10
11use crate::client::Client;
12use crate::error::OpenAIError;
13use crate::pagination::{HasId, List};
14use crate::request::RequestOptions;
15use crate::types::vector_stores::{
16 FileContentResponse, VectorStore, VectorStoreCreateParams, VectorStoreDeleted, VectorStoreFile,
17 VectorStoreFileBatch, VectorStoreFileBatchCreateParams, VectorStoreFileCreateParams,
18 VectorStoreFileDeleted, VectorStoreFileListParams, VectorStoreListParams,
19 VectorStoreSearchParams, VectorStoreSearchResponse, VectorStoreUpdateParams,
20};
21
22const BETA_HEADER: (&str, &str) = ("OpenAI-Beta", "assistants=v2");
24
25fn with_beta(mut options: RequestOptions) -> RequestOptions {
28 options
29 .extra_headers
30 .insert(0, (BETA_HEADER.0.to_string(), BETA_HEADER.1.to_string()));
31 options
32}
33
34fn beta_json<B: serde::Serialize>(body: &B) -> Result<RequestOptions, OpenAIError> {
36 Ok(with_beta(RequestOptions::json(serde_json::to_value(body)?)))
37}
38
39async fn paginate_all_beta<T: HasId + DeserializeOwned>(
41 client: &Client,
42 path: &str,
43 query: Vec<(String, String)>,
44) -> Result<Vec<T>, OpenAIError> {
45 client
46 .paginate_all_with_headers(
47 path,
48 query,
49 vec![(BETA_HEADER.0.to_string(), BETA_HEADER.1.to_string())],
50 )
51 .await
52}
53
54#[derive(Debug, Clone)]
56pub struct VectorStores {
57 client: Client,
58}
59
60impl VectorStores {
61 pub(crate) fn new(client: Client) -> Self {
62 Self { client }
63 }
64
65 pub async fn create(
67 &self,
68 params: VectorStoreCreateParams,
69 ) -> Result<VectorStore, OpenAIError> {
70 self.client
71 .execute(Method::POST, "/vector_stores", beta_json(¶ms)?)
72 .await
73 }
74
75 pub async fn retrieve(&self, vector_store_id: &str) -> Result<VectorStore, OpenAIError> {
77 self.client
78 .execute(
79 Method::GET,
80 &format!("/vector_stores/{vector_store_id}"),
81 with_beta(RequestOptions::default()),
82 )
83 .await
84 }
85
86 pub async fn update(
88 &self,
89 vector_store_id: &str,
90 params: VectorStoreUpdateParams,
91 ) -> Result<VectorStore, OpenAIError> {
92 self.client
93 .execute(
94 Method::POST,
95 &format!("/vector_stores/{vector_store_id}"),
96 beta_json(¶ms)?,
97 )
98 .await
99 }
100
101 pub async fn list(
103 &self,
104 params: Option<VectorStoreListParams>,
105 ) -> Result<List<VectorStore>, OpenAIError> {
106 let query = params.unwrap_or_default().to_query();
107 self.client
108 .execute(
109 Method::GET,
110 "/vector_stores",
111 with_beta(RequestOptions::query(query)),
112 )
113 .await
114 }
115
116 pub async fn list_all(&self) -> Result<Vec<VectorStore>, OpenAIError> {
118 paginate_all_beta(&self.client, "/vector_stores", Vec::new()).await
119 }
120
121 pub async fn delete(
123 &self,
124 vector_store_id: &str,
125 ) -> Result<VectorStoreDeleted, OpenAIError> {
126 self.client
127 .execute(
128 Method::DELETE,
129 &format!("/vector_stores/{vector_store_id}"),
130 with_beta(RequestOptions::default()),
131 )
132 .await
133 }
134
135 pub async fn search(
137 &self,
138 vector_store_id: &str,
139 params: VectorStoreSearchParams,
140 ) -> Result<List<VectorStoreSearchResponse>, OpenAIError> {
141 self.client
142 .execute(
143 Method::POST,
144 &format!("/vector_stores/{vector_store_id}/search"),
145 beta_json(¶ms)?,
146 )
147 .await
148 }
149
150 pub fn files(&self, vector_store_id: impl Into<String>) -> VectorStoreFiles {
152 VectorStoreFiles::new(self.client.clone(), vector_store_id.into())
153 }
154
155 pub fn file_batches(
157 &self,
158 vector_store_id: impl Into<String>,
159 ) -> VectorStoreFileBatches {
160 VectorStoreFileBatches::new(self.client.clone(), vector_store_id.into())
161 }
162}
163
164#[derive(Debug, Clone)]
166pub struct VectorStoreFiles {
167 client: Client,
168 vector_store_id: String,
169}
170
171impl VectorStoreFiles {
172 pub(crate) fn new(client: Client, vector_store_id: String) -> Self {
173 Self {
174 client,
175 vector_store_id,
176 }
177 }
178
179 fn base(&self) -> String {
180 format!("/vector_stores/{}/files", self.vector_store_id)
181 }
182
183 pub async fn create(
185 &self,
186 params: VectorStoreFileCreateParams,
187 ) -> Result<VectorStoreFile, OpenAIError> {
188 self.client
189 .execute(Method::POST, &self.base(), beta_json(¶ms)?)
190 .await
191 }
192
193 pub async fn retrieve(&self, file_id: &str) -> Result<VectorStoreFile, OpenAIError> {
195 self.client
196 .execute(
197 Method::GET,
198 &format!("{}/{file_id}", self.base()),
199 with_beta(RequestOptions::default()),
200 )
201 .await
202 }
203
204 pub async fn update(
206 &self,
207 file_id: &str,
208 attributes: serde_json::Value,
209 ) -> Result<VectorStoreFile, OpenAIError> {
210 let body = serde_json::json!({ "attributes": attributes });
211 self.client
212 .execute(
213 Method::POST,
214 &format!("{}/{file_id}", self.base()),
215 with_beta(RequestOptions::json(body)),
216 )
217 .await
218 }
219
220 pub async fn delete(
222 &self,
223 file_id: &str,
224 ) -> Result<VectorStoreFileDeleted, OpenAIError> {
225 self.client
226 .execute(
227 Method::DELETE,
228 &format!("{}/{file_id}", self.base()),
229 with_beta(RequestOptions::default()),
230 )
231 .await
232 }
233
234 pub async fn list(
236 &self,
237 params: Option<VectorStoreFileListParams>,
238 ) -> Result<List<VectorStoreFile>, OpenAIError> {
239 let query = params.unwrap_or_default().to_query();
240 self.client
241 .execute(
242 Method::GET,
243 &self.base(),
244 with_beta(RequestOptions::query(query)),
245 )
246 .await
247 }
248
249 pub async fn list_all(
251 &self,
252 params: Option<VectorStoreFileListParams>,
253 ) -> Result<Vec<VectorStoreFile>, OpenAIError> {
254 let query = params.unwrap_or_default().to_query();
255 paginate_all_beta(&self.client, &self.base(), query).await
256 }
257
258 pub async fn content(
260 &self,
261 file_id: &str,
262 ) -> Result<List<FileContentResponse>, OpenAIError> {
263 self.client
264 .execute(
265 Method::GET,
266 &format!("{}/{file_id}/content", self.base()),
267 with_beta(RequestOptions::default()),
268 )
269 .await
270 }
271}
272
273#[derive(Debug, Clone)]
275pub struct VectorStoreFileBatches {
276 client: Client,
277 vector_store_id: String,
278}
279
280impl VectorStoreFileBatches {
281 pub(crate) fn new(client: Client, vector_store_id: String) -> Self {
282 Self {
283 client,
284 vector_store_id,
285 }
286 }
287
288 fn base(&self) -> String {
289 format!("/vector_stores/{}/file_batches", self.vector_store_id)
290 }
291
292 pub async fn create(
294 &self,
295 params: VectorStoreFileBatchCreateParams,
296 ) -> Result<VectorStoreFileBatch, OpenAIError> {
297 self.client
298 .execute(Method::POST, &self.base(), beta_json(¶ms)?)
299 .await
300 }
301
302 pub async fn retrieve(
304 &self,
305 batch_id: &str,
306 ) -> Result<VectorStoreFileBatch, OpenAIError> {
307 self.client
308 .execute(
309 Method::GET,
310 &format!("{}/{batch_id}", self.base()),
311 with_beta(RequestOptions::default()),
312 )
313 .await
314 }
315
316 pub async fn cancel(
318 &self,
319 batch_id: &str,
320 ) -> Result<VectorStoreFileBatch, OpenAIError> {
321 self.client
322 .execute(
323 Method::POST,
324 &format!("{}/{batch_id}/cancel", self.base()),
325 with_beta(RequestOptions::default()),
326 )
327 .await
328 }
329
330 pub async fn list_files(
332 &self,
333 batch_id: &str,
334 params: Option<VectorStoreFileListParams>,
335 ) -> Result<List<VectorStoreFile>, OpenAIError> {
336 let query = params.unwrap_or_default().to_query();
337 self.client
338 .execute(
339 Method::GET,
340 &format!("{}/{batch_id}/files", self.base()),
341 with_beta(RequestOptions::query(query)),
342 )
343 .await
344 }
345}