1use std::collections::HashMap;
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
14pub struct Message {
15 pub role: Role,
17
18 pub parts: Vec<MessagePart>,
20
21 #[serde(rename = "messageId", skip_serializing_if = "Option::is_none")]
23 pub message_id: Option<String>,
24
25 #[serde(rename = "taskId", skip_serializing_if = "Option::is_none")]
27 pub task_id: Option<String>,
28
29 #[serde(rename = "contextId", skip_serializing_if = "Option::is_none")]
31 pub context_id: Option<String>,
32
33 #[serde(skip_serializing_if = "Option::is_none")]
35 pub metadata: Option<HashMap<String, Value>>,
36
37 #[serde(skip_serializing_if = "Option::is_none")]
39 pub extensions: Option<HashMap<String, Value>>,
40}
41
42impl Message {
43 pub fn new(role: Role, text: impl Into<String>) -> Self {
45 Self {
46 role,
47 parts: vec![MessagePart::Text { text: text.into() }],
48 message_id: None,
49 task_id: None,
50 context_id: None,
51 metadata: None,
52 extensions: None,
53 }
54 }
55
56 pub fn user(text: impl Into<String>) -> Self {
58 Self::new(Role::User, text)
59 }
60
61 pub fn assistant(text: impl Into<String>) -> Self {
63 Self::new(Role::Assistant, text)
64 }
65
66 pub fn builder() -> MessageBuilder {
68 MessageBuilder::new()
69 }
70
71 pub fn with_metadata(mut self, key: impl Into<String>, value: Value) -> Self {
73 self.metadata
74 .get_or_insert_with(HashMap::new)
75 .insert(key.into(), value);
76 self
77 }
78
79 pub fn with_extension(mut self, key: impl Into<String>, value: Value) -> Self {
81 self.extensions
82 .get_or_insert_with(HashMap::new)
83 .insert(key.into(), value);
84 self
85 }
86
87 pub fn with_part(mut self, part: MessagePart) -> Self {
89 self.parts.push(part);
90 self
91 }
92}
93
94#[derive(Debug, Default)]
96pub struct MessageBuilder {
97 role: Option<Role>,
98 parts: Vec<MessagePart>,
99 message_id: Option<String>,
100 task_id: Option<String>,
101 context_id: Option<String>,
102 metadata: Option<HashMap<String, Value>>,
103 extensions: Option<HashMap<String, Value>>,
104}
105
106impl MessageBuilder {
107 pub fn new() -> Self {
109 Self::default()
110 }
111
112 pub fn role(mut self, role: Role) -> Self {
114 self.role = Some(role);
115 self
116 }
117
118 pub fn parts(mut self, parts: Vec<MessagePart>) -> Self {
120 self.parts = parts;
121 self
122 }
123
124 pub fn part(mut self, part: MessagePart) -> Self {
126 self.parts.push(part);
127 self
128 }
129
130 pub fn message_id(mut self, id: impl Into<String>) -> Self {
132 self.message_id = Some(id.into());
133 self
134 }
135
136 pub fn task_id(mut self, id: impl Into<String>) -> Self {
138 self.task_id = Some(id.into());
139 self
140 }
141
142 pub fn context_id(mut self, id: impl Into<String>) -> Self {
144 self.context_id = Some(id.into());
145 self
146 }
147
148 pub fn metadata(mut self, key: impl Into<String>, value: Value) -> Self {
150 self.metadata
151 .get_or_insert_with(HashMap::new)
152 .insert(key.into(), value);
153 self
154 }
155
156 pub fn extension(mut self, key: impl Into<String>, value: Value) -> Self {
158 self.extensions
159 .get_or_insert_with(HashMap::new)
160 .insert(key.into(), value);
161 self
162 }
163
164 pub fn build(self) -> Message {
170 let role = self.role.expect("Message role is required");
171 assert!(!self.parts.is_empty(), "Message must have at least one part");
172
173 Message {
174 role,
175 parts: self.parts,
176 message_id: self.message_id,
177 task_id: self.task_id,
178 context_id: self.context_id,
179 metadata: self.metadata,
180 extensions: self.extensions,
181 }
182 }
183}
184
185#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
187#[serde(rename_all = "lowercase")]
188pub enum Role {
189 User,
191
192 Assistant,
194
195 System,
197}
198
199#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
203#[serde(untagged)]
204pub enum MessagePart {
205 Text {
207 text: String,
209 },
210
211 File {
213 #[serde(rename = "fileUri")]
215 file_uri: String,
216
217 #[serde(rename = "mimeType", skip_serializing_if = "Option::is_none")]
219 mime_type: Option<String>,
220 },
221
222 Data {
224 data: Value,
226
227 #[serde(rename = "mimeType", skip_serializing_if = "Option::is_none")]
229 mime_type: Option<String>,
230 },
231}
232
233impl MessagePart {
234 pub fn text(text: impl Into<String>) -> Self {
236 Self::Text { text: text.into() }
237 }
238
239 pub fn file(file_uri: impl Into<String>) -> Self {
241 Self::File {
242 file_uri: file_uri.into(),
243 mime_type: None,
244 }
245 }
246
247 pub fn file_with_type(file_uri: impl Into<String>, mime_type: impl Into<String>) -> Self {
249 Self::File {
250 file_uri: file_uri.into(),
251 mime_type: Some(mime_type.into()),
252 }
253 }
254
255 pub fn data(data: Value) -> Self {
257 Self::Data {
258 data,
259 mime_type: None,
260 }
261 }
262
263 pub fn data_with_type(data: Value, mime_type: impl Into<String>) -> Self {
265 Self::Data {
266 data,
267 mime_type: Some(mime_type.into()),
268 }
269 }
270}
271
272#[cfg(test)]
273mod tests {
274 use serde_json::json;
275
276 use super::*;
277
278 #[test]
279 fn test_message_creation() {
280 let msg = Message::user("Hello, agent!");
281 assert_eq!(msg.role, Role::User);
282 assert_eq!(msg.parts.len(), 1);
283
284 match &msg.parts[0] {
285 MessagePart::Text { text } => assert_eq!(text, "Hello, agent!"),
286 _ => panic!("Expected text part"),
287 }
288 }
289
290 #[test]
291 fn test_message_with_metadata() {
292 let msg = Message::user("Test")
293 .with_metadata("key", json!("value"))
294 .with_extension("ext", json!({"enabled": true}));
295
296 assert!(msg.metadata.is_some());
297 assert!(msg.extensions.is_some());
298 }
299
300 #[test]
301 fn test_message_serialization() {
302 let msg = Message::user("Test message");
303 let json = serde_json::to_string(&msg).unwrap();
304 assert!(json.contains("\"role\":\"user\""));
305 assert!(json.contains("\"text\":\"Test message\""));
306
307 let deserialized: Message = serde_json::from_str(&json).unwrap();
308 assert_eq!(msg, deserialized);
309 }
310
311 #[test]
312 fn test_message_part_types() {
313 let text = MessagePart::text("Hello");
314 let file = MessagePart::file("file://path/to/file");
315 let data = MessagePart::data(json!({"key": "value"}));
316
317 assert!(matches!(text, MessagePart::Text { .. }));
318 assert!(matches!(file, MessagePart::File { .. }));
319 assert!(matches!(data, MessagePart::Data { .. }));
320 }
321
322 #[test]
323 fn test_message_builder() {
324 let msg = Message::builder()
325 .role(Role::User)
326 .parts(vec![MessagePart::text("Hello")])
327 .message_id("msg-123")
328 .task_id("task-456")
329 .context_id("ctx-789")
330 .build();
331
332 assert_eq!(msg.role, Role::User);
333 assert_eq!(msg.parts.len(), 1);
334 assert_eq!(msg.message_id, Some("msg-123".to_string()));
335 assert_eq!(msg.task_id, Some("task-456".to_string()));
336 assert_eq!(msg.context_id, Some("ctx-789".to_string()));
337 }
338
339 #[test]
340 fn test_message_builder_with_part() {
341 let msg = Message::builder()
342 .role(Role::Assistant)
343 .part(MessagePart::text("First"))
344 .part(MessagePart::text("Second"))
345 .build();
346
347 assert_eq!(msg.parts.len(), 2);
348 }
349
350 #[test]
351 #[should_panic(expected = "Message role is required")]
352 fn test_message_builder_missing_role() {
353 Message::builder()
354 .parts(vec![MessagePart::text("Hello")])
355 .build();
356 }
357
358 #[test]
359 #[should_panic(expected = "Message must have at least one part")]
360 fn test_message_builder_no_parts() {
361 Message::builder().role(Role::User).build();
362 }
363
364 #[test]
365 fn test_message_serialization_with_ids() {
366 let msg = Message::builder()
367 .role(Role::User)
368 .parts(vec![MessagePart::text("Test")])
369 .message_id("msg-123")
370 .task_id("task-456")
371 .build();
372
373 let json = serde_json::to_string(&msg).unwrap();
374 assert!(json.contains("\"messageId\":\"msg-123\""));
375 assert!(json.contains("\"taskId\":\"task-456\""));
376
377 let deserialized: Message = serde_json::from_str(&json).unwrap();
378 assert_eq!(msg, deserialized);
379 }
380}