1use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use uuid::Uuid;
9
10use crate::impl_serializable;
11
12#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
14#[serde(untagged)]
15pub enum MessageContent {
16 Text(String),
18 Blocks(Vec<ContentBlock>),
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
24#[serde(tag = "type")]
25pub enum ContentBlock {
26 Text { text: String },
28 Image {
30 image_url: String,
31 alt_text: Option<String>,
32 },
33 Json { data: serde_json::Value },
35 ToolCall {
37 id: String,
38 name: String,
39 args: HashMap<String, serde_json::Value>,
40 },
41 ToolResult {
43 tool_call_id: String,
44 content: String,
45 },
46}
47
48pub trait BaseMessage: Send + Sync {
50 fn content(&self) -> &MessageContent;
52
53 fn message_type(&self) -> &str;
55
56 fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value>;
58
59 fn response_metadata(&self) -> &HashMap<String, serde_json::Value>;
61
62 fn name(&self) -> Option<&str>;
64
65 fn id(&self) -> Option<&str>;
67
68 fn text(&self) -> String {
70 match self.content() {
71 MessageContent::Text(text) => text.clone(),
72 MessageContent::Blocks(blocks) => blocks
73 .iter()
74 .filter_map(|block| match block {
75 ContentBlock::Text { text } => Some(text.clone()),
76 ContentBlock::ToolResult { content, .. } => Some(content.clone()),
77 _ => None,
78 })
79 .collect::<Vec<_>>()
80 .join(""),
81 }
82 }
83
84 fn is_human(&self) -> bool {
86 self.message_type() == "human"
87 }
88
89 fn is_ai(&self) -> bool {
91 self.message_type() == "ai"
92 }
93
94 fn is_system(&self) -> bool {
96 self.message_type() == "system"
97 }
98
99 fn is_tool(&self) -> bool {
101 self.message_type() == "tool"
102 }
103}
104
105#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
107pub struct HumanMessage {
108 pub content: MessageContent,
110 #[serde(default)]
112 pub additional_kwargs: HashMap<String, serde_json::Value>,
113 #[serde(default)]
115 pub response_metadata: HashMap<String, serde_json::Value>,
116 pub name: Option<String>,
118 pub id: Option<String>,
120}
121
122impl HumanMessage {
123 pub fn new(content: impl Into<String>) -> Self {
125 Self {
126 content: MessageContent::Text(content.into()),
127 additional_kwargs: HashMap::new(),
128 response_metadata: HashMap::new(),
129 name: None,
130 id: Some(Uuid::new_v4().to_string()),
131 }
132 }
133
134 pub fn new_with_blocks(content: Vec<ContentBlock>) -> Self {
136 Self {
137 content: MessageContent::Blocks(content),
138 additional_kwargs: HashMap::new(),
139 response_metadata: HashMap::new(),
140 name: None,
141 id: Some(Uuid::new_v4().to_string()),
142 }
143 }
144}
145
146impl BaseMessage for HumanMessage {
147 fn content(&self) -> &MessageContent {
148 &self.content
149 }
150
151 fn message_type(&self) -> &str {
152 "human"
153 }
154
155 fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value> {
156 &self.additional_kwargs
157 }
158
159 fn response_metadata(&self) -> &HashMap<String, serde_json::Value> {
160 &self.response_metadata
161 }
162
163 fn name(&self) -> Option<&str> {
164 self.name.as_deref()
165 }
166
167 fn id(&self) -> Option<&str> {
168 self.id.as_deref()
169 }
170}
171
172impl_serializable!(HumanMessage, ["ferriclink", "messages", "human"]);
173
174#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
176pub struct AIMessage {
177 pub content: MessageContent,
179 #[serde(default)]
181 pub additional_kwargs: HashMap<String, serde_json::Value>,
182 #[serde(default)]
184 pub response_metadata: HashMap<String, serde_json::Value>,
185 pub name: Option<String>,
187 pub id: Option<String>,
189}
190
191impl AIMessage {
192 pub fn new(content: impl Into<String>) -> Self {
194 Self {
195 content: MessageContent::Text(content.into()),
196 additional_kwargs: HashMap::new(),
197 response_metadata: HashMap::new(),
198 name: None,
199 id: Some(Uuid::new_v4().to_string()),
200 }
201 }
202
203 pub fn new_with_blocks(content: Vec<ContentBlock>) -> Self {
205 Self {
206 content: MessageContent::Blocks(content),
207 additional_kwargs: HashMap::new(),
208 response_metadata: HashMap::new(),
209 name: None,
210 id: Some(Uuid::new_v4().to_string()),
211 }
212 }
213}
214
215impl BaseMessage for AIMessage {
216 fn content(&self) -> &MessageContent {
217 &self.content
218 }
219
220 fn message_type(&self) -> &str {
221 "ai"
222 }
223
224 fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value> {
225 &self.additional_kwargs
226 }
227
228 fn response_metadata(&self) -> &HashMap<String, serde_json::Value> {
229 &self.response_metadata
230 }
231
232 fn name(&self) -> Option<&str> {
233 self.name.as_deref()
234 }
235
236 fn id(&self) -> Option<&str> {
237 self.id.as_deref()
238 }
239}
240
241impl_serializable!(AIMessage, ["ferriclink", "messages", "ai"]);
242
243#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
245pub struct SystemMessage {
246 pub content: MessageContent,
248 #[serde(default)]
250 pub additional_kwargs: HashMap<String, serde_json::Value>,
251 #[serde(default)]
253 pub response_metadata: HashMap<String, serde_json::Value>,
254 pub name: Option<String>,
256 pub id: Option<String>,
258}
259
260impl SystemMessage {
261 pub fn new(content: impl Into<String>) -> Self {
263 Self {
264 content: MessageContent::Text(content.into()),
265 additional_kwargs: HashMap::new(),
266 response_metadata: HashMap::new(),
267 name: None,
268 id: Some(Uuid::new_v4().to_string()),
269 }
270 }
271}
272
273impl BaseMessage for SystemMessage {
274 fn content(&self) -> &MessageContent {
275 &self.content
276 }
277
278 fn message_type(&self) -> &str {
279 "system"
280 }
281
282 fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value> {
283 &self.additional_kwargs
284 }
285
286 fn response_metadata(&self) -> &HashMap<String, serde_json::Value> {
287 &self.response_metadata
288 }
289
290 fn name(&self) -> Option<&str> {
291 self.name.as_deref()
292 }
293
294 fn id(&self) -> Option<&str> {
295 self.id.as_deref()
296 }
297}
298
299impl_serializable!(SystemMessage, ["ferriclink", "messages", "system"]);
300
301#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
303pub struct ToolMessage {
304 pub content: MessageContent,
306 pub tool_call_id: String,
308 #[serde(default)]
310 pub additional_kwargs: HashMap<String, serde_json::Value>,
311 #[serde(default)]
313 pub response_metadata: HashMap<String, serde_json::Value>,
314 pub name: Option<String>,
316 pub id: Option<String>,
318}
319
320impl ToolMessage {
321 pub fn new(content: impl Into<String>, tool_call_id: impl Into<String>) -> Self {
323 Self {
324 content: MessageContent::Text(content.into()),
325 tool_call_id: tool_call_id.into(),
326 additional_kwargs: HashMap::new(),
327 response_metadata: HashMap::new(),
328 name: None,
329 id: Some(Uuid::new_v4().to_string()),
330 }
331 }
332}
333
334impl BaseMessage for ToolMessage {
335 fn content(&self) -> &MessageContent {
336 &self.content
337 }
338
339 fn message_type(&self) -> &str {
340 "tool"
341 }
342
343 fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value> {
344 &self.additional_kwargs
345 }
346
347 fn response_metadata(&self) -> &HashMap<String, serde_json::Value> {
348 &self.response_metadata
349 }
350
351 fn name(&self) -> Option<&str> {
352 self.name.as_deref()
353 }
354
355 fn id(&self) -> Option<&str> {
356 self.id.as_deref()
357 }
358}
359
360impl_serializable!(ToolMessage, ["ferriclink", "messages", "tool"]);
361
362#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
364#[serde(tag = "type", content = "content")]
365pub enum AnyMessage {
366 Human(HumanMessage),
368 AI(AIMessage),
370 System(SystemMessage),
372 Tool(ToolMessage),
374}
375
376impl AnyMessage {
377 pub fn human(content: impl Into<String>) -> Self {
379 Self::Human(HumanMessage::new(content))
380 }
381
382 pub fn ai(content: impl Into<String>) -> Self {
384 Self::AI(AIMessage::new(content))
385 }
386
387 pub fn system(content: impl Into<String>) -> Self {
389 Self::System(SystemMessage::new(content))
390 }
391
392 pub fn tool(content: impl Into<String>, tool_call_id: impl Into<String>) -> Self {
394 Self::Tool(ToolMessage::new(content, tool_call_id))
395 }
396}
397
398impl BaseMessage for AnyMessage {
399 fn content(&self) -> &MessageContent {
400 match self {
401 AnyMessage::Human(msg) => msg.content(),
402 AnyMessage::AI(msg) => msg.content(),
403 AnyMessage::System(msg) => msg.content(),
404 AnyMessage::Tool(msg) => msg.content(),
405 }
406 }
407
408 fn message_type(&self) -> &str {
409 match self {
410 AnyMessage::Human(msg) => msg.message_type(),
411 AnyMessage::AI(msg) => msg.message_type(),
412 AnyMessage::System(msg) => msg.message_type(),
413 AnyMessage::Tool(msg) => msg.message_type(),
414 }
415 }
416
417 fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value> {
418 match self {
419 AnyMessage::Human(msg) => msg.additional_kwargs(),
420 AnyMessage::AI(msg) => msg.additional_kwargs(),
421 AnyMessage::System(msg) => msg.additional_kwargs(),
422 AnyMessage::Tool(msg) => msg.additional_kwargs(),
423 }
424 }
425
426 fn response_metadata(&self) -> &HashMap<String, serde_json::Value> {
427 match self {
428 AnyMessage::Human(msg) => msg.response_metadata(),
429 AnyMessage::AI(msg) => msg.response_metadata(),
430 AnyMessage::System(msg) => msg.response_metadata(),
431 AnyMessage::Tool(msg) => msg.response_metadata(),
432 }
433 }
434
435 fn name(&self) -> Option<&str> {
436 match self {
437 AnyMessage::Human(msg) => msg.name(),
438 AnyMessage::AI(msg) => msg.name(),
439 AnyMessage::System(msg) => msg.name(),
440 AnyMessage::Tool(msg) => msg.name(),
441 }
442 }
443
444 fn id(&self) -> Option<&str> {
445 match self {
446 AnyMessage::Human(msg) => msg.id(),
447 AnyMessage::AI(msg) => msg.id(),
448 AnyMessage::System(msg) => msg.id(),
449 AnyMessage::Tool(msg) => msg.id(),
450 }
451 }
452}
453
454impl_serializable!(AnyMessage, ["ferriclink", "messages", "any"]);
455
456pub fn get_buffer_string(messages: &[AnyMessage], human_prefix: &str, ai_prefix: &str) -> String {
458 let mut buffer = String::new();
459
460 for message in messages {
461 match message {
462 AnyMessage::Human(msg) => {
463 buffer.push_str(human_prefix);
464 buffer.push_str(": ");
465 buffer.push_str(&msg.text());
466 buffer.push('\n');
467 }
468 AnyMessage::AI(msg) => {
469 buffer.push_str(ai_prefix);
470 buffer.push_str(": ");
471 buffer.push_str(&msg.text());
472 buffer.push('\n');
473 }
474 AnyMessage::System(msg) => {
475 buffer.push_str("System: ");
476 buffer.push_str(&msg.text());
477 buffer.push('\n');
478 }
479 AnyMessage::Tool(msg) => {
480 buffer.push_str("Tool: ");
481 buffer.push_str(&msg.text());
482 buffer.push('\n');
483 }
484 }
485 }
486
487 buffer
488}
489
490#[cfg(test)]
491mod tests {
492 use super::*;
493 use crate::serializable::Serializable;
494
495 #[test]
496 fn test_human_message() {
497 let msg = HumanMessage::new("Hello, world!");
498 assert_eq!(msg.text(), "Hello, world!");
499 assert!(msg.is_human());
500 assert!(!msg.is_ai());
501 assert!(!msg.is_system());
502 assert!(!msg.is_tool());
503 }
504
505 #[test]
506 fn test_ai_message() {
507 let msg = AIMessage::new("Hello! How can I help you?");
508 assert_eq!(msg.text(), "Hello! How can I help you?");
509 assert!(!msg.is_human());
510 assert!(msg.is_ai());
511 assert!(!msg.is_system());
512 assert!(!msg.is_tool());
513 }
514
515 #[test]
516 fn test_system_message() {
517 let msg = SystemMessage::new("You are a helpful assistant.");
518 assert_eq!(msg.text(), "You are a helpful assistant.");
519 assert!(!msg.is_human());
520 assert!(!msg.is_ai());
521 assert!(msg.is_system());
522 assert!(!msg.is_tool());
523 }
524
525 #[test]
526 fn test_tool_message() {
527 let msg = ToolMessage::new("Tool result", "call_123");
528 assert_eq!(msg.text(), "Tool result");
529 assert!(!msg.is_human());
530 assert!(!msg.is_ai());
531 assert!(!msg.is_system());
532 assert!(msg.is_tool());
533 }
534
535 #[test]
536 fn test_any_message() {
537 let human = AnyMessage::human("Hello");
538 let ai = AnyMessage::ai("Hi there!");
539
540 assert!(human.is_human());
541 assert!(ai.is_ai());
542 }
543
544 #[test]
545 fn test_message_content_blocks() {
546 let blocks = vec![
547 ContentBlock::Text {
548 text: "Hello".to_string(),
549 },
550 ContentBlock::Image {
551 image_url: "https://example.com/image.jpg".to_string(),
552 alt_text: Some("An image".to_string()),
553 },
554 ];
555
556 let msg = HumanMessage::new_with_blocks(blocks);
557 assert_eq!(msg.text(), "Hello");
558 }
559
560 #[test]
561 fn test_get_buffer_string() {
562 let messages = vec![AnyMessage::human("Hello"), AnyMessage::ai("Hi there!")];
563
564 let buffer = get_buffer_string(&messages, "Human", "Assistant");
565 assert!(buffer.contains("Human: Hello"));
566 assert!(buffer.contains("Assistant: Hi there!"));
567 }
568
569 #[test]
570 fn test_serialization() {
571 let msg = HumanMessage::new("Test message");
572 let json = msg.to_json().unwrap();
573 let deserialized: HumanMessage = HumanMessage::from_json(&json).unwrap();
574 assert_eq!(msg, deserialized);
575 }
576}