Skip to main content

zai_rs/knowledge/
document_reembedding.rs

1use std::{collections::BTreeMap, sync::Arc};
2
3use serde::{Deserialize, Serialize};
4use validator::Validate;
5
6use crate::{
7    ZaiResult,
8    client::{
9        endpoints::{ApiBase, EndpointConfig, join_url, paths},
10        http::{HttpClient, HttpClientConfig, parse_typed_response},
11    },
12};
13
14/// Request body for re-embedding a document
15#[derive(Debug, Clone, Serialize, Deserialize, Validate, Default)]
16pub struct DocumentReembeddingBody {
17    /// Optional callback URL that will be called when re-embedding completes
18    #[serde(skip_serializing_if = "Option::is_none")]
19    #[validate(url)]
20    pub callback_url: Option<String>,
21    /// Optional callback headers key-value pairs
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub callback_header: Option<BTreeMap<String, String>>,
24}
25
26/// Re-embedding request (POST /llm-application/open/document/embedding/{id})
27pub struct DocumentReembeddingRequest {
28    /// Bearer API key
29    pub key: String,
30    url: String,
31    endpoint_config: EndpointConfig,
32    api_base: ApiBase,
33    document_id: String,
34    http_config: Arc<HttpClientConfig>,
35    body: DocumentReembeddingBody,
36}
37
38impl DocumentReembeddingRequest {
39    /// Create a new request for the specified document id
40    pub fn new(key: String, document_id: impl AsRef<str>) -> Self {
41        let document_id = document_id.as_ref().to_string();
42        let endpoint_config = EndpointConfig::default();
43        let api_base = ApiBase::LlmApplication;
44        let url = endpoint_config.url(
45            &api_base,
46            &join_url(paths::DOCUMENT_EMBEDDING, &document_id),
47        );
48        Self {
49            key,
50            url,
51            endpoint_config,
52            api_base,
53            document_id,
54            http_config: Arc::new(HttpClientConfig::default()),
55            body: DocumentReembeddingBody::default(),
56        }
57    }
58
59    fn rebuild_url(&mut self) {
60        self.url = self.endpoint_config.url(
61            &self.api_base,
62            &join_url(paths::DOCUMENT_EMBEDDING, &self.document_id),
63        );
64    }
65
66    pub fn with_base_url(mut self, base: impl Into<String>) -> Self {
67        self.api_base = ApiBase::Custom(base.into());
68        self.rebuild_url();
69        self
70    }
71
72    pub fn with_endpoint_config(mut self, endpoint_config: EndpointConfig) -> Self {
73        self.endpoint_config = endpoint_config;
74        self.rebuild_url();
75        self
76    }
77
78    pub fn with_http_config(mut self, config: HttpClientConfig) -> Self {
79        self.http_config = Arc::new(config);
80        self
81    }
82
83    /// Set callback URL
84    pub fn with_callback_url(mut self, url: impl Into<String>) -> Self {
85        self.body.callback_url = Some(url.into());
86        self
87    }
88
89    /// Set callback headers
90    pub fn with_callback_header(mut self, hdr: BTreeMap<String, String>) -> Self {
91        self.body.callback_header = Some(hdr);
92        self
93    }
94
95    /// Send POST request with JSON body and parse typed response
96    pub async fn send(&self) -> ZaiResult<DocumentReembeddingResponse> {
97        // validate body
98        self.body.validate()?;
99        let resp = self.post().await?;
100        let parsed = parse_typed_response::<DocumentReembeddingResponse>(resp).await?;
101        Ok(parsed)
102    }
103}
104
105impl HttpClient for DocumentReembeddingRequest {
106    type Body = DocumentReembeddingBody;
107    type ApiUrl = String;
108    type ApiKey = String;
109
110    fn api_url(&self) -> &Self::ApiUrl {
111        &self.url
112    }
113    fn api_key(&self) -> &Self::ApiKey {
114        &self.key
115    }
116    fn body(&self) -> &Self::Body {
117        &self.body
118    }
119
120    fn http_config(&self) -> Arc<HttpClientConfig> {
121        Arc::clone(&self.http_config)
122    }
123}
124
125/// Simple response envelope: { code, message, timestamp }
126#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
127pub struct DocumentReembeddingResponse {
128    #[serde(skip_serializing_if = "Option::is_none")]
129    pub code: Option<i64>,
130    #[serde(skip_serializing_if = "Option::is_none")]
131    pub message: Option<String>,
132    #[serde(skip_serializing_if = "Option::is_none")]
133    pub timestamp: Option<u64>,
134}