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

rust_genai/
documents.rs
