ferriclink_core/
messages.rs

1//! Message types for FerricLink Core
2//!
3//! This module provides the core message abstractions used throughout the FerricLink
4//! ecosystem, similar to LangChain's message system.
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use uuid::Uuid;
9
10use crate::impl_serializable;
11
12/// Content of a message, which can be either text or a list of content blocks
13#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
14#[serde(untagged)]
15pub enum MessageContent {
16    /// Simple text content
17    Text(String),
18    /// Complex content with multiple blocks
19    Blocks(Vec<ContentBlock>),
20}
21
22/// A content block within a message
23#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
24#[serde(tag = "type")]
25pub enum ContentBlock {
26    /// Text content block
27    Text { text: String },
28    /// Image content block
29    Image {
30        image_url: String,
31        alt_text: Option<String>,
32    },
33    /// JSON content block
34    Json { data: serde_json::Value },
35    /// Tool call content block
36    ToolCall {
37        id: String,
38        name: String,
39        args: HashMap<String, serde_json::Value>,
40    },
41    /// Tool result content block
42    ToolResult {
43        tool_call_id: String,
44        content: String,
45    },
46}
47
48/// Base message trait that all message types implement
49pub trait BaseMessage: Send + Sync {
50    /// Get the content of the message
51    fn content(&self) -> &MessageContent;
52
53    /// Get the message type
54    fn message_type(&self) -> &str;
55
56    /// Get additional kwargs
57    fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value>;
58
59    /// Get response metadata
60    fn response_metadata(&self) -> &HashMap<String, serde_json::Value>;
61
62    /// Get the message name
63    fn name(&self) -> Option<&str>;
64
65    /// Get the message ID
66    fn id(&self) -> Option<&str>;
67
68    /// Convert the message content to text
69    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    /// Check if this is a human message
85    fn is_human(&self) -> bool {
86        self.message_type() == "human"
87    }
88
89    /// Check if this is an AI message
90    fn is_ai(&self) -> bool {
91        self.message_type() == "ai"
92    }
93
94    /// Check if this is a system message
95    fn is_system(&self) -> bool {
96        self.message_type() == "system"
97    }
98
99    /// Check if this is a tool message
100    fn is_tool(&self) -> bool {
101        self.message_type() == "tool"
102    }
103}
104
105/// Human message
106#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
107pub struct HumanMessage {
108    /// The content of the message
109    pub content: MessageContent,
110    /// Additional keyword arguments
111    #[serde(default)]
112    pub additional_kwargs: HashMap<String, serde_json::Value>,
113    /// Response metadata
114    #[serde(default)]
115    pub response_metadata: HashMap<String, serde_json::Value>,
116    /// Optional name for the message
117    pub name: Option<String>,
118    /// Optional unique identifier
119    pub id: Option<String>,
120}
121
122impl HumanMessage {
123    /// Create a new human message with text content
124    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    /// Create a new human message with content blocks
135    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/// AI message
175#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
176pub struct AIMessage {
177    /// The content of the message
178    pub content: MessageContent,
179    /// Additional keyword arguments
180    #[serde(default)]
181    pub additional_kwargs: HashMap<String, serde_json::Value>,
182    /// Response metadata
183    #[serde(default)]
184    pub response_metadata: HashMap<String, serde_json::Value>,
185    /// Optional name for the message
186    pub name: Option<String>,
187    /// Optional unique identifier
188    pub id: Option<String>,
189}
190
191impl AIMessage {
192    /// Create a new AI message with text content
193    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    /// Create a new AI message with content blocks
204    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/// System message
244#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
245pub struct SystemMessage {
246    /// The content of the message
247    pub content: MessageContent,
248    /// Additional keyword arguments
249    #[serde(default)]
250    pub additional_kwargs: HashMap<String, serde_json::Value>,
251    /// Response metadata
252    #[serde(default)]
253    pub response_metadata: HashMap<String, serde_json::Value>,
254    /// Optional name for the message
255    pub name: Option<String>,
256    /// Optional unique identifier
257    pub id: Option<String>,
258}
259
260impl SystemMessage {
261    /// Create a new system message with text content
262    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/// Tool message
302#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
303pub struct ToolMessage {
304    /// The content of the message
305    pub content: MessageContent,
306    /// The tool call ID this message is responding to
307    pub tool_call_id: String,
308    /// Additional keyword arguments
309    #[serde(default)]
310    pub additional_kwargs: HashMap<String, serde_json::Value>,
311    /// Response metadata
312    #[serde(default)]
313    pub response_metadata: HashMap<String, serde_json::Value>,
314    /// Optional name for the message
315    pub name: Option<String>,
316    /// Optional unique identifier
317    pub id: Option<String>,
318}
319
320impl ToolMessage {
321    /// Create a new tool message
322    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/// Union type for all message types
363#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
364#[serde(tag = "type", content = "content")]
365pub enum AnyMessage {
366    /// Human message
367    Human(HumanMessage),
368    /// AI message
369    AI(AIMessage),
370    /// System message
371    System(SystemMessage),
372    /// Tool message
373    Tool(ToolMessage),
374}
375
376impl AnyMessage {
377    /// Create a human message
378    pub fn human(content: impl Into<String>) -> Self {
379        Self::Human(HumanMessage::new(content))
380    }
381
382    /// Create an AI message
383    pub fn ai(content: impl Into<String>) -> Self {
384        Self::AI(AIMessage::new(content))
385    }
386
387    /// Create a system message
388    pub fn system(content: impl Into<String>) -> Self {
389        Self::System(SystemMessage::new(content))
390    }
391
392    /// Create a tool message
393    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
456/// Helper function to convert messages to a string representation
457pub 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}