Skip to main content

zai_rs/knowledge/
create.rs

1use serde::{Deserialize, Deserializer, Serialize, Serializer};
2use validator::Validate;
3
4use crate::ZaiResult;
5use crate::client::ZaiClient;
6
7/// Embedding model id enum mapped to integer ids (plan §13.5).
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum EmbeddingId {
10    /// 3: Embedding-2
11    Embedding2,
12    /// 11: Embedding-3
13    Embedding3New,
14    /// 12: Embedding-3-pro
15    Embedding3Pro,
16}
17
18impl EmbeddingId {
19    /// Return the upstream integer id for this embedding model.
20    pub fn as_i64(&self) -> i64 {
21        match self {
22            EmbeddingId::Embedding2 => 3,
23            EmbeddingId::Embedding3New => 11,
24            EmbeddingId::Embedding3Pro => 12,
25        }
26    }
27}
28
29impl Serialize for EmbeddingId {
30    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
31        serializer.serialize_i64(self.as_i64())
32    }
33}
34
35impl<'de> Deserialize<'de> for EmbeddingId {
36    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
37        let v = i64::deserialize(deserializer)?;
38        match v {
39            3 => Ok(EmbeddingId::Embedding2),
40            11 => Ok(EmbeddingId::Embedding3New),
41            12 => Ok(EmbeddingId::Embedding3Pro),
42            other => Err(serde::de::Error::custom(format!(
43                "unsupported embedding_id: {other} (expected 3, 11 or 12)"
44            ))),
45        }
46    }
47}
48
49/// Background color enum
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
51#[serde(rename_all = "lowercase")]
52pub enum BackgroundColor {
53    /// Blue background.
54    Blue,
55    /// Red background.
56    Red,
57    /// Orange background.
58    Orange,
59    /// Purple background.
60    Purple,
61    /// Sky-blue background.
62    Sky,
63    /// Green background.
64    Green,
65    /// Yellow background.
66    Yellow,
67}
68
69/// Knowledge icon enum
70#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
71#[serde(rename_all = "lowercase")]
72pub enum KnowledgeIcon {
73    /// Question-mark icon.
74    Question,
75    /// Book icon.
76    Book,
77    /// Seal icon.
78    Seal,
79    /// Wrench icon.
80    Wrench,
81    /// Tag icon.
82    Tag,
83    /// Horn icon.
84    Horn,
85    /// House icon.
86    House,
87}
88
89/// Request body for creating a knowledge base
90#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
91pub struct CreateKnowledgeBody {
92    /// Embedding model id (3, 11 or 12; plan §13.5).
93    pub embedding_id: EmbeddingId,
94    /// Knowledge base name
95    #[validate(length(min = 1))]
96    pub name: String,
97    /// Knowledge base description (optional)
98    #[serde(skip_serializing_if = "Option::is_none")]
99    pub description: Option<String>,
100    /// Background color (optional; default blue on server)
101    #[serde(skip_serializing_if = "Option::is_none")]
102    pub background: Option<BackgroundColor>,
103    /// Icon name (optional; default question on server)
104    #[serde(skip_serializing_if = "Option::is_none")]
105    pub icon: Option<KnowledgeIcon>,
106    /// Embedding model name (optional). When given alongside `embedding_id`,
107    /// the two must be consistent (plan §13.5).
108    #[serde(skip_serializing_if = "Option::is_none")]
109    pub embedding_model: Option<String>,
110    /// Contextual retrieval flag (0/1; plan §13.5).
111    #[serde(skip_serializing_if = "Option::is_none")]
112    pub contextual: Option<u8>,
113}
114
115/// Create knowledge request (POST /llm-application/open/knowledge)
116///
117/// Credentials and transport live on the [`ZaiClient`], passed to
118/// [`send_via`](Self::send_via).
119pub struct CreateKnowledgeRequest {
120    body: CreateKnowledgeBody,
121}
122
123impl CreateKnowledgeRequest {
124    /// Build a create request with required fields
125    pub fn new(embedding_id: EmbeddingId, name: impl Into<String>) -> Self {
126        let body = CreateKnowledgeBody {
127            embedding_id,
128            name: name.into(),
129            description: None,
130            background: None,
131            icon: None,
132            embedding_model: None,
133            contextual: None,
134        };
135        Self { body }
136    }
137
138    /// Optional fields setters
139    /// Set the knowledge-base description.
140    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
141        self.body.description = Some(desc.into());
142        self
143    }
144    /// Set the background color.
145    pub fn with_background(mut self, bg: BackgroundColor) -> Self {
146        self.body.background = Some(bg);
147        self
148    }
149    /// Set the icon.
150    pub fn with_icon(mut self, icon: KnowledgeIcon) -> Self {
151        self.body.icon = Some(icon);
152        self
153    }
154    /// Set the embedding model name (optional). When given alongside
155    /// `embedding_id`, the two must be consistent (plan §13.5).
156    pub fn with_embedding_model(mut self, model: impl Into<String>) -> Self {
157        self.body.embedding_model = Some(model.into());
158        self
159    }
160    /// Set the contextual-retrieval flag (0 or 1; plan §13.5).
161    pub fn with_contextual(mut self, contextual: u8) -> Self {
162        self.body.contextual = Some(contextual);
163        self
164    }
165
166    /// Validate and send via a [`ZaiClient`], returning the typed response.
167    pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<CreateKnowledgeResponse> {
168        self.body.validate()?;
169        let route = crate::client::routes::KNOWLEDGE_CREATE;
170        let url = client.endpoints().resolve_route(route, &[])?;
171        client
172            .send_json::<_, CreateKnowledgeResponse>(route.method(), url, &self.body)
173            .await
174    }
175}
176
177/// Inner data of [`CreateKnowledgeResponse`] — the newly created id.
178#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
179pub struct CreateKnowledgeResponseData {
180    /// Newly created id
181    #[serde(skip_serializing_if = "Option::is_none")]
182    pub id: Option<String>,
183}
184
185/// Response of knowledge creation
186#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
187pub struct CreateKnowledgeResponse {
188    /// Created knowledge-base data.
189    #[serde(skip_serializing_if = "Option::is_none")]
190    pub data: Option<CreateKnowledgeResponseData>,
191    /// Business status code.
192    #[serde(skip_serializing_if = "Option::is_none")]
193    pub code: Option<i64>,
194    /// Human-readable message.
195    #[serde(skip_serializing_if = "Option::is_none")]
196    pub message: Option<String>,
197    /// Server timestamp.
198    #[serde(skip_serializing_if = "Option::is_none")]
199    pub timestamp: Option<u64>,
200}