gemini_rust/file_search/
document_builder.rs

1use std::sync::Arc;
2
3use async_stream::stream;
4use futures::Stream;
5
6use crate::client::{Error, GeminiClient};
7use crate::file_search::DocumentHandle;
8
9pub struct DocumentBuilder {
10    pub(crate) client: Arc<GeminiClient>,
11    pub(crate) store_name: String,
12}
13
14impl DocumentBuilder {
15    pub async fn get(self, document_id: &str) -> Result<DocumentHandle, Error> {
16        let document = self
17            .client
18            .get_document(&self.store_name, document_id)
19            .await?;
20
21        Ok(DocumentHandle::new(self.client, document))
22    }
23
24    pub async fn delete(self, document_id: &str, force: bool) -> Result<(), Error> {
25        self.client
26            .delete_document(&self.store_name, document_id, force)
27            .await
28    }
29
30    pub fn list(self, page_size: Option<u32>) -> impl Stream<Item = Result<DocumentHandle, Error>> {
31        stream! {
32            let mut page_token: Option<String> = None;
33            loop {
34                let response = self
35                    .client
36                    .list_documents(&self.store_name, page_size, page_token.as_deref())
37                    .await?;
38
39                for document in response.documents {
40                    yield Ok(DocumentHandle::new(self.client.clone(), document));
41                }
42
43                match response.next_page_token {
44                    Some(token) => page_token = Some(token),
45                    None => break,
46                }
47            }
48        }
49    }
50}