Skip to main content

lellm_core/
message.rs

1//! 消息与内容块类型。
2
3use serde::{Deserialize, Serialize};
4
5/// 纯文本块
6#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
7pub struct TextBlock {
8    pub text: String,
9}
10
11/// 思考块(Claude thinking / OpenAI reasoning)
12#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
13pub struct ThinkingBlock {
14    pub thinking: String,
15    /// 部分 provider 支持 redacted thinking
16    pub redacted: Option<String>,
17}
18
19/// 图片资源
20#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
21pub struct ImageSource {
22    /// base64 编码的图片数据
23    pub data: String,
24    /// MIME 类型,如 "image/png"
25    pub media_type: String,
26}
27
28/// LLM 请求的工具调用。
29#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
30pub struct ToolCall {
31    pub id: String,
32    pub name: String,
33    pub arguments: serde_json::Value,
34}
35
36/// 内容块 — Message 和 ChatResponse 的基本组成单元。
37/// 核心层极简,无 provider 特有标记。
38#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
39#[serde(tag = "type", rename_all = "snake_case")]
40pub enum ContentBlock {
41    Text(TextBlock),
42    Thinking(ThinkingBlock),
43    Image { source: ImageSource },
44    ToolCall(ToolCall),
45}
46
47impl ContentBlock {
48    pub fn text(s: String) -> Self {
49        ContentBlock::Text(TextBlock { text: s })
50    }
51
52    pub fn as_text(&self) -> Option<&str> {
53        match self {
54            ContentBlock::Text(block) => Some(&block.text),
55            _ => None,
56        }
57    }
58}
59
60/// 对话中的单条消息。
61#[derive(Debug, Clone, Serialize, Deserialize)]
62#[serde(tag = "type", rename_all = "snake_case")]
63pub enum Message {
64    System {
65        content: Vec<ContentBlock>,
66    },
67    User {
68        content: Vec<ContentBlock>,
69    },
70    Assistant {
71        content: Vec<ContentBlock>,
72    },
73    ToolResult {
74        tool_call_id: String,
75        content: Vec<ContentBlock>,
76    },
77}
78
79impl Message {
80    /// 返回适合 API 序列化的 role 字符串。
81    /// OpenAI / Anthropic 统一使用 `role` 字段。
82    pub fn role(&self) -> &str {
83        match self {
84            Message::System { .. } => "system",
85            Message::User { .. } => "user",
86            Message::Assistant { .. } => "assistant",
87            Message::ToolResult { .. } => "tool_result",
88        }
89    }
90
91    /// 返回内容块的引用(用于 provider 适配器序列化)
92    pub fn content(&self) -> &Vec<ContentBlock> {
93        match self {
94            Message::System { content }
95            | Message::User { content }
96            | Message::Assistant { content }
97            | Message::ToolResult { content, .. } => content,
98        }
99    }
100
101    /// 提取所有文本块拼接为字符串
102    pub fn extract_text(&self) -> String {
103        match self {
104            Message::System { content } => Self::join_text(content),
105            Message::User { content } => Self::join_text(content),
106            Message::Assistant { content } => Self::join_text(content),
107            Message::ToolResult { content, .. } => Self::join_text(content),
108        }
109    }
110
111    fn join_text(blocks: &[ContentBlock]) -> String {
112        blocks
113            .iter()
114            .filter_map(|b| b.as_text().map(|s| s.to_string()))
115            .collect::<Vec<_>>()
116            .join("")
117    }
118
119    /// 提取所有 ToolCall
120    pub fn extract_tool_calls(&self) -> Vec<ToolCall> {
121        match self {
122            Message::Assistant { content } => content
123                .iter()
124                .filter_map(|b| {
125                    if let ContentBlock::ToolCall(tc) = b {
126                        Some(tc.clone())
127                    } else {
128                        None
129                    }
130                })
131                .collect(),
132            _ => Vec::new(),
133        }
134    }
135}
136
137/// 便捷函数:创建纯文本块
138pub fn text_block(s: String) -> Vec<ContentBlock> {
139    vec![ContentBlock::text(s)]
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[test]
147    fn test_content_block_text() {
148        let block = ContentBlock::text("hello".to_string());
149        assert_eq!(block.as_text(), Some("hello"));
150    }
151
152    #[test]
153    fn test_content_block_tool_call_no_as_text() {
154        let block = ContentBlock::ToolCall(ToolCall {
155            id: "1".into(),
156            name: "test".into(),
157            arguments: serde_json::json!({}),
158        });
159        assert_eq!(block.as_text(), None);
160    }
161
162    #[test]
163    fn test_message_extract_text() {
164        let msg = Message::User {
165            content: text_block("hello world".to_string()),
166        };
167        assert_eq!(msg.extract_text(), "hello world");
168    }
169
170    #[test]
171    fn test_message_extract_tool_calls() {
172        let tc = ToolCall {
173            id: "1".into(),
174            name: "test".into(),
175            arguments: serde_json::json!({}),
176        };
177        let msg = Message::Assistant {
178            content: vec![ContentBlock::ToolCall(tc.clone())],
179        };
180        let calls = msg.extract_tool_calls();
181        assert_eq!(calls.len(), 1);
182        assert_eq!(calls[0].name, "test");
183    }
184}