rust_genai/
documents.rs

1//! Documents API surface.
2
3use std::sync::Arc;
4use std::time::Duration;
5
6use reqwest::header::{HeaderName, HeaderValue};
7use rust_genai_types::documents::{
8    DeleteDocumentConfig, Document, GetDocumentConfig, ListDocumentsConfig, ListDocumentsResponse,
9};
10
11use crate::client::{Backend, ClientInner};
12use crate::error::{Error, Result};
13
14#[derive(Clone)]
15pub struct Documents {
16    pub(crate) inner: Arc<ClientInner>,
17}
18
19impl Documents {
20    pub(crate) fn new(inner: Arc<ClientInner>) -> Self {
21        Self { inner }
22    }
23
24    /// 获取 Document。
25    pub async fn get(&self, name: impl AsRef<str>) -> Result<Document> {
26        self.get_with_config(name, GetDocumentConfig::default())
27            .await
28    }
29
30    /// 获取 Document（带配置）。
31    pub async fn get_with_config(
32        &self,
33        name: impl AsRef<str>,
34        mut config: GetDocumentConfig,
35    ) -> Result<Document> {
36        ensure_gemini_backend(&self.inner)?;
37        let http_options = config.http_options.take();
38        let name = normalize_document_name(name.as_ref())?;
39        let url = build_document_url(&self.inner, &name, http_options.as_ref())?;
40        let mut request = self.inner.http.get(url);
41        request = apply_http_options(request, http_options.as_ref())?;
42
43        let response = self.inner.send(request).await?;
44        if !response.status().is_success() {
45            return Err(Error::ApiError {
46                status: response.status().as_u16(),
47                message: response.text().await.unwrap_or_default(),
48            });
49        }
50        Ok(response.json::<Document>().await?)
51    }
52
53    /// 删除 Document。
54    pub async fn delete(&self, name: impl AsRef<str>) -> Result<()> {
55        self.delete_with_config(name, DeleteDocumentConfig::default())
56            .await
57    }
58
59    /// 删除 Document（带配置）。
60    pub async fn delete_with_config(
61        &self,
62        name: impl AsRef<str>,
63        mut config: DeleteDocumentConfig,
64    ) -> Result<()> {
65        ensure_gemini_backend(&self.inner)?;
66        let http_options = config.http_options.take();
67        let name = normalize_document_name(name.as_ref())?;
68        let url = build_document_url(&self.inner, &name, http_options.as_ref())?;
69        let url = add_delete_query_params(url, &config)?;
70        let mut request = self.inner.http.delete(url);
71        request = apply_http_options(request, http_options.as_ref())?;
72
73        let response = self.inner.send(request).await?;
74        if !response.status().is_success() {
75            return Err(Error::ApiError {
76                status: response.status().as_u16(),
77                message: response.text().await.unwrap_or_default(),
78            });
79        }
80        Ok(())
81    }
82
83    /// 列出 Documents。
84    pub async fn list(&self, parent: impl AsRef<str>) -> Result<ListDocumentsResponse> {
85        self.list_with_config(parent, ListDocumentsConfig::default())
86            .await
87    }
88
89    /// 列出 Documents（带配置）。
90    pub async fn list_with_config(
91        &self,
92        parent: impl AsRef<str>,
93        mut config: ListDocumentsConfig,
94    ) -> Result<ListDocumentsResponse> {
95        ensure_gemini_backend(&self.inner)?;
96        let http_options = config.http_options.take();
97        let parent = normalize_file_search_store_name(parent.as_ref());
98        let url = build_documents_url(&self.inner, &parent, http_options.as_ref())?;
99        let url = add_list_query_params(url, &config)?;
100        let mut request = self.inner.http.get(url);
101        request = apply_http_options(request, http_options.as_ref())?;
102
103        let response = self.inner.send(request).await?;
104        if !response.status().is_success() {
105            return Err(Error::ApiError {
106                status: response.status().as_u16(),
107                message: response.text().await.unwrap_or_default(),
108            });
109        }
110        Ok(response.json::<ListDocumentsResponse>().await?)
111    }
112
113    /// 列出所有 Documents（自动翻页）。
114    pub async fn all(&self, parent: impl AsRef<str>) -> Result<Vec<Document>> {
115        self.all_with_config(parent, ListDocumentsConfig::default())
116            .await
117    }
118
119    /// 列出所有 Documents（带配置，自动翻页）。
120    pub async fn all_with_config(
121        &self,
122        parent: impl AsRef<str>,
123        mut config: ListDocumentsConfig,
124    ) -> Result<Vec<Document>> {
125        let mut docs = Vec::new();
126        let http_options = config.http_options.clone();
127        loop {
128            let mut page_config = config.clone();
129            page_config.http_options = http_options.clone();
130            let response = self.list_with_config(parent.as_ref(), page_config).await?;
131            if let Some(items) = response.documents {
132                docs.extend(items);
133            }
134            match response.next_page_token {
135                Some(token) if !token.is_empty() => {
136                    config.page_token = Some(token);
137                }
138                _ => break,
139            }
140        }
141        Ok(docs)
142    }
143}
144
145fn ensure_gemini_backend(inner: &ClientInner) -> Result<()> {
146    if inner.config.backend == Backend::VertexAi {
147        return Err(Error::InvalidConfig {
148            message: "Documents API is only supported in Gemini API".into(),
149        });
150    }
151    Ok(())
152}
153
154fn normalize_file_search_store_name(name: &str) -> String {
155    if name.starts_with("fileSearchStores/") {
156        name.to_string()
157    } else {
158        format!("fileSearchStores/{name}")
159    }
160}
161
162fn normalize_document_name(name: &str) -> Result<String> {
163    if name.contains("/documents/") {
164        Ok(name.to_string())
165    } else {
166        Err(Error::InvalidConfig {
167            message: format!(
168                "Document name must be a full resource name, e.g. fileSearchStores/xxx/documents/yyy (got {name})"
169            ),
170        })
171    }
172}
173
174fn build_document_url(
175    inner: &ClientInner,
176    name: &str,
177    http_options: Option<&rust_genai_types::http::HttpOptions>,
178) -> Result<String> {
179    let base = http_options
180        .and_then(|opts| opts.base_url.as_deref())
181        .unwrap_or(&inner.api_client.base_url);
182    let version = http_options
183        .and_then(|opts| opts.api_version.as_deref())
184        .unwrap_or(&inner.api_client.api_version);
185    Ok(format!("{base}{version}/{name}"))
186}
187
188fn build_documents_url(
189    inner: &ClientInner,
190    parent: &str,
191    http_options: Option<&rust_genai_types::http::HttpOptions>,
192) -> Result<String> {
193    let base = http_options
194        .and_then(|opts| opts.base_url.as_deref())
195        .unwrap_or(&inner.api_client.base_url);
196    let version = http_options
197        .and_then(|opts| opts.api_version.as_deref())
198        .unwrap_or(&inner.api_client.api_version);
199    Ok(format!("{base}{version}/{parent}/documents"))
200}
201
202fn add_list_query_params(url: String, config: &ListDocumentsConfig) -> Result<String> {
203    let mut url = reqwest::Url::parse(&url).map_err(|err| Error::InvalidConfig {
204        message: err.to_string(),
205    })?;
206    {
207        let mut pairs = url.query_pairs_mut();
208        if let Some(page_size) = config.page_size {
209            pairs.append_pair("pageSize", &page_size.to_string());
210        }
211        if let Some(page_token) = &config.page_token {
212            pairs.append_pair("pageToken", page_token);
213        }
214    }
215    Ok(url.to_string())
216}
217
218fn add_delete_query_params(url: String, config: &DeleteDocumentConfig) -> Result<String> {
219    let mut url = reqwest::Url::parse(&url).map_err(|err| Error::InvalidConfig {
220        message: err.to_string(),
221    })?;
222    {
223        let mut pairs = url.query_pairs_mut();
224        if let Some(force) = config.force {
225            pairs.append_pair("force", &force.to_string());
226        }
227    }
228    Ok(url.to_string())
229}
230
231fn apply_http_options(
232    mut request: reqwest::RequestBuilder,
233    http_options: Option<&rust_genai_types::http::HttpOptions>,
234) -> Result<reqwest::RequestBuilder> {
235    if let Some(options) = http_options {
236        if let Some(timeout) = options.timeout {
237            request = request.timeout(Duration::from_millis(timeout));
238        }
239        if let Some(headers) = &options.headers {
240            for (key, value) in headers {
241                let name =
242                    HeaderName::from_bytes(key.as_bytes()).map_err(|_| Error::InvalidConfig {
243                        message: format!("Invalid header name: {key}"),
244                    })?;
245                let value = HeaderValue::from_str(value).map_err(|_| Error::InvalidConfig {
246                    message: format!("Invalid header value for {key}"),
247                })?;
248                request = request.header(name, value);
249            }
250        }
251    }
252    Ok(request)
253}
rust_genai/documents.rs

rust_genai/
documents.rs
