Skip to main content

openai_compat/resources/
vector_stores.rs

1//! `/vector_stores` resource, mirroring
2//! `openai-python/src/openai/resources/vector_stores/`. Covers the store CRUD +
3//! search endpoints plus the `.files(id)` and `.file_batches(id)` sub-resources.
4//!
5//! Every vector-store endpoint requires the beta header
6//! `OpenAI-Beta: assistants=v2`; it is attached to every request here.
7
8use 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
22/// The beta header sent on every request in this module.
23const BETA_HEADER: (&str, &str) = ("OpenAI-Beta", "assistants=v2");
24
25/// Prepend the beta header so a caller-supplied header of the same name
26/// (added later) would still win, matching the Python header-merge semantics.
27fn 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
34/// Build JSON request options from a serializable body, plus the beta header.
35fn beta_json<B: serde::Serialize>(body: &B) -> Result<RequestOptions, OpenAIError> {
36    Ok(with_beta(RequestOptions::json(serde_json::to_value(body)?)))
37}
38
39/// Paginate a cursor list endpoint while sending the beta header on every page.
40async 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/// The vector stores resource.
55#[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    /// Create a vector store.
66    pub async fn create(
67        &self,
68        params: VectorStoreCreateParams,
69    ) -> Result<VectorStore, OpenAIError> {
70        self.client
71            .execute(Method::POST, "/vector_stores", beta_json(&params)?)
72            .await
73    }
74
75    /// Retrieve a vector store by id.
76    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    /// Update a vector store (POST, not PATCH).
87    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(&params)?,
97            )
98            .await
99    }
100
101    /// List vector stores (single page).
102    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    /// List all vector stores, following pagination cursors.
117    pub async fn list_all(&self) -> Result<Vec<VectorStore>, OpenAIError> {
118        paginate_all_beta(&self.client, "/vector_stores", Vec::new()).await
119    }
120
121    /// Delete a vector store.
122    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    /// Search a vector store for relevant chunks (POST, returns a single page).
136    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(&params)?,
146            )
147            .await
148    }
149
150    /// Files sub-resource bound to a vector store id.
151    pub fn files(&self, vector_store_id: impl Into<String>) -> VectorStoreFiles {
152        VectorStoreFiles::new(self.client.clone(), vector_store_id.into())
153    }
154
155    /// File batches sub-resource bound to a vector store id.
156    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/// The files sub-resource of a single vector store.
165#[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    /// Attach a file to the vector store.
184    pub async fn create(
185        &self,
186        params: VectorStoreFileCreateParams,
187    ) -> Result<VectorStoreFile, OpenAIError> {
188        self.client
189            .execute(Method::POST, &self.base(), beta_json(&params)?)
190            .await
191    }
192
193    /// Retrieve a vector store file.
194    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    /// Update a vector store file's attributes (POST).
205    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    /// Remove a file from the vector store.
221    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    /// List files in the vector store (single page).
235    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    /// List all files in the vector store, following pagination cursors.
250    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    /// Retrieve the parsed content of a vector store file.
259    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/// The file batches sub-resource of a single vector store.
274#[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    /// Create a file batch.
293    pub async fn create(
294        &self,
295        params: VectorStoreFileBatchCreateParams,
296    ) -> Result<VectorStoreFileBatch, OpenAIError> {
297        self.client
298            .execute(Method::POST, &self.base(), beta_json(&params)?)
299            .await
300    }
301
302    /// Retrieve a file batch.
303    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    /// Cancel a file batch (POST, no body).
317    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    /// List the files in a file batch (single page).
331    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}