1use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use std::collections::HashMap;
6
7use lc_shared::tools::ToolCall;
8
9use super::audio::AudioContent;
10use super::file::FileContent;
11use super::image::ImageContent;
12
13#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
15#[serde(rename_all = "lowercase")]
16pub enum MessageType {
17 System,
18 Human,
19 AI,
20 Tool { tool_call_id: String },
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct Message {
26 pub content: String,
27
28 #[serde(default)]
30 pub images: Vec<ImageContent>,
31
32 #[serde(default)]
34 pub audio: Vec<AudioContent>,
35
36 #[serde(default)]
38 pub files: Vec<FileContent>,
39
40 #[serde(rename = "type")]
41 pub message_type: MessageType,
42
43 #[serde(skip_serializing_if = "Option::is_none")]
44 pub name: Option<String>,
45
46 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
47 pub additional_kwargs: HashMap<String, Value>,
48
49 #[serde(skip_serializing_if = "Option::is_none")]
50 pub id: Option<String>,
51
52 #[serde(skip_serializing_if = "Option::is_none")]
53 pub tool_calls: Option<Vec<ToolCall>>,
54}
55
56impl Message {
57 pub fn system(content: impl Into<String>) -> Self {
59 Self {
60 content: content.into(),
61 images: Vec::new(),
62 audio: Vec::new(),
63 files: Vec::new(),
64 message_type: MessageType::System,
65 name: None,
66 additional_kwargs: HashMap::new(),
67 id: None,
68 tool_calls: None,
69 }
70 }
71
72 pub fn human(content: impl Into<String>) -> Self {
74 Self {
75 content: content.into(),
76 images: Vec::new(),
77 audio: Vec::new(),
78 files: Vec::new(),
79 message_type: MessageType::Human,
80 name: None,
81 additional_kwargs: HashMap::new(),
82 id: None,
83 tool_calls: None,
84 }
85 }
86
87 pub fn human_with_image(content: impl Into<String>, image_url: impl Into<String>) -> Self {
89 Self {
90 content: content.into(),
91 images: vec![ImageContent::from_url(image_url)],
92 audio: Vec::new(),
93 files: Vec::new(),
94 message_type: MessageType::Human,
95 name: None,
96 additional_kwargs: HashMap::new(),
97 id: None,
98 tool_calls: None,
99 }
100 }
101
102 pub fn human_with_images(content: impl Into<String>, images: Vec<ImageContent>) -> Self {
104 Self {
105 content: content.into(),
106 images,
107 audio: Vec::new(),
108 files: Vec::new(),
109 message_type: MessageType::Human,
110 name: None,
111 additional_kwargs: HashMap::new(),
112 id: None,
113 tool_calls: None,
114 }
115 }
116
117 pub fn human_with_audio(content: impl Into<String>, audio: AudioContent) -> Self {
119 Self {
120 content: content.into(),
121 images: Vec::new(),
122 audio: vec![audio],
123 files: Vec::new(),
124 message_type: MessageType::Human,
125 name: None,
126 additional_kwargs: HashMap::new(),
127 id: None,
128 tool_calls: None,
129 }
130 }
131
132 pub fn human_with_file(content: impl Into<String>, file: FileContent) -> Self {
134 Self {
135 content: content.into(),
136 images: Vec::new(),
137 audio: Vec::new(),
138 files: vec![file],
139 message_type: MessageType::Human,
140 name: None,
141 additional_kwargs: HashMap::new(),
142 id: None,
143 tool_calls: None,
144 }
145 }
146
147 pub fn ai(content: impl Into<String>) -> Self {
149 Self {
150 content: content.into(),
151 images: Vec::new(),
152 audio: Vec::new(),
153 files: Vec::new(),
154 message_type: MessageType::AI,
155 name: None,
156 additional_kwargs: HashMap::new(),
157 id: None,
158 tool_calls: None,
159 }
160 }
161
162 pub fn ai_with_tool_calls(content: impl Into<String>, tool_calls: Vec<ToolCall>) -> Self {
164 Self {
165 content: content.into(),
166 images: Vec::new(),
167 audio: Vec::new(),
168 files: Vec::new(),
169 message_type: MessageType::AI,
170 name: None,
171 additional_kwargs: HashMap::new(),
172 id: None,
173 tool_calls: Some(tool_calls),
174 }
175 }
176
177 pub fn tool(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
179 Self {
180 content: content.into(),
181 images: Vec::new(),
182 audio: Vec::new(),
183 files: Vec::new(),
184 message_type: MessageType::Tool {
185 tool_call_id: tool_call_id.into(),
186 },
187 name: None,
188 additional_kwargs: HashMap::new(),
189 id: None,
190 tool_calls: None,
191 }
192 }
193
194 pub fn with_name(mut self, name: impl Into<String>) -> Self {
196 self.name = Some(name.into());
197 self
198 }
199
200 pub fn with_id(mut self, id: impl Into<String>) -> Self {
202 self.id = Some(id.into());
203 self
204 }
205
206 pub fn with_additional_kwarg(mut self, key: impl Into<String>, value: Value) -> Self {
208 self.additional_kwargs.insert(key.into(), value);
209 self
210 }
211
212 pub fn with_image(mut self, image: ImageContent) -> Self {
214 self.images.push(image);
215 self
216 }
217
218 pub fn with_audio(mut self, audio: AudioContent) -> Self {
220 self.audio.push(audio);
221 self
222 }
223
224 pub fn with_file(mut self, file: FileContent) -> Self {
226 self.files.push(file);
227 self
228 }
229
230 pub fn has_images(&self) -> bool {
232 !self.images.is_empty()
233 }
234
235 pub fn has_audio(&self) -> bool {
237 !self.audio.is_empty()
238 }
239
240 pub fn has_files(&self) -> bool {
242 !self.files.is_empty()
243 }
244
245 pub fn is_multimodal(&self) -> bool {
247 self.has_images() || self.has_audio() || self.has_files()
248 }
249
250 pub fn type_str(&self) -> String {
255 match &self.message_type {
256 MessageType::System => "system".to_string(),
257 MessageType::Human => "human".to_string(),
258 MessageType::AI => "ai".to_string(),
259 MessageType::Tool { tool_call_id } => format!("tool:{tool_call_id}"),
260 }
261 }
262
263 pub fn has_tool_calls(&self) -> bool {
265 self.tool_calls.as_deref().is_some_and(|t| !t.is_empty())
266 }
267
268 pub fn get_tool_calls(&self) -> Option<&[ToolCall]> {
270 self.tool_calls.as_deref()
271 }
272}
273
274#[cfg(test)]
275mod tests {
276 use super::*;
277
278 #[test]
279 fn test_human_with_image() {
280 let msg = Message::human_with_image("描述这张图", "https://example.com/img.jpg");
281 assert_eq!(msg.content, "描述这张图");
282 assert_eq!(msg.images.len(), 1);
283 assert_eq!(msg.images[0].url, "https://example.com/img.jpg");
284 assert!(msg.has_images());
285 }
286
287 #[test]
288 fn test_human_no_images_by_default() {
289 let msg = Message::human("纯文本");
290 assert!(msg.images.is_empty());
291 assert!(!msg.has_images());
292 }
293
294 #[test]
295 fn test_with_image_builder() {
296 let msg = Message::human("看图")
297 .with_image(ImageContent::from_url("https://example.com/a.png"))
298 .with_image(ImageContent::from_base64("abc"));
299 assert_eq!(msg.images.len(), 2);
300 }
301
302 #[test]
303 fn test_message_deserialize_without_images_field() {
304 let json = r#"{"content":"hi","type":"human"}"#;
306 let msg: Message = serde_json::from_str(json).unwrap();
307 assert_eq!(msg.content, "hi");
308 assert!(msg.images.is_empty());
309 }
310
311 #[test]
312 fn test_human_with_images_multiple() {
313 let msg = Message::human_with_images(
314 "多图",
315 vec![
316 ImageContent::from_url("https://example.com/1.jpg"),
317 ImageContent::from_url("https://example.com/2.jpg"),
318 ],
319 );
320 assert_eq!(msg.images.len(), 2);
321 }
322
323 #[test]
324 fn test_system_ai_no_images() {
325 assert!(Message::system("s").images.is_empty());
326 assert!(Message::ai("a").images.is_empty());
327 assert!(Message::tool("id", "c").images.is_empty());
328 }
329
330 #[test]
331 fn test_type_str_includes_tool_call_id() {
332 assert_eq!(Message::system("s").type_str(), "system");
333 assert_eq!(Message::human("h").type_str(), "human");
334 assert_eq!(Message::ai("a").type_str(), "ai");
335 assert_eq!(
336 Message::tool("call_123", "result").type_str(),
337 "tool:call_123"
338 );
339 }
340
341 #[test]
342 fn test_has_tool_calls_empty_and_present() {
343 let with_calls = Message::ai_with_tool_calls(
344 "call tool",
345 vec![ToolCall::new("call_1", "weather", r#"{"city":"beijing"}"#)],
346 );
347 assert!(with_calls.has_tool_calls());
348 assert_eq!(with_calls.get_tool_calls().unwrap().len(), 1);
349
350 assert!(!Message::ai("plain").has_tool_calls());
352 let empty = Message::ai_with_tool_calls("no calls", vec![]);
353 assert!(!empty.has_tool_calls());
354 }
355}