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) const fn new(inner: Arc<ClientInner>) -> Self {
21        Self { inner }
22    }
23
24    /// 获取 Document。
25    ///
26    /// # Errors
27    /// 当请求失败或响应解析失败时返回错误。
28    pub async fn get(&self, name: impl AsRef<str>) -> Result<Document> {
29        self.get_with_config(name, GetDocumentConfig::default())
30            .await
31    }
32
33    /// 获取 Document（带配置）。
34    ///
35    /// # Errors
36    /// 当请求失败或响应解析失败时返回错误。
37    pub async fn get_with_config(
38        &self,
39        name: impl AsRef<str>,
40        mut config: GetDocumentConfig,
41    ) -> Result<Document> {
42        ensure_gemini_backend(&self.inner)?;
43        let http_options = config.http_options.take();
44        let name = normalize_document_name(name.as_ref())?;
45        let url = build_document_url(&self.inner, &name, http_options.as_ref());
46        let mut request = self.inner.http.get(url);
47        request = apply_http_options(request, http_options.as_ref())?;
48
49        let response = self.inner.send(request).await?;
50        if !response.status().is_success() {
51            return Err(Error::ApiError {
52                status: response.status().as_u16(),
53                message: response.text().await.unwrap_or_default(),
54            });
55        }
56        Ok(response.json::<Document>().await?)
57    }
58
59    /// 删除 Document。
60    ///
61    /// # Errors
62    /// 当请求失败或响应解析失败时返回错误。
63    pub async fn delete(&self, name: impl AsRef<str>) -> Result<()> {
64        self.delete_with_config(name, DeleteDocumentConfig::default())
65            .await
66    }
67
68    /// 删除 Document（带配置）。
69    ///
70    /// # Errors
71    /// 当请求失败或响应解析失败时返回错误。
72    pub async fn delete_with_config(
73        &self,
74        name: impl AsRef<str>,
75        mut config: DeleteDocumentConfig,
76    ) -> Result<()> {
77        ensure_gemini_backend(&self.inner)?;
78        let http_options = config.http_options.take();
79        let name = normalize_document_name(name.as_ref())?;
80        let url = build_document_url(&self.inner, &name, http_options.as_ref());
81        let url = add_delete_query_params(&url, &config)?;
82        let mut request = self.inner.http.delete(url);
83        request = apply_http_options(request, http_options.as_ref())?;
84
85        let response = self.inner.send(request).await?;
86        if !response.status().is_success() {
87            return Err(Error::ApiError {
88                status: response.status().as_u16(),
89                message: response.text().await.unwrap_or_default(),
90            });
91        }
92        Ok(())
93    }
94
95    /// 列出 Documents。
96    ///
97    /// # Errors
98    /// 当请求失败或响应解析失败时返回错误。
99    pub async fn list(&self, parent: impl AsRef<str>) -> Result<ListDocumentsResponse> {
100        self.list_with_config(parent, ListDocumentsConfig::default())
101            .await
102    }
103
104    /// 列出 Documents（带配置）。
105    ///
106    /// # Errors
107    /// 当请求失败或响应解析失败时返回错误。
108    pub async fn list_with_config(
109        &self,
110        parent: impl AsRef<str>,
111        mut config: ListDocumentsConfig,
112    ) -> Result<ListDocumentsResponse> {
113        ensure_gemini_backend(&self.inner)?;
114        let http_options = config.http_options.take();
115        let parent = normalize_file_search_store_name(parent.as_ref());
116        let url = build_documents_url(&self.inner, &parent, http_options.as_ref());
117        let url = add_list_query_params(&url, &config)?;
118        let mut request = self.inner.http.get(url);
119        request = apply_http_options(request, http_options.as_ref())?;
120
121        let response = self.inner.send(request).await?;
122        if !response.status().is_success() {
123            return Err(Error::ApiError {
124                status: response.status().as_u16(),
125                message: response.text().await.unwrap_or_default(),
126            });
127        }
128        Ok(response.json::<ListDocumentsResponse>().await?)
129    }
130
131    /// 列出所有 Documents（自动翻页）。
132    ///
133    /// # Errors
134    /// 当请求失败或响应解析失败时返回错误。
135    pub async fn all(&self, parent: impl AsRef<str>) -> Result<Vec<Document>> {
136        self.all_with_config(parent, ListDocumentsConfig::default())
137            .await
138    }
139
140    /// 列出所有 Documents（带配置，自动翻页）。
141    ///
142    /// # Errors
143    /// 当请求失败或响应解析失败时返回错误。
144    pub async fn all_with_config(
145        &self,
146        parent: impl AsRef<str>,
147        mut config: ListDocumentsConfig,
148    ) -> Result<Vec<Document>> {
149        let mut docs = Vec::new();
150        let http_options = config.http_options.clone();
151        loop {
152            let mut page_config = config.clone();
153            page_config.http_options.clone_from(&http_options);
154            let response = self.list_with_config(parent.as_ref(), page_config).await?;
155            if let Some(items) = response.documents {
156                docs.extend(items);
157            }
158            match response.next_page_token {
159                Some(token) if !token.is_empty() => {
160                    config.page_token = Some(token);
161                }
162                _ => break,
163            }
164        }
165        Ok(docs)
166    }
167}
168
169fn ensure_gemini_backend(inner: &ClientInner) -> Result<()> {
170    if inner.config.backend == Backend::VertexAi {
171        return Err(Error::InvalidConfig {
172            message: "Documents API is only supported in Gemini API".into(),
173        });
174    }
175    Ok(())
176}
177
178fn normalize_file_search_store_name(name: &str) -> String {
179    if name.starts_with("fileSearchStores/") {
180        name.to_string()
181    } else {
182        format!("fileSearchStores/{name}")
183    }
184}
185
186fn normalize_document_name(name: &str) -> Result<String> {
187    if name.contains("/documents/") {
188        Ok(name.to_string())
189    } else {
190        Err(Error::InvalidConfig {
191            message: format!(
192                "Document name must be a full resource name, e.g. fileSearchStores/xxx/documents/yyy (got {name})"
193            ),
194        })
195    }
196}
197
198fn build_document_url(
199    inner: &ClientInner,
200    name: &str,
201    http_options: Option<&rust_genai_types::http::HttpOptions>,
202) -> String {
203    let base = http_options
204        .and_then(|opts| opts.base_url.as_deref())
205        .unwrap_or(&inner.api_client.base_url);
206    let version = http_options
207        .and_then(|opts| opts.api_version.as_deref())
208        .unwrap_or(&inner.api_client.api_version);
209    format!("{base}{version}/{name}")
210}
211
212fn build_documents_url(
213    inner: &ClientInner,
214    parent: &str,
215    http_options: Option<&rust_genai_types::http::HttpOptions>,
216) -> String {
217    let base = http_options
218        .and_then(|opts| opts.base_url.as_deref())
219        .unwrap_or(&inner.api_client.base_url);
220    let version = http_options
221        .and_then(|opts| opts.api_version.as_deref())
222        .unwrap_or(&inner.api_client.api_version);
223    format!("{base}{version}/{parent}/documents")
224}
225
226fn add_list_query_params(url: &str, config: &ListDocumentsConfig) -> Result<String> {
227    let mut url = reqwest::Url::parse(url).map_err(|err| Error::InvalidConfig {
228        message: err.to_string(),
229    })?;
230    {
231        let mut pairs = url.query_pairs_mut();
232        if let Some(page_size) = config.page_size {
233            pairs.append_pair("pageSize", &page_size.to_string());
234        }
235        if let Some(page_token) = &config.page_token {
236            pairs.append_pair("pageToken", page_token);
237        }
238    }
239    Ok(url.to_string())
240}
241
242fn add_delete_query_params(url: &str, config: &DeleteDocumentConfig) -> Result<String> {
243    let mut url = reqwest::Url::parse(url).map_err(|err| Error::InvalidConfig {
244        message: err.to_string(),
245    })?;
246    {
247        let mut pairs = url.query_pairs_mut();
248        if let Some(force) = config.force {
249            pairs.append_pair("force", &force.to_string());
250        }
251    }
252    Ok(url.to_string())
253}
254
255fn apply_http_options(
256    mut request: reqwest::RequestBuilder,
257    http_options: Option<&rust_genai_types::http::HttpOptions>,
258) -> Result<reqwest::RequestBuilder> {
259    if let Some(options) = http_options {
260        if let Some(timeout) = options.timeout {
261            request = request.timeout(Duration::from_millis(timeout));
262        }
263        if let Some(headers) = &options.headers {
264            for (key, value) in headers {
265                let name =
266                    HeaderName::from_bytes(key.as_bytes()).map_err(|_| Error::InvalidConfig {
267                        message: format!("Invalid header name: {key}"),
268                    })?;
269                let value = HeaderValue::from_str(value).map_err(|_| Error::InvalidConfig {
270                    message: format!("Invalid header value for {key}"),
271                })?;
272                request = request.header(name, value);
273            }
274        }
275    }
276    Ok(request)
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use crate::test_support::test_client_inner;
283    use std::collections::HashMap;
284
285    #[test]
286    fn test_normalize_document_names_and_urls() {
287        assert_eq!(
288            normalize_file_search_store_name("store"),
289            "fileSearchStores/store"
290        );
291        assert!(normalize_document_name("fileSearchStores/store/documents/doc1").is_ok());
292        assert!(normalize_document_name("invalid").is_err());
293
294        let gemini = test_client_inner(Backend::GeminiApi);
295        let url = build_documents_url(&gemini, "fileSearchStores/store", None);
296        assert!(url.contains("/v1beta/fileSearchStores/store/documents"));
297
298        let url = build_document_url(&gemini, "fileSearchStores/store/documents/doc1", None);
299        assert!(url.ends_with("/v1beta/fileSearchStores/store/documents/doc1"));
300    }
301
302    #[test]
303    fn test_query_params_and_backend_check() {
304        let url = add_list_query_params(
305            "https://example.com/docs",
306            &ListDocumentsConfig {
307                page_size: Some(3),
308                page_token: Some("t".to_string()),
309                ..Default::default()
310            },
311        )
312        .unwrap();
313        assert!(url.contains("pageSize=3"));
314        assert!(url.contains("pageToken=t"));
315
316        let url = add_delete_query_params(
317            "https://example.com/docs/1",
318            &DeleteDocumentConfig {
319                force: Some(true),
320                ..Default::default()
321            },
322        )
323        .unwrap();
324        assert!(url.contains("force=true"));
325
326        let vertex = test_client_inner(Backend::VertexAi);
327        let err = ensure_gemini_backend(&vertex).unwrap_err();
328        assert!(matches!(err, Error::InvalidConfig { .. }));
329    }
330
331    #[test]
332    fn test_apply_http_options_invalid_header_value() {
333        let client = reqwest::Client::new();
334        let request = client.get("https://example.com");
335        let mut headers = HashMap::new();
336        headers.insert("x-test".to_string(), "bad\nvalue".to_string());
337        let options = rust_genai_types::http::HttpOptions {
338            headers: Some(headers),
339            ..Default::default()
340        };
341        let err = apply_http_options(request, Some(&options)).unwrap_err();
342        assert!(matches!(err, Error::InvalidConfig { .. }));
343    }
344
345    #[test]
346    fn test_apply_http_options_success_path() {
347        let client = reqwest::Client::new();
348        let request = client.get("https://example.com");
349        let mut headers = HashMap::new();
350        headers.insert("x-ok".to_string(), "ok".to_string());
351        let options = rust_genai_types::http::HttpOptions {
352            headers: Some(headers),
353            timeout: Some(1500),
354            ..Default::default()
355        };
356        let request = apply_http_options(request, Some(&options)).unwrap();
357        let built = request.build().unwrap();
358        assert_eq!(built.headers().get("x-ok").unwrap(), "ok");
359    }
360
361    #[test]
362    fn test_add_query_params_invalid_url_and_header_name() {
363        let err = add_list_query_params("://bad", &ListDocumentsConfig::default()).unwrap_err();
364        assert!(matches!(err, Error::InvalidConfig { .. }));
365        let err = add_delete_query_params("://bad", &DeleteDocumentConfig::default()).unwrap_err();
366        assert!(matches!(err, Error::InvalidConfig { .. }));
367
368        let client = reqwest::Client::new();
369        let request = client.get("https://example.com");
370        let mut headers = HashMap::new();
371        headers.insert("bad header".to_string(), "ok".to_string());
372        let options = rust_genai_types::http::HttpOptions {
373            headers: Some(headers),
374            ..Default::default()
375        };
376        let err = apply_http_options(request, Some(&options)).unwrap_err();
377        assert!(matches!(err, Error::InvalidConfig { .. }));
378    }
379
380    #[test]
381    fn test_normalize_file_search_store_name_with_prefix() {
382        assert_eq!(
383            normalize_file_search_store_name("fileSearchStores/store"),
384            "fileSearchStores/store"
385        );
386    }
387}
rust_genai/documents.rs

rust_genai/
documents.rs
