1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
5pub enum Message {
6 Human(HumanMessage),
8 Assistant(AssistantMessage),
10 System(SystemMessage),
12}
13
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
16pub struct HumanMessage {
17 pub content: MessageContent,
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
22pub struct AssistantMessage {
23 pub content: MessageContent,
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
28pub struct SystemMessage {
29 pub content: String,
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
34pub enum MessageContent {
35 Text(String),
37 Mixed {
39 text: Option<String>,
40 tool_calls: Vec<ToolCall>,
41 },
42}
43
44impl MessageContent {
45 pub fn text<S: Into<String>>(text: S) -> Self {
47 Self::Text(text.into())
48 }
49
50 pub fn mixed<S: Into<String>>(text: Option<S>, tool_calls: Vec<ToolCall>) -> Self {
52 Self::Mixed {
53 text: text.map(|t| t.into()),
54 tool_calls,
55 }
56 }
57
58 pub fn as_text(&self) -> Option<&str> {
60 match self {
61 Self::Text(text) => Some(text),
62 Self::Mixed { text, .. } => text.as_deref(),
63 }
64 }
65
66 pub fn to_text(&self) -> String {
68 match self {
69 Self::Text(text) => text.clone(),
70 Self::Mixed { text, .. } => text.clone().unwrap_or_default(),
71 }
72 }
73
74 pub fn has_tool_calls(&self) -> bool {
76 matches!(self, Self::Mixed { tool_calls, .. } if !tool_calls.is_empty())
77 }
78
79 pub fn tool_calls(&self) -> &[ToolCall] {
81 match self {
82 Self::Text(_) => &[],
83 Self::Mixed { tool_calls, .. } => tool_calls,
84 }
85 }
86}
87
88impl From<String> for MessageContent {
89 fn from(text: String) -> Self {
90 Self::Text(text)
91 }
92}
93
94impl From<&str> for MessageContent {
95 fn from(text: &str) -> Self {
96 Self::Text(text.to_string())
97 }
98}
99
100impl Message {
101 pub fn human<S: Into<String>>(content: S) -> Self {
103 Self::Human(HumanMessage {
104 content: MessageContent::text(content),
105 })
106 }
107
108 pub fn assistant<S: Into<String>>(content: S) -> Self {
110 Self::Assistant(AssistantMessage {
111 content: MessageContent::text(content),
112 })
113 }
114
115 pub fn system<S: Into<String>>(content: S) -> Self {
117 Self::System(SystemMessage {
118 content: content.into(),
119 })
120 }
121
122 pub fn role(&self) -> &'static str {
124 match self {
125 Self::Human(_) => "user",
126 Self::Assistant(_) => "assistant",
127 Self::System(_) => "system",
128 }
129 }
130
131 pub fn content_text(&self) -> String {
133 match self {
134 Self::Human(msg) => msg.content.to_text(),
135 Self::Assistant(msg) => msg.content.to_text(),
136 Self::System(msg) => msg.content.clone(),
137 }
138 }
139}
140
141impl HumanMessage {
142 pub fn new<S: Into<String>>(content: S) -> Self {
144 Self {
145 content: MessageContent::text(content),
146 }
147 }
148
149 pub fn with_tool_calls<S: Into<String>>(text: Option<S>, tool_calls: Vec<ToolCall>) -> Self {
151 Self {
152 content: MessageContent::mixed(text, tool_calls),
153 }
154 }
155}
156
157impl AssistantMessage {
158 pub fn new<S: Into<String>>(content: S) -> Self {
160 Self {
161 content: MessageContent::text(content),
162 }
163 }
164
165 pub fn with_tool_calls<S: Into<String>>(text: Option<S>, tool_calls: Vec<ToolCall>) -> Self {
167 Self {
168 content: MessageContent::mixed(text, tool_calls),
169 }
170 }
171}
172
173impl SystemMessage {
174 pub fn new<S: Into<String>>(content: S) -> Self {
176 Self {
177 content: content.into(),
178 }
179 }
180}
181
182#[cfg(test)]
183mod tests {
184 use super::*;
185
186 #[test]
187 fn test_message_content_text() {
188 let content = MessageContent::text("Hello world");
189 assert_eq!(content.as_text(), Some("Hello world"));
190 assert_eq!(content.to_text(), "Hello world");
191 assert!(!content.has_tool_calls());
192 assert!(content.tool_calls().is_empty());
193 }
194
195 #[test]
196 fn test_message_content_mixed() {
197 let tool_call = ToolCall {
198 id: "call_1".to_string(),
199 call_id: Some("call_1".to_string()),
200 function: ToolFunction {
201 name: "test_function".to_string(),
202 arguments: serde_json::json!({"param": "value"}),
203 },
204 };
205
206 let content = MessageContent::mixed(Some("Hello"), vec![tool_call.clone()]);
207 assert_eq!(content.as_text(), Some("Hello"));
208 assert_eq!(content.to_text(), "Hello");
209 assert!(content.has_tool_calls());
210 assert_eq!(content.tool_calls().len(), 1);
211 assert_eq!(content.tool_calls()[0].id, "call_1");
212 }
213
214 #[test]
215 fn test_message_content_from_string() {
216 let content: MessageContent = "Test message".into();
217 assert_eq!(content.as_text(), Some("Test message"));
218 }
219
220 #[test]
221 fn test_message_constructors() {
222 let human_msg = Message::human("Hello");
223 assert_eq!(human_msg.role(), "user");
224 assert_eq!(human_msg.content_text(), "Hello");
225
226 let assistant_msg = Message::assistant("Hi there");
227 assert_eq!(assistant_msg.role(), "assistant");
228 assert_eq!(assistant_msg.content_text(), "Hi there");
229
230 let system_msg = Message::system("You are helpful");
231 assert_eq!(system_msg.role(), "system");
232 assert_eq!(system_msg.content_text(), "You are helpful");
233 }
234
235 #[test]
236 fn test_human_message_constructors() {
237 let msg = HumanMessage::new("Hello");
238 assert_eq!(msg.content.to_text(), "Hello");
239
240 let tool_call = ToolCall {
241 id: "call_1".to_string(),
242 call_id: None,
243 function: ToolFunction {
244 name: "test".to_string(),
245 arguments: serde_json::json!({}),
246 },
247 };
248
249 let msg_with_tools = HumanMessage::with_tool_calls(Some("Text"), vec![tool_call]);
250 assert_eq!(msg_with_tools.content.to_text(), "Text");
251 assert!(msg_with_tools.content.has_tool_calls());
252 }
253
254 #[test]
255 fn test_assistant_message_constructors() {
256 let msg = AssistantMessage::new("Response");
257 assert_eq!(msg.content.to_text(), "Response");
258
259 let tool_call = ToolCall {
260 id: "call_1".to_string(),
261 call_id: None,
262 function: ToolFunction {
263 name: "test".to_string(),
264 arguments: serde_json::json!({}),
265 },
266 };
267
268 let msg_with_tools = AssistantMessage::with_tool_calls(Some("Response"), vec![tool_call]);
269 assert_eq!(msg_with_tools.content.to_text(), "Response");
270 assert!(msg_with_tools.content.has_tool_calls());
271 }
272
273 #[test]
274 fn test_message_serialization() {
275 let msg = Message::human("Test message");
276 let serialized = serde_json::to_string(&msg).unwrap();
277 let deserialized: Message = serde_json::from_str(&serialized).unwrap();
278
279 match deserialized {
280 Message::Human(human_msg) => {
281 assert_eq!(human_msg.content.to_text(), "Test message");
282 }
283 _ => panic!("Expected human message"),
284 }
285 }
286}
287
288#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
289pub enum HumanContent {
290 Text(Text),
291 ToolCall(ToolCall),
292}
293
294#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
296pub struct Text {
297 pub text: String,
298}
299
300#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
302pub struct ToolCall {
303 pub id: String,
304 pub call_id: Option<String>,
305 pub function: ToolFunction,
306}
307
308#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
310pub struct ToolFunction {
311 pub name: String,
312 pub arguments: serde_json::Value,
313}