Skip to main content

zai_rs/knowledge/
document_upload_file.rs

1use std::{collections::BTreeMap, path::PathBuf, sync::Arc};
2
3use validator::Validate;
4
5use super::types::UploadFileResponse;
6use crate::client::{
7    endpoints::{ApiBase, EndpointConfig, join_url, paths},
8    http::{HttpClient, HttpClientConfig, parse_typed_response, send_multipart_request},
9};
10
11/// Slice type (knowledge_type)
12#[derive(Debug, Clone, Copy)]
13pub enum DocumentSliceType {
14    /// 1: Title-paragraph slicing (txt, doc, pdf, url, docx, ppt, pptx, md)
15    TitleParagraph = 1,
16    /// 2: Q&A slicing (txt, doc, pdf, url, docx, ppt, pptx, md)
17    QaPair = 2,
18    /// 3: Line slicing (xls, xlsx, csv)
19    Line = 3,
20    /// 5: Custom slicing (txt, doc, pdf, url, docx, ppt, pptx, md)
21    Custom = 5,
22    /// 6: Page slicing (pdf, ppt, pptx)
23    Page = 6,
24    /// 7: Single slice (xls, xlsx, csv)
25    Single = 7,
26}
27impl DocumentSliceType {
28    fn as_i64(self) -> i64 {
29        self as i64
30    }
31}
32
33/// Optional parameters for file upload
34#[derive(Debug, Clone, Default, Validate)]
35pub struct UploadFileOptions {
36    /// Document type; if omitted, the server parses dynamically
37    pub knowledge_type: Option<DocumentSliceType>,
38    /// Custom slicing rules; used when knowledge_type = 5
39    pub custom_separator: Option<Vec<String>>,
40    /// Custom slice size; used when knowledge_type = 5; valid range: 20..=2000
41    #[validate(range(min = 20, max = 2000))]
42    pub sentence_size: Option<u32>,
43    /// Whether to parse images
44    pub parse_image: Option<bool>,
45    /// Callback URL
46    #[validate(url)]
47    pub callback_url: Option<String>,
48    /// Callback headers
49    pub callback_header: Option<BTreeMap<String, String>>,
50    /// Document word number limit (must be numeric string per API)
51    pub word_num_limit: Option<String>,
52    /// Request id
53    #[validate(length(min = 1))]
54    pub req_id: Option<String>,
55}
56
57/// File upload request (multipart/form-data)
58pub struct DocumentUploadFileRequest {
59    /// Bearer API key
60    pub key: String,
61    url: String,
62    endpoint_config: EndpointConfig,
63    api_base: ApiBase,
64    knowledge_id: String,
65    http_config: Arc<HttpClientConfig>,
66    files: Vec<PathBuf>,
67    options: UploadFileOptions,
68    _body: (),
69}
70
71impl DocumentUploadFileRequest {
72    /// Create a new request for a specific knowledge base id
73    pub fn new(key: String, knowledge_id: impl AsRef<str>) -> Self {
74        let knowledge_id = knowledge_id.as_ref().to_string();
75        let endpoint_config = EndpointConfig::default();
76        let api_base = ApiBase::LlmApplication;
77        let url = endpoint_config.url(
78            &api_base,
79            &join_url(paths::DOCUMENT_UPLOAD_DOCUMENT, &knowledge_id),
80        );
81        Self {
82            key,
83            url,
84            endpoint_config,
85            api_base,
86            knowledge_id,
87            http_config: Arc::new(HttpClientConfig::default()),
88            files: Vec::new(),
89            options: UploadFileOptions::default(),
90            _body: (),
91        }
92    }
93
94    fn rebuild_url(&mut self) {
95        self.url = self.endpoint_config.url(
96            &self.api_base,
97            &join_url(paths::DOCUMENT_UPLOAD_DOCUMENT, &self.knowledge_id),
98        );
99    }
100
101    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
102        self.api_base = ApiBase::Custom(base_url.into());
103        self.rebuild_url();
104        self
105    }
106
107    pub fn with_endpoint_config(mut self, endpoint_config: EndpointConfig) -> Self {
108        self.endpoint_config = endpoint_config;
109        self.rebuild_url();
110        self
111    }
112
113    pub fn with_http_config(mut self, config: HttpClientConfig) -> Self {
114        self.http_config = Arc::new(config);
115        self
116    }
117
118    /// Add a local file path to upload
119    pub fn add_file_path(mut self, path: impl Into<PathBuf>) -> Self {
120        self.files.push(path.into());
121        self
122    }
123
124    /// Set optional parameters
125    pub fn with_options(mut self, opts: UploadFileOptions) -> Self {
126        self.options = opts;
127        self
128    }
129
130    /// Mutable access to options for incremental configuration
131    pub fn options_mut(&mut self) -> &mut UploadFileOptions {
132        &mut self.options
133    }
134
135    /// Validate cross-field constraints not expressible via `validator`
136    fn validate_cross(&self) -> crate::ZaiResult<()> {
137        // When knowledge_type is Custom (5), sentence_size should be within 20..=2000
138        if let Some(DocumentSliceType::Custom) = self.options.knowledge_type {
139            // sentence_size recommended; API shows default 300; we ensure range if provided
140            if let Some(sz) = self.options.sentence_size
141                && !(20..=2000).contains(&sz)
142            {
143                return Err(crate::client::error::ZaiError::ApiError {
144                    code: 1200,
145                    message: "sentence_size must be 20..=2000 when knowledge_type=Custom (5)"
146                        .to_string(),
147                });
148            }
149        }
150        if let Some(ref w) = self.options.word_num_limit
151            && !w.chars().all(|c| c.is_ascii_digit())
152        {
153            return Err(crate::client::error::ZaiError::ApiError {
154                code: 1200,
155                message: "word_num_limit must be numeric string".to_string(),
156            });
157        }
158        if self.files.is_empty() {
159            return Err(crate::client::error::ZaiError::ApiError {
160                code: 1200,
161                message: "at least one file path must be provided".to_string(),
162            });
163        }
164        Ok(())
165    }
166
167    /// Send multipart request and parse typed response
168    pub async fn send(&self) -> crate::ZaiResult<UploadFileResponse> {
169        // Field validations
170        self.options.validate()?;
171        self.validate_cross()?;
172
173        let resp = self.post().await?;
174        parse_typed_response::<UploadFileResponse>(resp).await
175    }
176}
177
178impl HttpClient for DocumentUploadFileRequest {
179    type Body = (); // unused
180    type ApiUrl = String;
181    type ApiKey = String;
182
183    fn api_url(&self) -> &Self::ApiUrl {
184        &self.url
185    }
186    fn api_key(&self) -> &Self::ApiKey {
187        &self.key
188    }
189    fn body(&self) -> &Self::Body {
190        &self._body
191    }
192
193    // Override POST to send multipart/form-data
194
195    fn post(
196        &self,
197    ) -> impl std::future::Future<Output = crate::ZaiResult<reqwest::Response>> + Send {
198        let url = self.url.clone();
199        let key = self.key.clone();
200        let config = self.http_config.clone();
201        let files = self.files.clone();
202        let opts = self.options.clone();
203        async move {
204            let mut file_parts = Vec::new();
205            for path in files.iter() {
206                let fname = path
207                    .file_name()
208                    .and_then(|s| s.to_str())
209                    .map(|s| s.to_string())
210                    .unwrap_or_else(|| "upload.bin".to_string());
211                file_parts.push((fname, std::fs::read(path)?));
212            }
213            send_multipart_request(reqwest::Method::POST, url, key, config, move || {
214                let mut form = reqwest::multipart::Form::new();
215
216                if let Some(t) = opts.knowledge_type {
217                    form = form.text("knowledge_type", t.as_i64().to_string());
218                }
219                if let Some(seps) = opts.custom_separator.as_ref() {
220                    let s = serde_json::to_string(seps).unwrap_or_else(|_| "[]".to_string());
221                    form = form.text("custom_separator", s);
222                }
223                if let Some(sz) = opts.sentence_size {
224                    form = form.text("sentence_size", sz.to_string());
225                }
226                if let Some(pi) = opts.parse_image {
227                    form = form.text("parse_image", if pi { "true" } else { "false" }.to_string());
228                }
229                if let Some(u) = opts.callback_url.as_ref() {
230                    form = form.text("callback_url", u.clone());
231                }
232                if let Some(h) = opts.callback_header.as_ref() {
233                    let s = serde_json::to_string(h).unwrap_or_else(|_| "{}".to_string());
234                    form = form.text("callback_header", s);
235                }
236                if let Some(w) = opts.word_num_limit.as_ref() {
237                    form = form.text("word_num_limit", w.clone());
238                }
239                if let Some(r) = opts.req_id.as_ref() {
240                    form = form.text("req_id", r.clone());
241                }
242
243                for (fname, bytes) in file_parts.iter() {
244                    let part =
245                        reqwest::multipart::Part::bytes(bytes.clone()).file_name(fname.clone());
246                    form = form.part("files", part);
247                }
248                Ok(form)
249            })
250            .await
251        }
252    }
253
254    fn http_config(&self) -> Arc<HttpClientConfig> {
255        self.http_config.clone()
256    }
257}