1use serde::{Deserialize, Deserializer, Serialize, Serializer};
2use validator::Validate;
3
4use super::types::KnowledgeResponse;
5use crate::ZaiResult;
6use crate::client::ZaiClient;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum EmbeddingId {
11 Embedding2,
13 Embedding3New,
15 Embedding3Pro,
17}
18
19impl EmbeddingId {
20 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 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
61#[serde(rename_all = "lowercase")]
62pub enum BackgroundColor {
63 Blue,
65 Red,
67 Orange,
69 Purple,
71 Sky,
73 Green,
75 Yellow,
77}
78
79#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
81#[serde(rename_all = "lowercase")]
82pub enum KnowledgeIcon {
83 Question,
85 Book,
87 Seal,
89 Wrench,
91 Tag,
93 Horn,
95 House,
97}
98
99#[derive(Clone, Serialize, Deserialize, Validate)]
101pub struct KnowledgeCreateBody {
102 pub embedding_id: EmbeddingId,
104 #[validate(length(min = 1))]
106 pub name: String,
107 #[serde(skip_serializing_if = "Option::is_none")]
109 pub description: Option<String>,
110 #[serde(skip_serializing_if = "Option::is_none")]
112 pub background: Option<BackgroundColor>,
113 #[serde(skip_serializing_if = "Option::is_none")]
115 pub icon: Option<KnowledgeIcon>,
116 #[serde(skip_serializing_if = "Option::is_none")]
119 pub embedding_model: Option<String>,
120 #[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
162pub struct KnowledgeCreateRequest {
167 body: KnowledgeCreateBody,
168}
169
170impl KnowledgeCreateRequest {
171 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 pub fn with_description(mut self, desc: impl Into<String>) -> Self {
187 self.body.description = Some(desc.into());
188 self
189 }
190 pub fn with_background(mut self, bg: BackgroundColor) -> Self {
192 self.body.background = Some(bg);
193 self
194 }
195 pub fn with_icon(mut self, icon: KnowledgeIcon) -> Self {
197 self.body.icon = Some(icon);
198 self
199 }
200 pub fn with_embedding_model(mut self, model: impl Into<String>) -> Self {
203 self.body.embedding_model = Some(model.into());
204 self
205 }
206 pub fn with_contextual(mut self, contextual: u8) -> Self {
210 self.body.contextual = Some(contextual);
211 self
212 }
213
214 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 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#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
238pub struct KnowledgeCreateData {
239 #[serde(skip_serializing_if = "Option::is_none")]
241 pub id: Option<String>,
242}
243
244pub 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}