Skip to main content

zai_rs/knowledge/
document_reembedding.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4use validator::Validate;
5
6use super::types::KnowledgeOperationResponse;
7use crate::ZaiResult;
8use crate::client::ZaiClient;
9
10/// Request body for re-embedding a document
11#[derive(Clone, Serialize, Deserialize, Validate, Default)]
12pub struct DocumentReembedBody {
13    /// Optional callback URL that will be called when re-embedding completes
14    #[serde(skip_serializing_if = "Option::is_none")]
15    #[validate(url)]
16    pub callback_url: Option<String>,
17    /// Optional callback headers key-value pairs
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub callback_header: Option<BTreeMap<String, String>>,
20}
21
22impl std::fmt::Debug for DocumentReembedBody {
23    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        formatter
25            .debug_struct("DocumentReembedBody")
26            .field("callback_url_configured", &self.callback_url.is_some())
27            .field(
28                "callback_header_entries",
29                &self.callback_header.as_ref().map(BTreeMap::len),
30            )
31            .finish()
32    }
33}
34
35/// Re-embedding request (POST /llm-application/open/document/embedding/{id})
36///
37/// Credentials and transport live on the [`ZaiClient`], passed to
38/// [`send_via`](Self::send_via).
39pub struct DocumentReembedRequest {
40    document_id: String,
41    body: DocumentReembedBody,
42}
43
44impl DocumentReembedRequest {
45    /// Create a new request for the specified document id
46    pub fn new(document_id: impl Into<String>) -> Self {
47        Self {
48            document_id: document_id.into(),
49            body: DocumentReembedBody::default(),
50        }
51    }
52
53    /// Set callback URL
54    pub fn with_callback_url(mut self, url: impl Into<String>) -> Self {
55        self.body.callback_url = Some(url.into());
56        self
57    }
58
59    /// Set callback headers
60    pub fn with_callback_header(mut self, hdr: BTreeMap<String, String>) -> Self {
61        self.body.callback_header = Some(hdr);
62        self
63    }
64
65    /// Validate the target identifier and optional callback URL without I/O.
66    pub fn validate(&self) -> ZaiResult<()> {
67        crate::client::validation::require_non_blank(&self.document_id, "document_id")?;
68        self.body.validate()?;
69        Ok(())
70    }
71
72    /// Send the POST request via a [`ZaiClient`] and parse the typed response.
73    pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<DocumentReembedResponse> {
74        self.validate()?;
75        let route = crate::client::routes::DOCUMENTS_REEMBED;
76        let url = client
77            .endpoints()
78            .resolve_route(route, &[&self.document_id])?;
79        if self.body.callback_url.is_none() && self.body.callback_header.is_none() {
80            client
81                .send_empty::<DocumentReembedResponse>(route.method(), url)
82                .await
83        } else {
84            client
85                .send_json::<_, DocumentReembedResponse>(route.method(), url, &self.body)
86                .await
87        }
88    }
89}
90
91/// Simple response envelope without a data payload.
92pub type DocumentReembedResponse = KnowledgeOperationResponse;
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn validates_identifier_and_callback_url() {
100        assert!(DocumentReembedRequest::new(" ").validate().is_err());
101        assert!(
102            DocumentReembedRequest::new("doc-1")
103                .with_callback_url("not a URL")
104                .validate()
105                .is_err()
106        );
107        assert!(DocumentReembedRequest::new("doc-1").validate().is_ok());
108    }
109
110    #[test]
111    fn body_debug_redacts_callback_url_and_headers() {
112        let body = DocumentReembedBody {
113            callback_url: Some("https://private.example/callback".to_owned()),
114            callback_header: Some(BTreeMap::from([(
115                "Authorization".to_owned(),
116                "private-token".to_owned(),
117            )])),
118        };
119        let debug = format!("{body:?}");
120        for secret in ["private.example", "Authorization", "private-token"] {
121            assert!(!debug.contains(secret));
122        }
123        assert!(debug.contains("callback_header_entries: Some(1)"));
124    }
125}