gemini_rust/file_search/
store_handle.rs

1use std::sync::Arc;
2use tracing::instrument;
3
4use crate::client::{Error, GeminiClient};
5use crate::file_search::model::FileSearchStore;
6use crate::file_search::{DocumentBuilder, ImportBuilder, UploadBuilder};
7
8/// A handle for managing a file search store.
9///
10/// Provides methods to upload files, import files, manage documents,
11/// and delete the store. The store persists indefinitely until explicitly deleted.
12///
13/// # Example
14///
15/// ```no_run
16/// use gemini_rust::prelude::*;
17/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
18/// let client = Gemini::new("API_KEY")?;
19///
20/// // Create a store
21/// let store = client
22///     .create_file_search_store()
23///     .with_display_name("My Store")
24///     .execute()
25///     .await?;
26///
27/// // Upload a file
28/// let data = b"Sample document content";
29/// let operation = store
30///     .upload(data.to_vec())
31///     .with_display_name("Sample Doc")
32///     .execute()
33///     .await?;
34/// # Ok(())
35/// # }
36/// ```
37#[derive(Debug, Clone)]
38pub struct FileSearchStoreHandle {
39    client: Arc<GeminiClient>,
40    store: FileSearchStore,
41}
42
43impl FileSearchStoreHandle {
44    pub fn new(client: Arc<GeminiClient>, store: FileSearchStore) -> Self {
45        Self { client, store }
46    }
47
48    pub fn name(&self) -> &str {
49        &self.store.name
50    }
51
52    pub fn display_name(&self) -> Option<&str> {
53        self.store.display_name.as_deref()
54    }
55
56    pub fn active_documents_count(&self) -> Option<i64> {
57        self.store.active_documents_count
58    }
59
60    pub fn pending_documents_count(&self) -> Option<i64> {
61        self.store.pending_documents_count
62    }
63
64    pub fn failed_documents_count(&self) -> Option<i64> {
65        self.store.failed_documents_count
66    }
67
68    pub fn size_bytes(&self) -> Option<i64> {
69        self.store.size_bytes
70    }
71
72    pub fn store(&self) -> &FileSearchStore {
73        &self.store
74    }
75
76    #[instrument(skip_all, fields(store.name = %self.store.name))]
77    pub async fn refresh(&mut self) -> Result<(), Error> {
78        self.store = self.client.get_file_search_store(&self.store.name).await?;
79        Ok(())
80    }
81
82    #[instrument(skip_all, fields(store.name = %self.store.name, force))]
83    pub async fn delete(self, force: bool) -> Result<(), Error> {
84        self.client
85            .delete_file_search_store(&self.store.name, force)
86            .await
87    }
88
89    pub fn upload(&self, file_data: Vec<u8>) -> UploadBuilder {
90        UploadBuilder {
91            client: self.client.clone(),
92            store_name: self.store.name.clone(),
93            file_data,
94            display_name: None,
95            mime_type: None,
96            custom_metadata: None,
97            chunking_config: None,
98        }
99    }
100
101    pub fn import_file(&self, file_name: String) -> ImportBuilder {
102        ImportBuilder {
103            client: self.client.clone(),
104            store_name: self.store.name.clone(),
105            file_name,
106            custom_metadata: None,
107            chunking_config: None,
108        }
109    }
110
111    pub fn documents(&self) -> DocumentBuilder {
112        DocumentBuilder {
113            client: self.client.clone(),
114            store_name: self.store.name.clone(),
115        }
116    }
117}