1use std::{collections::BTreeMap, path::PathBuf};
2
3use serde::{Deserialize, Deserializer, Serialize, Serializer};
4use validator::Validate;
5
6use super::types::DocumentUploadResponse;
7use crate::client::ZaiClient;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11#[repr(i64)]
12pub enum DocumentSliceType {
13 TitleParagraph = 1,
15 QaPair = 2,
17 Line = 3,
19 Custom = 5,
21 Page = 6,
23 Single = 7,
25}
26impl DocumentSliceType {
27 pub const fn as_i64(self) -> i64 {
29 self as i64
30 }
31}
32
33impl Serialize for DocumentSliceType {
34 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
35 where
36 S: Serializer,
37 {
38 serializer.serialize_i64(self.as_i64())
39 }
40}
41
42impl<'de> Deserialize<'de> for DocumentSliceType {
43 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
44 where
45 D: Deserializer<'de>,
46 {
47 match i64::deserialize(deserializer)? {
48 1 => Ok(Self::TitleParagraph),
49 2 => Ok(Self::QaPair),
50 3 => Ok(Self::Line),
51 5 => Ok(Self::Custom),
52 6 => Ok(Self::Page),
53 7 => Ok(Self::Single),
54 value => Err(serde::de::Error::custom(format!(
55 "unsupported document slice type: {value}"
56 ))),
57 }
58 }
59}
60
61#[derive(Clone, Default, Validate)]
63pub struct DocumentUploadOptions {
64 pub knowledge_type: Option<DocumentSliceType>,
66 pub custom_separator: Option<Vec<String>>,
68 #[validate(range(min = 20, max = 2000))]
70 pub sentence_size: Option<u32>,
71 pub parse_image: Option<bool>,
73 #[validate(url)]
75 pub callback_url: Option<String>,
76 pub callback_header: Option<BTreeMap<String, String>>,
78 pub word_num_limit: Option<String>,
80 #[validate(length(min = 1))]
82 pub req_id: Option<String>,
83}
84
85impl std::fmt::Debug for DocumentUploadOptions {
86 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
87 formatter
88 .debug_struct("DocumentUploadOptions")
89 .field("knowledge_type", &self.knowledge_type)
90 .field(
91 "custom_separator_count",
92 &self.custom_separator.as_ref().map(Vec::len),
93 )
94 .field("sentence_size", &self.sentence_size)
95 .field("parse_image", &self.parse_image)
96 .field("callback_url_configured", &self.callback_url.is_some())
97 .field(
98 "callback_header_entries",
99 &self.callback_header.as_ref().map(BTreeMap::len),
100 )
101 .field("word_num_limit_configured", &self.word_num_limit.is_some())
102 .field("req_id", &self.req_id.as_ref().map(|_| "[REDACTED]"))
103 .finish()
104 }
105}
106
107pub struct DocumentUploadRequest {
112 knowledge_id: String,
113 files: Vec<PathBuf>,
114 options: DocumentUploadOptions,
115}
116
117impl DocumentUploadRequest {
118 pub fn new(knowledge_id: impl Into<String>) -> Self {
120 Self {
121 knowledge_id: knowledge_id.into(),
122 files: Vec::new(),
123 options: DocumentUploadOptions::default(),
124 }
125 }
126
127 pub fn add_file_path(mut self, path: impl Into<PathBuf>) -> Self {
129 self.files.push(path.into());
130 self
131 }
132
133 pub fn with_options(mut self, opts: DocumentUploadOptions) -> Self {
135 self.options = opts;
136 self
137 }
138
139 fn validate_cross(&self) -> crate::ZaiResult<()> {
141 crate::client::validation::require_non_blank(&self.knowledge_id, "knowledge_id")?;
142 if let Some(ref w) = self.options.word_num_limit
143 && (w.is_empty() || !w.chars().all(|c| c.is_ascii_digit()))
144 {
145 return Err(crate::client::validation::invalid(
146 "word_num_limit must be numeric string",
147 ));
148 }
149 if self.files.is_empty() {
150 return Err(crate::client::validation::invalid(
151 "at least one file path must be provided",
152 ));
153 }
154 if self
155 .options
156 .custom_separator
157 .as_ref()
158 .is_some_and(|separators| {
159 separators.is_empty() || separators.iter().any(|value| value.trim().is_empty())
160 })
161 {
162 return Err(crate::client::validation::invalid(
163 "custom_separator must contain at least one non-blank separator",
164 ));
165 }
166 if (self.options.custom_separator.is_some() || self.options.sentence_size.is_some())
167 && self.options.knowledge_type != Some(DocumentSliceType::Custom)
168 {
169 return Err(crate::client::validation::invalid(
170 "custom_separator and sentence_size require knowledge_type=Custom",
171 ));
172 }
173 if self
174 .options
175 .req_id
176 .as_deref()
177 .is_some_and(|request_id| request_id.trim().is_empty())
178 {
179 return Err(crate::client::validation::invalid(
180 "req_id must not be blank",
181 ));
182 }
183 Ok(())
184 }
185
186 pub async fn send_via(&self, client: &ZaiClient) -> crate::ZaiResult<DocumentUploadResponse> {
188 self.options.validate()?;
189 self.validate_cross()?;
190
191 let route = crate::client::routes::DOCUMENTS_UPLOAD;
192 let url = client
193 .endpoints()
194 .resolve_route(route, &[&self.knowledge_id])?;
195 let mut factory = crate::client::transport::multipart::MultipartBodyFactory::new();
196 if let Some(t) = self.options.knowledge_type {
197 factory = factory.field("knowledge_type", t.as_i64().to_string())?;
198 }
199 if let Some(seps) = self.options.custom_separator.as_ref() {
200 factory = factory.field("custom_separator", serde_json::to_string(seps)?)?;
201 }
202 if let Some(sz) = self.options.sentence_size {
203 factory = factory.field("sentence_size", sz.to_string())?;
204 }
205 if let Some(pi) = self.options.parse_image {
206 factory = factory.field("parse_image", pi.to_string())?;
207 }
208 if let Some(url) = self.options.callback_url.as_ref() {
209 factory = factory.field("callback_url", url.as_str())?;
210 }
211 if let Some(header) = self.options.callback_header.as_ref() {
212 factory = factory.field("callback_header", serde_json::to_string(header)?)?;
213 }
214 if let Some(limit) = self.options.word_num_limit.as_ref() {
215 factory = factory.field("word_num_limit", limit.as_str())?;
216 }
217 if let Some(req_id) = self.options.req_id.as_ref() {
218 factory = factory.field("req_id", req_id.as_str())?;
219 }
220 for path in &self.files {
221 let part = crate::client::transport::multipart::FilePart::from_path(path)?;
222 factory = factory.file_named("files", part)?;
223 }
224 client
225 .send_multipart::<DocumentUploadResponse>(route.method(), url, &factory)
226 .await
227 }
228}
229
230#[cfg(test)]
231mod tests {
232 use super::*;
233
234 #[test]
235 fn options_debug_redacts_callbacks_separators_and_request_id() {
236 let options = DocumentUploadOptions {
237 knowledge_type: Some(DocumentSliceType::Custom),
238 custom_separator: Some(vec!["private-separator".to_owned()]),
239 callback_url: Some("https://private.example/callback".to_owned()),
240 callback_header: Some(BTreeMap::from([(
241 "Authorization".to_owned(),
242 "private-token".to_owned(),
243 )])),
244 word_num_limit: Some("private-limit".to_owned()),
245 req_id: Some("private-request".to_owned()),
246 ..DocumentUploadOptions::default()
247 };
248 let debug = format!("{options:?}");
249 for secret in [
250 "private-separator",
251 "private.example",
252 "Authorization",
253 "private-token",
254 "private-limit",
255 "private-request",
256 ] {
257 assert!(!debug.contains(secret));
258 }
259 assert!(debug.contains("callback_header_entries: Some(1)"));
260 }
261
262 #[test]
263 fn whitespace_only_custom_separator_is_rejected() {
264 let request = DocumentUploadRequest::new("knowledge-1")
265 .add_file_path("document.txt")
266 .with_options(DocumentUploadOptions {
267 knowledge_type: Some(DocumentSliceType::Custom),
268 custom_separator: Some(vec![" ".to_owned()]),
269 ..DocumentUploadOptions::default()
270 });
271 assert!(request.validate_cross().is_err());
272 }
273}