use std::{collections::BTreeMap, path::PathBuf};
use validator::Validate;
use super::types::UploadFileResponse;
use crate::client::ZaiClient;
#[derive(Debug, Clone, Copy)]
pub enum DocumentSliceType {
TitleParagraph = 1,
QaPair = 2,
Line = 3,
Custom = 5,
Page = 6,
Single = 7,
}
impl DocumentSliceType {
fn as_i64(self) -> i64 {
self as i64
}
}
#[derive(Debug, Clone, Default, Validate)]
pub struct UploadFileOptions {
pub knowledge_type: Option<DocumentSliceType>,
pub custom_separator: Option<Vec<String>>,
#[validate(range(min = 20, max = 2000))]
pub sentence_size: Option<u32>,
pub parse_image: Option<bool>,
#[validate(url)]
pub callback_url: Option<String>,
pub callback_header: Option<BTreeMap<String, String>>,
pub word_num_limit: Option<String>,
#[validate(length(min = 1))]
pub req_id: Option<String>,
}
pub struct DocumentUploadFileRequest {
knowledge_id: String,
files: Vec<PathBuf>,
options: UploadFileOptions,
}
impl DocumentUploadFileRequest {
pub fn new(knowledge_id: impl Into<String>) -> Self {
Self {
knowledge_id: knowledge_id.into(),
files: Vec::new(),
options: UploadFileOptions::default(),
}
}
pub fn add_file_path(mut self, path: impl Into<PathBuf>) -> Self {
self.files.push(path.into());
self
}
pub fn with_options(mut self, opts: UploadFileOptions) -> Self {
self.options = opts;
self
}
pub fn options_mut(&mut self) -> &mut UploadFileOptions {
&mut self.options
}
fn validate_cross(&self) -> crate::ZaiResult<()> {
if let Some(DocumentSliceType::Custom) = self.options.knowledge_type {
if let Some(sz) = self.options.sentence_size
&& !(20..=2000).contains(&sz)
{
return Err(crate::client::error::ZaiError::ApiError {
code: crate::client::error::codes::SDK_VALIDATION,
message: "sentence_size must be 20..=2000 when knowledge_type=Custom (5)"
.to_string(),
});
}
}
if let Some(ref w) = self.options.word_num_limit
&& !w.chars().all(|c| c.is_ascii_digit())
{
return Err(crate::client::error::ZaiError::ApiError {
code: crate::client::error::codes::SDK_VALIDATION,
message: "word_num_limit must be numeric string".to_string(),
});
}
if self.files.is_empty() {
return Err(crate::client::error::ZaiError::ApiError {
code: crate::client::error::codes::SDK_VALIDATION,
message: "at least one file path must be provided".to_string(),
});
}
Ok(())
}
pub async fn send_via(&self, client: &ZaiClient) -> crate::ZaiResult<UploadFileResponse> {
self.options.validate()?;
self.validate_cross()?;
let route = crate::client::routes::DOCUMENTS_UPLOAD;
let url = client
.endpoints()
.resolve_route(route, &[&self.knowledge_id])?;
let mut factory = crate::client::transport::multipart::MultipartBodyFactory::new();
if let Some(t) = self.options.knowledge_type {
factory = factory.field("knowledge_type", t.as_i64().to_string())?;
}
if let Some(seps) = self.options.custom_separator.as_ref() {
factory = factory.field("custom_separator", serde_json::to_string(seps)?)?;
}
if let Some(sz) = self.options.sentence_size {
factory = factory.field("sentence_size", sz.to_string())?;
}
if let Some(pi) = self.options.parse_image {
factory = factory.field("parse_image", pi.to_string())?;
}
if let Some(url) = self.options.callback_url.as_ref() {
factory = factory.field("callback_url", url.clone())?;
}
if let Some(header) = self.options.callback_header.as_ref() {
factory = factory.field("callback_header", serde_json::to_string(header)?)?;
}
if let Some(limit) = self.options.word_num_limit.as_ref() {
factory = factory.field("word_num_limit", limit.clone())?;
}
if let Some(req_id) = self.options.req_id.as_ref() {
factory = factory.field("req_id", req_id.clone())?;
}
for path in &self.files {
let part = crate::client::transport::multipart::FilePart::from_path(path)?;
factory = factory.file_named("files", part)?;
}
client
.send_multipart::<UploadFileResponse>(route.method(), url, &factory)
.await
}
}