Skip to main content

zai_rs/knowledge/
create.rs

1use serde::{Deserialize, Deserializer, Serialize, Serializer};
2use validator::Validate;
3
4use super::types::KnowledgeResponse;
5use crate::ZaiResult;
6use crate::client::ZaiClient;
7
8/// Embedding model identifier encoded as the service's integer value.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum EmbeddingId {
11    /// 3: Embedding-2
12    Embedding2,
13    /// 11: Embedding-3
14    Embedding3New,
15    /// 12: Embedding-3-pro
16    Embedding3Pro,
17}
18
19impl EmbeddingId {
20    /// Return the upstream integer id for this embedding model.
21    pub fn as_i64(&self) -> i64 {
22        match self {
23            EmbeddingId::Embedding2 => 3,
24            EmbeddingId::Embedding3New => 11,
25            EmbeddingId::Embedding3Pro => 12,
26        }
27    }
28
29    /// Return the canonical model name paired with this identifier.
30    pub fn as_model_name(&self) -> &'static str {
31        match self {
32            EmbeddingId::Embedding2 => "Embedding-2",
33            EmbeddingId::Embedding3New => "Embedding-3",
34            EmbeddingId::Embedding3Pro => "Embedding-3-pro",
35        }
36    }
37}
38
39impl Serialize for EmbeddingId {
40    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
41        serializer.serialize_i64(self.as_i64())
42    }
43}
44
45impl<'de> Deserialize<'de> for EmbeddingId {
46    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
47        let v = i64::deserialize(deserializer)?;
48        match v {
49            3 => Ok(EmbeddingId::Embedding2),
50            11 => Ok(EmbeddingId::Embedding3New),
51            12 => Ok(EmbeddingId::Embedding3Pro),
52            other => Err(serde::de::Error::custom(format!(
53                "unsupported embedding_id: {other} (expected 3, 11 or 12)"
54            ))),
55        }
56    }
57}
58
59/// Background color enum
60#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
61#[serde(rename_all = "lowercase")]
62pub enum BackgroundColor {
63    /// Blue background.
64    Blue,
65    /// Red background.
66    Red,
67    /// Orange background.
68    Orange,
69    /// Purple background.
70    Purple,
71    /// Sky-blue background.
72    Sky,
73    /// Green background.
74    Green,
75    /// Yellow background.
76    Yellow,
77}
78
79/// Knowledge icon enum
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
81#[serde(rename_all = "lowercase")]
82pub enum KnowledgeIcon {
83    /// Question-mark icon.
84    Question,
85    /// Book icon.
86    Book,
87    /// Seal icon.
88    Seal,
89    /// Wrench icon.
90    Wrench,
91    /// Tag icon.
92    Tag,
93    /// Horn icon.
94    Horn,
95    /// House icon.
96    House,
97}
98
99/// Request body for creating a knowledge base
100#[derive(Clone, Serialize, Deserialize, Validate)]
101pub struct KnowledgeCreateBody {
102    /// Embedding model ID (3, 11, or 12).
103    pub embedding_id: EmbeddingId,
104    /// Knowledge base name
105    #[validate(length(min = 1))]
106    pub name: String,
107    /// Knowledge base description (optional)
108    #[serde(skip_serializing_if = "Option::is_none")]
109    pub description: Option<String>,
110    /// Background color (optional; default blue on server)
111    #[serde(skip_serializing_if = "Option::is_none")]
112    pub background: Option<BackgroundColor>,
113    /// Icon name (optional; default question on server)
114    #[serde(skip_serializing_if = "Option::is_none")]
115    pub icon: Option<KnowledgeIcon>,
116    /// Embedding model name (optional). When given alongside `embedding_id`,
117    /// the service requires the two to be consistent.
118    #[serde(skip_serializing_if = "Option::is_none")]
119    pub embedding_model: Option<String>,
120    /// Contextual retrieval flag (`0` or `1`).
121    #[serde(skip_serializing_if = "Option::is_none")]
122    #[validate(range(max = 1))]
123    pub contextual: Option<u8>,
124}
125
126impl std::fmt::Debug for KnowledgeCreateBody {
127    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128        formatter
129            .debug_struct("KnowledgeCreateBody")
130            .field("embedding_id", &self.embedding_id)
131            .field("name", &"[REDACTED]")
132            .field("description_configured", &self.description.is_some())
133            .field("background", &self.background)
134            .field("icon", &self.icon)
135            .field(
136                "embedding_model_configured",
137                &self.embedding_model.is_some(),
138            )
139            .field("contextual", &self.contextual)
140            .finish()
141    }
142}
143
144impl KnowledgeCreateBody {
145    fn validate_embedding_pair(&self) -> ZaiResult<()> {
146        if let Some(model) = self.embedding_model.as_deref()
147            && model != self.embedding_id.as_model_name()
148        {
149            return Err(crate::ZaiError::ApiError {
150                code: crate::client::error::codes::SDK_VALIDATION,
151                message: format!(
152                    "embedding_id {} requires embedding_model '{}'",
153                    self.embedding_id.as_i64(),
154                    self.embedding_id.as_model_name()
155                ),
156            });
157        }
158        Ok(())
159    }
160}
161
162/// Create knowledge request (POST /llm-application/open/knowledge)
163///
164/// Credentials and transport live on the [`ZaiClient`], passed to
165/// [`send_via`](Self::send_via).
166pub struct KnowledgeCreateRequest {
167    body: KnowledgeCreateBody,
168}
169
170impl KnowledgeCreateRequest {
171    /// Build a create request with required fields
172    pub fn new(embedding_id: EmbeddingId, name: impl Into<String>) -> Self {
173        let body = KnowledgeCreateBody {
174            embedding_id,
175            name: name.into(),
176            description: None,
177            background: None,
178            icon: None,
179            embedding_model: None,
180            contextual: None,
181        };
182        Self { body }
183    }
184
185    /// Set the knowledge-base description.
186    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
187        self.body.description = Some(desc.into());
188        self
189    }
190    /// Set the background color.
191    pub fn with_background(mut self, bg: BackgroundColor) -> Self {
192        self.body.background = Some(bg);
193        self
194    }
195    /// Set the icon.
196    pub fn with_icon(mut self, icon: KnowledgeIcon) -> Self {
197        self.body.icon = Some(icon);
198        self
199    }
200    /// Set the embedding model name. It must match the configured
201    /// [`EmbeddingId`]; [`Self::send_via`] validates the pair locally.
202    pub fn with_embedding_model(mut self, model: impl Into<String>) -> Self {
203        self.body.embedding_model = Some(model.into());
204        self
205    }
206    /// Set the contextual-retrieval flag (`0` or `1`).
207    ///
208    /// [`Self::send_via`] rejects values other than `0` and `1`.
209    pub fn with_contextual(mut self, contextual: u8) -> Self {
210        self.body.contextual = Some(contextual);
211        self
212    }
213
214    /// Validate all documented request constraints without performing I/O.
215    pub fn validate(&self) -> ZaiResult<()> {
216        self.body.validate()?;
217        if self.body.name.trim().is_empty() {
218            return Err(crate::client::validation::invalid(
219                "knowledge name must not be blank",
220            ));
221        }
222        self.body.validate_embedding_pair()
223    }
224
225    /// Validate and send via a [`ZaiClient`], returning the typed response.
226    pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<KnowledgeCreateResponse> {
227        self.validate()?;
228        let route = crate::client::routes::KNOWLEDGE_CREATE;
229        let url = client.endpoints().resolve_route(route, &[])?;
230        client
231            .send_json::<_, KnowledgeCreateResponse>(route.method(), url, &self.body)
232            .await
233    }
234}
235
236/// Inner data of [`KnowledgeCreateResponse`] — the newly created id.
237#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
238pub struct KnowledgeCreateData {
239    /// Newly created id
240    #[serde(skip_serializing_if = "Option::is_none")]
241    pub id: Option<String>,
242}
243
244/// Response of knowledge creation.
245pub type KnowledgeCreateResponse = KnowledgeResponse<KnowledgeCreateData>;
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    #[test]
252    fn rejects_blank_names_and_mismatched_embedding_models() {
253        assert!(
254            KnowledgeCreateRequest::new(EmbeddingId::Embedding3New, " \t")
255                .validate()
256                .is_err()
257        );
258        assert!(
259            KnowledgeCreateRequest::new(EmbeddingId::Embedding3New, "docs")
260                .with_embedding_model("Embedding-2")
261                .validate()
262                .is_err()
263        );
264        assert!(
265            KnowledgeCreateRequest::new(EmbeddingId::Embedding3New, "docs")
266                .with_embedding_model("Embedding-3")
267                .validate()
268                .is_ok()
269        );
270    }
271
272    #[test]
273    fn request_body_debug_redacts_names_and_descriptions() {
274        let body = KnowledgeCreateBody {
275            embedding_id: EmbeddingId::Embedding3New,
276            name: "private-name".to_owned(),
277            description: Some("private-description".to_owned()),
278            background: Some(BackgroundColor::Blue),
279            icon: Some(KnowledgeIcon::Book),
280            embedding_model: Some("private-model-string".to_owned()),
281            contextual: Some(1),
282        };
283        let debug = format!("{body:?}");
284        for secret in [
285            "private-name",
286            "private-description",
287            "private-model-string",
288        ] {
289            assert!(!debug.contains(secret));
290        }
291    }
292}