agent_base/types/
message.rs

1//! Message types for LLM conversations.
2//!
3//! `ChatMessage`, `ImageAttachment`, `ImageDetail`, and `ToolCallMessage` are
4//! re-exported from `llm-trait`. `MessageRole` is re-exported from
5//! `agent-types`. `Message` remains here as it is agent-runtime specific.
6
7pub use agent_types::MessageRole;
8pub use llm_trait::message::{ChatMessage, ImageAttachment, ImageDetail, ToolCallMessage};
9
10use serde::{Deserialize, Serialize};
11
12#[derive(Clone, Debug, Serialize, Deserialize)]
13pub struct Message {
14    pub role: MessageRole,
15    pub content: String,
16}
17
18impl From<&ChatMessage> for Message {
19    fn from(cm: &ChatMessage) -> Self {
20        match cm {
21            ChatMessage::System { content, .. } => Message {
22                role: MessageRole::System,
23                content: content.clone(),
24            },
25            ChatMessage::User { content, .. } => Message {
26                role: MessageRole::User,
27                content: content.clone(),
28            },
29            ChatMessage::Assistant { content, .. } => Message {
30                role: MessageRole::Assistant,
31                content: content.clone().unwrap_or_default(),
32            },
33            ChatMessage::Tool { content, .. } => Message {
34                role: MessageRole::Tool,
35                content: content.clone(),
36            },
37            ChatMessage::Custom { role: _, data } => Message {
38                role: MessageRole::User,
39                content: data.to_string(),
40            },
41        }
42    }
43}
44
45/// Callback that transforms the message list before it is sent to the LLM.
46///
47/// The default implementation filters out [`ChatMessage::Custom`] variants because
48/// most providers don't understand application-specific message types. Consumers
49/// can override this to inject custom serialization logic for their message types.
50pub type ConvertToLlmFn = std::sync::Arc<dyn Fn(&[ChatMessage]) -> Vec<ChatMessage> + Send + Sync>;
51
52/// Default conversion that strips [`ChatMessage::Custom`] messages.
53pub fn default_convert_to_llm(messages: &[ChatMessage]) -> Vec<ChatMessage> {
54    messages
55        .iter()
56        .filter(|m| !matches!(m, ChatMessage::Custom { .. }))
57        .cloned()
58        .collect()
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    #[test]
66    fn test_default_convert_to_llm_filters_custom() {
67        let messages = vec![
68            ChatMessage::system("You are a helpful assistant."),
69            ChatMessage::user("Hello"),
70            ChatMessage::Custom {
71                role: "artifact".to_string(),
72                data: serde_json::json!({"id": "abc123"}),
73            },
74            ChatMessage::assistant("Hi there!"),
75            ChatMessage::Custom {
76                role: "notification".to_string(),
77                data: serde_json::json!({"level": "info"}),
78            },
79            ChatMessage::tool("call_1", "result"),
80        ];
81
82        let filtered = default_convert_to_llm(&messages);
83
84        assert_eq!(filtered.len(), 4);
85        assert!(matches!(filtered[0], ChatMessage::System { .. }));
86        assert!(matches!(filtered[1], ChatMessage::User { .. }));
87        assert!(matches!(filtered[2], ChatMessage::Assistant { .. }));
88        assert!(matches!(filtered[3], ChatMessage::Tool { .. }));
89    }
90
91    #[test]
92    fn test_default_convert_to_llm_no_custom() {
93        let messages = vec![
94            ChatMessage::system("sys"),
95            ChatMessage::user("usr"),
96            ChatMessage::assistant("asst"),
97        ];
98
99        let filtered = default_convert_to_llm(&messages);
100        assert_eq!(filtered.len(), 3);
101    }
102
103    #[test]
104    fn test_custom_convert_to_llm_preserves_selected() {
105        let messages = vec![
106            ChatMessage::system("sys"),
107            ChatMessage::Custom {
108                role: "artifact".to_string(),
109                data: serde_json::json!({"id": "x"}),
110            },
111            ChatMessage::user("usr"),
112        ];
113
114        let convert = |msgs: &[ChatMessage]| -> Vec<ChatMessage> {
115            msgs.iter()
116                .filter(|m| match m {
117                    ChatMessage::Custom { role, .. } => role == "artifact",
118                    _ => true,
119                })
120                .cloned()
121                .collect()
122        };
123
124        let filtered = convert(&messages);
125        assert_eq!(filtered.len(), 3);
126        assert!(matches!(filtered[0], ChatMessage::System { .. }));
127        assert!(matches!(filtered[1], ChatMessage::Custom { .. }));
128        assert!(matches!(filtered[2], ChatMessage::User { .. }));
129
130        let messages2 = vec![
131            ChatMessage::system("sys"),
132            ChatMessage::Custom {
133                role: "notification".to_string(),
134                data: serde_json::json!({"level": "info"}),
135            },
136        ];
137        let filtered2 = convert(&messages2);
138        assert_eq!(filtered2.len(), 1);
139        assert!(matches!(filtered2[0], ChatMessage::System { .. }));
140    }
141
142    #[test]
143    fn test_custom_message_is_ephemeral_false() {
144        let custom = ChatMessage::Custom {
145            role: "artifact".to_string(),
146            data: serde_json::json!({}),
147        };
148        assert!(!custom.is_ephemeral());
149    }
150
151    #[test]
152    fn test_custom_message_serialization_roundtrip() {
153        let custom = ChatMessage::Custom {
154            role: "artifact".to_string(),
155            data: serde_json::json!({"id": "test-123", "content": "hello"}),
156        };
157        let json_str = serde_json::to_string(&custom).unwrap();
158        let deserialized: ChatMessage = serde_json::from_str(&json_str).unwrap();
159        match deserialized {
160            ChatMessage::Custom { role, data } => {
161                assert_eq!(role, "artifact");
162                assert_eq!(data["id"], "test-123");
163                assert_eq!(data["content"], "hello");
164            }
165            _ => panic!("Expected Custom variant"),
166        }
167    }
168}
agent_base/types/message.rs

agent_base/types/
message.rs
