Skip to main content

zai_rs/knowledge/
update.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4use validator::Validate;
5
6use super::create::{BackgroundColor, EmbeddingId, KnowledgeIcon};
7use super::types::KnowledgeOperationResponse;
8use crate::client::ZaiClient;
9
10/// Update body for editing a knowledge base
11#[derive(Clone, Default, Serialize, Deserialize, Validate)]
12pub struct KnowledgeUpdateBody {
13    /// Embedding model ID (3, 11, or 12).
14    #[serde(skip_serializing_if = "Option::is_none")]
15    pub embedding_id: Option<EmbeddingId>,
16    /// Embedding model name. When `embedding_id` is also set, both values must
17    /// identify the same model.
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub embedding_model: Option<String>,
20    /// Contextual retrieval flag (`0` or `1`).
21    #[serde(skip_serializing_if = "Option::is_none")]
22    #[validate(range(max = 1))]
23    pub contextual: Option<u8>,
24    /// Knowledge base name
25    #[serde(skip_serializing_if = "Option::is_none")]
26    #[validate(length(min = 1))]
27    pub name: Option<String>,
28    /// Knowledge base description
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub description: Option<String>,
31    /// Background color
32    #[serde(skip_serializing_if = "Option::is_none")]
33    pub background: Option<BackgroundColor>,
34    /// Icon name
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub icon: Option<KnowledgeIcon>,
37    /// Callback URL when rebuilding is required
38    #[serde(skip_serializing_if = "Option::is_none")]
39    #[validate(url)]
40    pub callback_url: Option<String>,
41    /// Callback headers as key-value
42    #[serde(skip_serializing_if = "Option::is_none")]
43    pub callback_header: Option<BTreeMap<String, String>>,
44}
45
46impl std::fmt::Debug for KnowledgeUpdateBody {
47    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        formatter
49            .debug_struct("KnowledgeUpdateBody")
50            .field("embedding_id", &self.embedding_id)
51            .field(
52                "embedding_model_configured",
53                &self.embedding_model.is_some(),
54            )
55            .field("contextual", &self.contextual)
56            .field("name_configured", &self.name.is_some())
57            .field("description_configured", &self.description.is_some())
58            .field("background", &self.background)
59            .field("icon", &self.icon)
60            .field("callback_url_configured", &self.callback_url.is_some())
61            .field(
62                "callback_header_entries",
63                &self.callback_header.as_ref().map(BTreeMap::len),
64            )
65            .finish()
66    }
67}
68
69impl KnowledgeUpdateBody {
70    fn validate_embedding_pair(&self) -> crate::ZaiResult<()> {
71        let Some(model) = self.embedding_model.as_deref() else {
72            return Ok(());
73        };
74        if !matches!(model, "Embedding-2" | "Embedding-3" | "Embedding-3-pro") {
75            return Err(crate::client::validation::invalid(
76                "embedding_model must be one of: Embedding-2, Embedding-3, Embedding-3-pro",
77            ));
78        }
79        if let Some(id) = self.embedding_id
80            && model != id.as_model_name()
81        {
82            return Err(crate::client::validation::invalid(format!(
83                "embedding_id {} requires embedding_model '{}'",
84                id.as_i64(),
85                id.as_model_name()
86            )));
87        }
88        Ok(())
89    }
90}
91
92/// Knowledge update request (PUT /llm-application/open/knowledge/{id})
93///
94/// Credentials and transport live on the [`ZaiClient`], passed to
95/// [`send_via`](Self::send_via).
96pub struct KnowledgeUpdateRequest {
97    id: String,
98    body: KnowledgeUpdateBody,
99}
100
101impl KnowledgeUpdateRequest {
102    /// Build update request targeting a specific id with empty body.
103    pub fn new(id: impl Into<String>) -> Self {
104        Self {
105            id: id.into(),
106            body: KnowledgeUpdateBody::default(),
107        }
108    }
109
110    /// Set the embedding model id.
111    pub fn with_embedding_id(mut self, id: EmbeddingId) -> Self {
112        self.body.embedding_id = Some(id);
113        self
114    }
115    /// Set the embedding model name.
116    pub fn with_embedding_model(mut self, model: impl Into<String>) -> Self {
117        self.body.embedding_model = Some(model.into());
118        self
119    }
120    /// Enable (`1`) or disable (`0`) contextual retrieval.
121    pub fn with_contextual(mut self, contextual: u8) -> Self {
122        self.body.contextual = Some(contextual);
123        self
124    }
125    /// Set the knowledge-base name.
126    pub fn with_name(mut self, name: impl Into<String>) -> Self {
127        self.body.name = Some(name.into());
128        self
129    }
130    /// Set the description.
131    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
132        self.body.description = Some(desc.into());
133        self
134    }
135    /// Set the background color.
136    pub fn with_background(mut self, bg: BackgroundColor) -> Self {
137        self.body.background = Some(bg);
138        self
139    }
140    /// Set the icon.
141    pub fn with_icon(mut self, icon: KnowledgeIcon) -> Self {
142        self.body.icon = Some(icon);
143        self
144    }
145    /// Set the callback URL (notified when rebuilding completes).
146    pub fn with_callback_url(mut self, url: impl Into<String>) -> Self {
147        self.body.callback_url = Some(url.into());
148        self
149    }
150    /// Set the callback headers.
151    pub fn with_callback_header(mut self, headers: BTreeMap<String, String>) -> Self {
152        self.body.callback_header = Some(headers);
153        self
154    }
155
156    /// Validate the target identifier and update body without performing I/O.
157    pub fn validate(&self) -> crate::ZaiResult<()> {
158        crate::client::validation::require_non_blank(&self.id, "knowledge_id")?;
159        self.body.validate()?;
160        if self
161            .body
162            .name
163            .as_deref()
164            .is_some_and(|name| name.trim().is_empty())
165        {
166            return Err(crate::client::validation::invalid(
167                "knowledge name must not be blank",
168            ));
169        }
170        self.body.validate_embedding_pair()
171    }
172
173    /// Send the update request via a [`ZaiClient`] and parse the typed response.
174    pub async fn send_via(&self, client: &ZaiClient) -> crate::ZaiResult<KnowledgeUpdateResponse> {
175        self.validate()?;
176        let route = crate::client::routes::KNOWLEDGE_UPDATE;
177        let url = client.endpoints().resolve_route(route, &[&self.id])?;
178        client
179            .send_json::<_, KnowledgeUpdateResponse>(route.method(), url, &self.body)
180            .await
181    }
182}
183
184/// Update response envelope without a data payload.
185pub type KnowledgeUpdateResponse = KnowledgeOperationResponse;
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn rejects_blank_target_and_blank_name_without_inventing_required_fields() {
193        assert!(
194            KnowledgeUpdateRequest::new(" ")
195                .with_description("value")
196                .validate()
197                .is_err()
198        );
199        assert!(KnowledgeUpdateRequest::new("kb-1").validate().is_ok());
200        assert!(
201            KnowledgeUpdateRequest::new("kb-1")
202                .with_name(" \t")
203                .validate()
204                .is_err()
205        );
206        assert!(
207            KnowledgeUpdateRequest::new("kb-1")
208                .with_name("docs")
209                .validate()
210                .is_ok()
211        );
212        assert_eq!(
213            serde_json::to_value(&KnowledgeUpdateRequest::new("kb-1").body).unwrap(),
214            serde_json::json!({})
215        );
216    }
217
218    #[test]
219    fn request_body_debug_redacts_text_callbacks_and_headers() {
220        let body = KnowledgeUpdateBody {
221            embedding_model: Some("private-model".to_owned()),
222            name: Some("private-name".to_owned()),
223            description: Some("private-description".to_owned()),
224            callback_url: Some("https://private.example/callback".to_owned()),
225            callback_header: Some(BTreeMap::from([(
226                "Authorization".to_owned(),
227                "private-token".to_owned(),
228            )])),
229            ..KnowledgeUpdateBody::default()
230        };
231        let debug = format!("{body:?}");
232        for secret in [
233            "private-model",
234            "private-name",
235            "private-description",
236            "private.example",
237            "Authorization",
238            "private-token",
239        ] {
240            assert!(!debug.contains(secret));
241        }
242        assert!(debug.contains("callback_header_entries: Some(1)"));
243    }
244}