Skip to main content

zai_rs/model/text_tokenizer/
request.rs

1use serde::{Deserialize, Serialize};
2
3/// Tokenizer-capable models
4#[derive(Debug, Clone, Serialize, Deserialize)]
5#[serde(rename_all = "kebab-case")]
6#[derive(Default)]
7pub enum TokenizerModel {
8    #[serde(rename = "glm-4-plus")]
9    #[default]
10    Glm4Plus,
11    #[serde(rename = "glm-4-0520")]
12    Glm40520,
13    #[serde(rename = "glm-4-long")]
14    Glm4Long,
15    #[serde(rename = "glm-4-air")]
16    Glm4Air,
17    #[serde(rename = "glm-4-flash")]
18    Glm4Flash,
19}
20
21/// One message item for tokenizer input
22#[derive(Debug, Clone, Serialize, Deserialize)]
23#[serde(tag = "role", rename_all = "lowercase")]
24pub enum TokenizerMessage {
25    /// 用户消息
26    User { content: String },
27    /// 系统消息
28    System { content: String },
29    /// 助手消息(content 可选)
30    Assistant {
31        #[serde(skip_serializing_if = "Option::is_none")]
32        content: Option<String>,
33    },
34    /// 工具消息
35    Tool { content: String },
36}
37
38/// Request body for tokenizer
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct TokenizerBody {
41    /// 调用的模型代码(默认 glm-4-plus)
42    pub model: TokenizerModel,
43    /// 对话消息列表(至少 1 条)
44    pub messages: Vec<TokenizerMessage>,
45    /// 客户端请求 ID
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub request_id: Option<String>,
48    /// 终端用户 ID
49    #[serde(skip_serializing_if = "Option::is_none")]
50    pub user_id: Option<String>,
51}
52
53impl TokenizerBody {
54    pub fn new(model: TokenizerModel, messages: Vec<TokenizerMessage>) -> Self {
55        Self {
56            model,
57            messages,
58            request_id: None,
59            user_id: None,
60        }
61    }
62    pub fn with_request_id(mut self, v: impl Into<String>) -> Self {
63        self.request_id = Some(v.into());
64        self
65    }
66    pub fn with_user_id(mut self, v: impl Into<String>) -> Self {
67        self.user_id = Some(v.into());
68        self
69    }
70}