openai_ergonomic/builders/
vector_stores.rs

1//! Vector Stores API builders.
2//!
3//! This module provides ergonomic builders for `OpenAI` Vector Stores API operations,
4//! including creating and managing vector stores for RAG (Retrieval-Augmented Generation) use cases.
5//!
6//! Vector stores are used to store and search through documents that can be used
7//! by assistants with file search capabilities.
8
9use std::collections::HashMap;
10
11/// Builder for creating a new vector store.
12///
13/// Vector stores are collections of files that can be searched through
14/// using semantic similarity. They're commonly used for RAG applications.
15#[derive(Debug, Clone)]
16pub struct VectorStoreBuilder {
17    name: Option<String>,
18    file_ids: Vec<String>,
19    expires_after: Option<VectorStoreExpirationPolicy>,
20    metadata: HashMap<String, String>,
21}
22
23/// Expiration policy for vector stores.
24#[derive(Debug, Clone)]
25pub struct VectorStoreExpirationPolicy {
26    /// Number of days after which the vector store expires.
27    pub days: i32,
28}
29
30impl VectorStoreBuilder {
31    /// Create a new vector store builder.
32    ///
33    /// # Examples
34    ///
35    /// ```rust
36    /// use openai_ergonomic::builders::vector_stores::VectorStoreBuilder;
37    ///
38    /// let builder = VectorStoreBuilder::new()
39    ///     .name("My Knowledge Base")
40    ///     .file_ids(vec!["file-123".to_string(), "file-456".to_string()]);
41    /// ```
42    #[must_use]
43    pub fn new() -> Self {
44        Self {
45            name: None,
46            file_ids: Vec::new(),
47            expires_after: None,
48            metadata: HashMap::new(),
49        }
50    }
51
52    /// Set the vector store's name.
53    #[must_use]
54    pub fn name(mut self, name: impl Into<String>) -> Self {
55        self.name = Some(name.into());
56        self
57    }
58
59    /// Set the file IDs to include in the vector store.
60    #[must_use]
61    pub fn file_ids(mut self, file_ids: Vec<String>) -> Self {
62        self.file_ids = file_ids;
63        self
64    }
65
66    /// Add a single file ID to the vector store.
67    #[must_use]
68    pub fn add_file(mut self, file_id: impl Into<String>) -> Self {
69        self.file_ids.push(file_id.into());
70        self
71    }
72
73    /// Add multiple file IDs to the vector store.
74    #[must_use]
75    pub fn add_files<I, S>(mut self, file_ids: I) -> Self
76    where
77        I: IntoIterator<Item = S>,
78        S: Into<String>,
79    {
80        self.file_ids
81            .extend(file_ids.into_iter().map(std::convert::Into::into));
82        self
83    }
84
85    /// Clear all file IDs from the vector store.
86    #[must_use]
87    pub fn clear_files(mut self) -> Self {
88        self.file_ids.clear();
89        self
90    }
91
92    /// Set expiration policy for the vector store.
93    #[must_use]
94    pub fn expires_after_days(mut self, days: i32) -> Self {
95        self.expires_after = Some(VectorStoreExpirationPolicy { days });
96        self
97    }
98
99    /// Add metadata to the vector store.
100    #[must_use]
101    pub fn metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
102        self.metadata.insert(key.into(), value.into());
103        self
104    }
105
106    /// Get the name of this vector store.
107    #[must_use]
108    pub fn name_ref(&self) -> Option<&str> {
109        self.name.as_deref()
110    }
111
112    /// Get the file IDs for this vector store.
113    #[must_use]
114    pub fn file_ids_ref(&self) -> &[String] {
115        &self.file_ids
116    }
117
118    /// Get the expiration policy for this vector store.
119    #[must_use]
120    pub fn expires_after_ref(&self) -> Option<&VectorStoreExpirationPolicy> {
121        self.expires_after.as_ref()
122    }
123
124    /// Get the metadata for this vector store.
125    #[must_use]
126    pub fn metadata_ref(&self) -> &HashMap<String, String> {
127        &self.metadata
128    }
129
130    /// Check if the vector store has any files.
131    #[must_use]
132    pub fn has_files(&self) -> bool {
133        !self.file_ids.is_empty()
134    }
135
136    /// Get the number of files in the vector store.
137    #[must_use]
138    pub fn file_count(&self) -> usize {
139        self.file_ids.len()
140    }
141}
142
143impl Default for VectorStoreBuilder {
144    fn default() -> Self {
145        Self::new()
146    }
147}
148
149/// Builder for vector store file operations.
150#[derive(Debug, Clone)]
151pub struct VectorStoreFileBuilder {
152    vector_store_id: String,
153    file_id: String,
154}
155
156impl VectorStoreFileBuilder {
157    /// Create a new vector store file builder.
158    #[must_use]
159    pub fn new(vector_store_id: impl Into<String>, file_id: impl Into<String>) -> Self {
160        Self {
161            vector_store_id: vector_store_id.into(),
162            file_id: file_id.into(),
163        }
164    }
165
166    /// Get the vector store ID.
167    #[must_use]
168    pub fn vector_store_id(&self) -> &str {
169        &self.vector_store_id
170    }
171
172    /// Get the file ID.
173    #[must_use]
174    pub fn file_id(&self) -> &str {
175        &self.file_id
176    }
177}
178
179/// Builder for searching through vector stores.
180#[derive(Debug, Clone)]
181pub struct VectorStoreSearchBuilder {
182    vector_store_id: String,
183    query: String,
184    limit: Option<i32>,
185    filter: HashMap<String, String>,
186}
187
188impl VectorStoreSearchBuilder {
189    /// Create a new vector store search builder.
190    #[must_use]
191    pub fn new(vector_store_id: impl Into<String>, query: impl Into<String>) -> Self {
192        Self {
193            vector_store_id: vector_store_id.into(),
194            query: query.into(),
195            limit: None,
196            filter: HashMap::new(),
197        }
198    }
199
200    /// Set the maximum number of results to return.
201    #[must_use]
202    pub fn limit(mut self, limit: i32) -> Self {
203        self.limit = Some(limit);
204        self
205    }
206
207    /// Add a filter to the search.
208    #[must_use]
209    pub fn filter(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
210        self.filter.insert(key.into(), value.into());
211        self
212    }
213
214    /// Get the vector store ID for this search.
215    #[must_use]
216    pub fn vector_store_id(&self) -> &str {
217        &self.vector_store_id
218    }
219
220    /// Get the search query.
221    #[must_use]
222    pub fn query(&self) -> &str {
223        &self.query
224    }
225
226    /// Get the search limit.
227    #[must_use]
228    pub fn limit_ref(&self) -> Option<i32> {
229        self.limit
230    }
231
232    /// Get the search filters.
233    #[must_use]
234    pub fn filter_ref(&self) -> &HashMap<String, String> {
235        &self.filter
236    }
237}
238
239/// Helper function to create a simple vector store with a name.
240#[must_use]
241pub fn simple_vector_store(name: impl Into<String>) -> VectorStoreBuilder {
242    VectorStoreBuilder::new().name(name)
243}
244
245/// Helper function to create a vector store with files.
246#[must_use]
247pub fn vector_store_with_files(
248    name: impl Into<String>,
249    file_ids: Vec<String>,
250) -> VectorStoreBuilder {
251    VectorStoreBuilder::new().name(name).file_ids(file_ids)
252}
253
254/// Helper function to create a temporary vector store that expires after a specified number of days.
255#[must_use]
256pub fn temporary_vector_store(
257    name: impl Into<String>,
258    expires_after_days: i32,
259) -> VectorStoreBuilder {
260    VectorStoreBuilder::new()
261        .name(name)
262        .expires_after_days(expires_after_days)
263}
264
265/// Helper function to add a file to a vector store.
266#[must_use]
267pub fn add_file_to_vector_store(
268    vector_store_id: impl Into<String>,
269    file_id: impl Into<String>,
270) -> VectorStoreFileBuilder {
271    VectorStoreFileBuilder::new(vector_store_id, file_id)
272}
273
274/// Helper function to search through a vector store.
275#[must_use]
276pub fn search_vector_store(
277    vector_store_id: impl Into<String>,
278    query: impl Into<String>,
279) -> VectorStoreSearchBuilder {
280    VectorStoreSearchBuilder::new(vector_store_id, query)
281}
282
283/// Helper function to search with a limit.
284#[must_use]
285pub fn search_vector_store_with_limit(
286    vector_store_id: impl Into<String>,
287    query: impl Into<String>,
288    limit: i32,
289) -> VectorStoreSearchBuilder {
290    VectorStoreSearchBuilder::new(vector_store_id, query).limit(limit)
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296
297    #[test]
298    fn test_vector_store_builder() {
299        let builder = VectorStoreBuilder::new()
300            .name("Test Store")
301            .add_file("file-1")
302            .add_file("file-2")
303            .expires_after_days(30)
304            .metadata("key", "value");
305
306        assert_eq!(builder.name_ref(), Some("Test Store"));
307        assert_eq!(builder.file_count(), 2);
308        assert_eq!(builder.file_ids_ref(), &["file-1", "file-2"]);
309        assert!(builder.has_files());
310        assert!(builder.expires_after_ref().is_some());
311        assert_eq!(builder.expires_after_ref().unwrap().days, 30);
312        assert_eq!(builder.metadata_ref().len(), 1);
313    }
314
315    #[test]
316    fn test_vector_store_builder_with_file_ids() {
317        let file_ids = vec![
318            "file-1".to_string(),
319            "file-2".to_string(),
320            "file-3".to_string(),
321        ];
322        let builder = VectorStoreBuilder::new()
323            .name("Bulk Files Store")
324            .file_ids(file_ids.clone());
325
326        assert_eq!(builder.name_ref(), Some("Bulk Files Store"));
327        assert_eq!(builder.file_ids_ref(), file_ids.as_slice());
328        assert_eq!(builder.file_count(), 3);
329        assert!(builder.has_files());
330    }
331
332    #[test]
333    fn test_vector_store_file_builder() {
334        let builder = VectorStoreFileBuilder::new("vs-123", "file-456");
335        assert_eq!(builder.vector_store_id(), "vs-123");
336        assert_eq!(builder.file_id(), "file-456");
337    }
338
339    #[test]
340    fn test_vector_store_search_builder() {
341        let builder = VectorStoreSearchBuilder::new("vs-123", "search query")
342            .limit(10)
343            .filter("category", "documentation");
344
345        assert_eq!(builder.vector_store_id(), "vs-123");
346        assert_eq!(builder.query(), "search query");
347        assert_eq!(builder.limit_ref(), Some(10));
348        assert_eq!(builder.filter_ref().len(), 1);
349        assert_eq!(
350            builder.filter_ref().get("category"),
351            Some(&"documentation".to_string())
352        );
353    }
354
355    #[test]
356    fn test_simple_vector_store_helper() {
357        let builder = simple_vector_store("Simple Store");
358        assert_eq!(builder.name_ref(), Some("Simple Store"));
359        assert!(!builder.has_files());
360    }
361
362    #[test]
363    fn test_vector_store_with_files_helper() {
364        let file_ids = vec!["file-1".to_string(), "file-2".to_string()];
365        let builder = vector_store_with_files("Files Store", file_ids.clone());
366        assert_eq!(builder.name_ref(), Some("Files Store"));
367        assert_eq!(builder.file_ids_ref(), file_ids.as_slice());
368        assert!(builder.has_files());
369    }
370
371    #[test]
372    fn test_temporary_vector_store_helper() {
373        let builder = temporary_vector_store("Temp Store", 7);
374        assert_eq!(builder.name_ref(), Some("Temp Store"));
375        assert!(builder.expires_after_ref().is_some());
376        assert_eq!(builder.expires_after_ref().unwrap().days, 7);
377    }
378
379    #[test]
380    fn test_add_file_to_vector_store_helper() {
381        let builder = add_file_to_vector_store("vs-123", "file-456");
382        assert_eq!(builder.vector_store_id(), "vs-123");
383        assert_eq!(builder.file_id(), "file-456");
384    }
385
386    #[test]
387    fn test_search_vector_store_helper() {
388        let builder = search_vector_store("vs-123", "test query");
389        assert_eq!(builder.vector_store_id(), "vs-123");
390        assert_eq!(builder.query(), "test query");
391        assert!(builder.limit_ref().is_none());
392    }
393
394    #[test]
395    fn test_search_vector_store_with_limit_helper() {
396        let builder = search_vector_store_with_limit("vs-123", "test query", 5);
397        assert_eq!(builder.vector_store_id(), "vs-123");
398        assert_eq!(builder.query(), "test query");
399        assert_eq!(builder.limit_ref(), Some(5));
400    }
401
402    #[test]
403    fn test_vector_store_builder_default() {
404        let builder = VectorStoreBuilder::default();
405        assert!(builder.name_ref().is_none());
406        assert!(!builder.has_files());
407        assert!(builder.expires_after_ref().is_none());
408        assert!(builder.metadata_ref().is_empty());
409    }
410
411    #[test]
412    fn test_vector_store_expiration_policy() {
413        let policy = VectorStoreExpirationPolicy { days: 15 };
414        assert_eq!(policy.days, 15);
415    }
416
417    #[test]
418    fn test_vector_store_builder_add_files() {
419        let builder = VectorStoreBuilder::new()
420            .name("Multi-File Store")
421            .add_file("file-1")
422            .add_files(vec!["file-2", "file-3", "file-4"])
423            .add_file("file-5");
424
425        assert_eq!(builder.name_ref(), Some("Multi-File Store"));
426        assert_eq!(builder.file_count(), 5);
427        assert_eq!(
428            builder.file_ids_ref(),
429            &["file-1", "file-2", "file-3", "file-4", "file-5"]
430        );
431        assert!(builder.has_files());
432    }
433
434    #[test]
435    fn test_vector_store_builder_clear_files() {
436        let builder = VectorStoreBuilder::new()
437            .add_files(vec!["file-1", "file-2", "file-3"])
438            .clear_files()
439            .add_file("file-new");
440
441        assert_eq!(builder.file_count(), 1);
442        assert_eq!(builder.file_ids_ref(), &["file-new"]);
443        assert!(builder.has_files());
444    }
445}
openai_ergonomic/builders/vector_stores.rs

openai_ergonomic/builders/
vector_stores.rs
