use std::sync::Arc;
use serde::{Deserialize, Serialize};
use validator::Validate;
use super::types::UploadUrlResponse;
use crate::{
ZaiResult,
client::{
endpoints::{ApiBase, EndpointConfig, paths},
http::{HttpClient, HttpClientConfig, parse_typed_response},
},
};
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
pub struct UploadUrlDetail {
#[validate(url)]
pub url: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub knowledge_type: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub custom_separator: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1))]
pub sentence_size: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(url)]
pub callback_url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub callback_header: Option<std::collections::BTreeMap<String, String>>,
}
impl UploadUrlDetail {
pub fn new(url: impl Into<String>) -> Self {
Self {
url: url.into(),
knowledge_type: None,
custom_separator: None,
sentence_size: None,
callback_url: None,
callback_header: None,
}
}
pub fn with_knowledge_type(mut self, t: i64) -> Self {
self.knowledge_type = Some(t);
self
}
pub fn with_custom_separator(mut self, seps: Vec<String>) -> Self {
self.custom_separator = Some(seps);
self
}
pub fn with_sentence_size(mut self, size: u32) -> Self {
self.sentence_size = Some(size);
self
}
pub fn with_callback_url(mut self, url: impl Into<String>) -> Self {
self.callback_url = Some(url.into());
self
}
pub fn with_callback_header(
mut self,
headers: std::collections::BTreeMap<String, String>,
) -> Self {
self.callback_header = Some(headers);
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
pub struct UploadUrlBody {
#[validate(length(min = 1))]
pub upload_detail: Vec<UploadUrlDetail>,
#[validate(length(min = 1))]
pub knowledge_id: String,
}
impl UploadUrlBody {
pub fn new(knowledge_id: impl Into<String>) -> Self {
Self {
upload_detail: Vec::new(),
knowledge_id: knowledge_id.into(),
}
}
pub fn add_detail(mut self, detail: UploadUrlDetail) -> Self {
self.upload_detail.push(detail);
self
}
pub fn add_url(mut self, url: impl Into<String>) -> Self {
self.upload_detail.push(UploadUrlDetail::new(url));
self
}
}
pub struct DocumentUploadUrlRequest {
pub key: String,
url: String,
endpoint_config: EndpointConfig,
api_base: ApiBase,
http_config: Arc<HttpClientConfig>,
body: UploadUrlBody,
}
impl DocumentUploadUrlRequest {
pub fn new(key: String, body: UploadUrlBody) -> Self {
let endpoint_config = EndpointConfig::default();
let api_base = ApiBase::LlmApplication;
let url = endpoint_config.url(&api_base, paths::DOCUMENT_UPLOAD_URL);
Self {
key,
url,
endpoint_config,
api_base,
http_config: Arc::new(HttpClientConfig::default()),
body,
}
}
fn rebuild_url(&mut self) {
self.url = self
.endpoint_config
.url(&self.api_base, paths::DOCUMENT_UPLOAD_URL);
}
pub fn with_base_url(mut self, base: impl Into<String>) -> Self {
self.api_base = ApiBase::Custom(base.into());
self.rebuild_url();
self
}
pub fn with_endpoint_config(mut self, endpoint_config: EndpointConfig) -> Self {
self.endpoint_config = endpoint_config;
self.rebuild_url();
self
}
pub fn with_http_config(mut self, config: HttpClientConfig) -> Self {
self.http_config = Arc::new(config);
self
}
pub fn body_mut(&mut self) -> &mut UploadUrlBody {
&mut self.body
}
pub async fn send(&self) -> ZaiResult<UploadUrlResponse> {
self.body.validate()?;
let resp = self.post().await?;
let parsed = parse_typed_response::<UploadUrlResponse>(resp).await?;
Ok(parsed)
}
}
impl HttpClient for DocumentUploadUrlRequest {
type Body = UploadUrlBody;
type ApiUrl = String;
type ApiKey = String;
fn api_url(&self) -> &Self::ApiUrl {
&self.url
}
fn api_key(&self) -> &Self::ApiKey {
&self.key
}
fn body(&self) -> &Self::Body {
&self.body
}
fn http_config(&self) -> Arc<HttpClientConfig> {
Arc::clone(&self.http_config)
}
}