Skip to main content

zai_rs/knowledge/
document_upload_url.rs

1use std::sync::Arc;
2
3use serde::{Deserialize, Serialize};
4use validator::Validate;
5
6use super::types::UploadUrlResponse;
7use crate::{
8    ZaiResult,
9    client::{
10        endpoints::{ApiBase, EndpointConfig, paths},
11        http::{HttpClient, HttpClientConfig, parse_typed_response},
12    },
13};
14
15/// Single URL upload detail
16#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
17pub struct UploadUrlDetail {
18    /// Source URL to fetch
19    #[validate(url)]
20    pub url: String,
21    /// Slice type (integer)
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub knowledge_type: Option<i64>,
24    /// Custom separators
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub custom_separator: Option<Vec<String>>,
27    /// Sentence size
28    #[serde(skip_serializing_if = "Option::is_none")]
29    #[validate(range(min = 1))]
30    pub sentence_size: Option<u32>,
31    /// Callback URL
32    #[serde(skip_serializing_if = "Option::is_none")]
33    #[validate(url)]
34    pub callback_url: Option<String>,
35    /// Callback headers
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub callback_header: Option<std::collections::BTreeMap<String, String>>,
38}
39
40impl UploadUrlDetail {
41    pub fn new(url: impl Into<String>) -> Self {
42        Self {
43            url: url.into(),
44            knowledge_type: None,
45            custom_separator: None,
46            sentence_size: None,
47            callback_url: None,
48            callback_header: None,
49        }
50    }
51    pub fn with_knowledge_type(mut self, t: i64) -> Self {
52        self.knowledge_type = Some(t);
53        self
54    }
55    pub fn with_custom_separator(mut self, seps: Vec<String>) -> Self {
56        self.custom_separator = Some(seps);
57        self
58    }
59    pub fn with_sentence_size(mut self, size: u32) -> Self {
60        self.sentence_size = Some(size);
61        self
62    }
63    pub fn with_callback_url(mut self, url: impl Into<String>) -> Self {
64        self.callback_url = Some(url.into());
65        self
66    }
67    pub fn with_callback_header(
68        mut self,
69        headers: std::collections::BTreeMap<String, String>,
70    ) -> Self {
71        self.callback_header = Some(headers);
72        self
73    }
74}
75
76/// Upload URL request body
77#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
78pub struct UploadUrlBody {
79    /// Upload detail list (at least 1)
80    #[validate(length(min = 1))]
81    pub upload_detail: Vec<UploadUrlDetail>,
82    /// Knowledge base id
83    #[validate(length(min = 1))]
84    pub knowledge_id: String,
85}
86
87impl UploadUrlBody {
88    pub fn new(knowledge_id: impl Into<String>) -> Self {
89        Self {
90            upload_detail: Vec::new(),
91            knowledge_id: knowledge_id.into(),
92        }
93    }
94    pub fn add_detail(mut self, detail: UploadUrlDetail) -> Self {
95        self.upload_detail.push(detail);
96        self
97    }
98    pub fn add_url(mut self, url: impl Into<String>) -> Self {
99        self.upload_detail.push(UploadUrlDetail::new(url));
100        self
101    }
102}
103
104/// Upload URL request (POST /llm-application/open/document/upload_url)
105pub struct DocumentUploadUrlRequest {
106    /// Bearer API key
107    pub key: String,
108    url: String,
109    endpoint_config: EndpointConfig,
110    api_base: ApiBase,
111    http_config: Arc<HttpClientConfig>,
112    body: UploadUrlBody,
113}
114
115impl DocumentUploadUrlRequest {
116    pub fn new(key: String, body: UploadUrlBody) -> Self {
117        let endpoint_config = EndpointConfig::default();
118        let api_base = ApiBase::LlmApplication;
119        let url = endpoint_config.url(&api_base, paths::DOCUMENT_UPLOAD_URL);
120        Self {
121            key,
122            url,
123            endpoint_config,
124            api_base,
125            http_config: Arc::new(HttpClientConfig::default()),
126            body,
127        }
128    }
129
130    fn rebuild_url(&mut self) {
131        self.url = self
132            .endpoint_config
133            .url(&self.api_base, paths::DOCUMENT_UPLOAD_URL);
134    }
135
136    pub fn with_base_url(mut self, base: impl Into<String>) -> Self {
137        self.api_base = ApiBase::Custom(base.into());
138        self.rebuild_url();
139        self
140    }
141
142    pub fn with_endpoint_config(mut self, endpoint_config: EndpointConfig) -> Self {
143        self.endpoint_config = endpoint_config;
144        self.rebuild_url();
145        self
146    }
147
148    pub fn with_http_config(mut self, config: HttpClientConfig) -> Self {
149        self.http_config = Arc::new(config);
150        self
151    }
152
153    pub fn body_mut(&mut self) -> &mut UploadUrlBody {
154        &mut self.body
155    }
156
157    /// Validate and send
158    pub async fn send(&self) -> ZaiResult<UploadUrlResponse> {
159        self.body.validate()?;
160        let resp = self.post().await?;
161        let parsed = parse_typed_response::<UploadUrlResponse>(resp).await?;
162        Ok(parsed)
163    }
164}
165
166impl HttpClient for DocumentUploadUrlRequest {
167    type Body = UploadUrlBody;
168    type ApiUrl = String;
169    type ApiKey = String;
170
171    fn api_url(&self) -> &Self::ApiUrl {
172        &self.url
173    }
174    fn api_key(&self) -> &Self::ApiKey {
175        &self.key
176    }
177    fn body(&self) -> &Self::Body {
178        &self.body
179    }
180
181    fn http_config(&self) -> Arc<HttpClientConfig> {
182        Arc::clone(&self.http_config)
183    }
184}