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