1use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
11#[serde(rename_all = "lowercase")]
12pub enum MessageRole {
13 System,
15 User,
17 Assistant,
19 Tool,
21}
22
23impl MessageRole {
24 pub fn can_initiate(&self) -> bool {
26 matches!(self, MessageRole::System | MessageRole::User)
27 }
28
29 pub fn can_respond(&self) -> bool {
31 matches!(self, MessageRole::Assistant | MessageRole::Tool)
32 }
33}
34
35impl std::fmt::Display for MessageRole {
36 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37 match self {
38 MessageRole::System => write!(f, "system"),
39 MessageRole::User => write!(f, "user"),
40 MessageRole::Assistant => write!(f, "assistant"),
41 MessageRole::Tool => write!(f, "tool"),
42 }
43 }
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
48#[serde(untagged)]
49pub enum MessageContent {
50 Text(String),
52
53 MultiModal {
55 text: Option<String>,
56 attachments: Vec<ContentAttachment>,
57 },
58}
59
60impl MessageContent {
61 pub fn text(content: impl Into<String>) -> Self {
63 Self::Text(content.into())
64 }
65
66 pub fn multi_modal(text: impl Into<String>) -> Self {
68 Self::MultiModal {
69 text: Some(text.into()),
70 attachments: Vec::new(),
71 }
72 }
73
74 pub fn with_attachment(mut self, attachment: ContentAttachment) -> Self {
76 match &mut self {
77 Self::MultiModal { attachments, .. } => {
78 attachments.push(attachment);
79 }
80 Self::Text(text) => {
81 let text = text.clone();
82 self = Self::MultiModal {
83 text: Some(text),
84 attachments: vec![attachment],
85 };
86 }
87 }
88 self
89 }
90
91 pub fn text_content(&self) -> Option<&str> {
93 match self {
94 Self::Text(text) => Some(text),
95 Self::MultiModal { text, .. } => text.as_deref(),
96 }
97 }
98
99 pub fn attachments(&self) -> &[ContentAttachment] {
101 match self {
102 Self::Text(_) => &[],
103 Self::MultiModal { attachments, .. } => attachments,
104 }
105 }
106
107 pub fn is_empty(&self) -> bool {
109 match self {
110 Self::Text(text) => text.is_empty(),
111 Self::MultiModal { text, attachments } => {
112 text.as_ref().map_or(true, |t| t.is_empty()) && attachments.is_empty()
113 }
114 }
115 }
116}
117
118impl From<String> for MessageContent {
119 fn from(text: String) -> Self {
120 Self::Text(text)
121 }
122}
123
124impl From<&str> for MessageContent {
125 fn from(text: &str) -> Self {
126 Self::Text(text.to_string())
127 }
128}
129
130#[derive(Debug, Clone, Serialize, Deserialize)]
132pub struct ContentAttachment {
133 pub attachment_type: AttachmentType,
135
136 pub content: AttachmentContent,
138
139 pub metadata: Option<HashMap<String, serde_json::Value>>,
141}
142
143impl ContentAttachment {
144 pub fn image_base64(mime_type: impl Into<String>, data: impl Into<String>) -> Self {
146 Self {
147 attachment_type: AttachmentType::Image,
148 content: AttachmentContent::Base64 {
149 mime_type: mime_type.into(),
150 data: data.into(),
151 },
152 metadata: None,
153 }
154 }
155
156 pub fn image_url(url: impl Into<String>) -> Self {
158 Self {
159 attachment_type: AttachmentType::Image,
160 content: AttachmentContent::Url { url: url.into() },
161 metadata: None,
162 }
163 }
164
165 pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
167 self.metadata
168 .get_or_insert_with(HashMap::new)
169 .insert(key.into(), value);
170 self
171 }
172}
173
174#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
176#[serde(rename_all = "lowercase")]
177pub enum AttachmentType {
178 Image,
180 Audio,
182 Video,
184 Document,
186 Other,
188}
189
190#[derive(Debug, Clone, Serialize, Deserialize)]
192#[serde(tag = "type", rename_all = "lowercase")]
193pub enum AttachmentContent {
194 Base64 { mime_type: String, data: String },
196
197 Url { url: String },
199
200 #[serde(skip)]
202 Bytes { mime_type: String, data: Vec<u8> },
203}
204
205#[derive(Debug, Clone, Serialize, Deserialize)]
207pub struct ChatMessage {
208 pub role: MessageRole,
210
211 pub content: MessageContent,
213
214 pub name: Option<String>,
216
217 pub tool_calls: Option<Vec<ToolCall>>,
219
220 pub tool_call_id: Option<String>,
222
223 pub metadata: HashMap<String, serde_json::Value>,
225
226 #[serde(with = "chrono::serde::ts_seconds_option")]
228 pub timestamp: Option<chrono::DateTime<chrono::Utc>>,
229}
230
231impl ChatMessage {
232 pub fn new(role: MessageRole, content: impl Into<MessageContent>) -> Self {
234 Self {
235 role,
236 content: content.into(),
237 name: None,
238 tool_calls: None,
239 tool_call_id: None,
240 metadata: HashMap::new(),
241 timestamp: Some(chrono::Utc::now()),
242 }
243 }
244
245 pub fn system(content: impl Into<MessageContent>) -> Self {
247 Self::new(MessageRole::System, content)
248 }
249
250 pub fn user(content: impl Into<MessageContent>) -> Self {
252 Self::new(MessageRole::User, content)
253 }
254
255 pub fn assistant(content: impl Into<MessageContent>) -> Self {
257 Self::new(MessageRole::Assistant, content)
258 }
259
260 pub fn tool(tool_call_id: impl Into<String>, content: impl Into<MessageContent>) -> Self {
262 Self {
263 role: MessageRole::Tool,
264 content: content.into(),
265 name: None,
266 tool_calls: None,
267 tool_call_id: Some(tool_call_id.into()),
268 metadata: HashMap::new(),
269 timestamp: Some(chrono::Utc::now()),
270 }
271 }
272
273 pub fn with_name(mut self, name: impl Into<String>) -> Self {
275 self.name = Some(name.into());
276 self
277 }
278
279 pub fn with_tool_calls(mut self, tool_calls: Vec<ToolCall>) -> Self {
281 self.tool_calls = Some(tool_calls);
282 self
283 }
284
285 pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
287 self.metadata.insert(key.into(), value);
288 self
289 }
290
291 pub fn text(&self) -> Option<&str> {
293 self.content.text_content()
294 }
295
296 pub fn is_empty(&self) -> bool {
298 self.content.is_empty()
299 }
300
301 pub fn len(&self) -> usize {
303 self.text().map_or(0, |t| t.len())
304 }
305}
306
307#[derive(Debug, Clone, Serialize, Deserialize)]
309pub struct ToolCall {
310 pub id: String,
312
313 #[serde(rename = "type")]
315 pub call_type: ToolCallType,
316
317 pub function: ToolFunction,
319}
320
321impl ToolCall {
322 pub fn function(
324 id: impl Into<String>,
325 name: impl Into<String>,
326 arguments: serde_json::Value,
327 ) -> Self {
328 Self {
329 id: id.into(),
330 call_type: ToolCallType::Function,
331 function: ToolFunction {
332 name: name.into(),
333 arguments,
334 },
335 }
336 }
337}
338
339#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
341#[serde(rename_all = "lowercase")]
342pub enum ToolCallType {
343 Function,
345}
346
347#[derive(Debug, Clone, Serialize, Deserialize)]
349pub struct ToolFunction {
350 pub name: String,
352
353 pub arguments: serde_json::Value,
355}
356
357pub struct MessageBuilder {
359 message: ChatMessage,
360}
361
362impl MessageBuilder {
363 pub fn new(role: MessageRole) -> Self {
365 Self {
366 message: ChatMessage {
367 role,
368 content: MessageContent::Text(String::new()),
369 name: None,
370 tool_calls: None,
371 tool_call_id: None,
372 metadata: HashMap::new(),
373 timestamp: Some(chrono::Utc::now()),
374 },
375 }
376 }
377
378 pub fn content(mut self, content: impl Into<MessageContent>) -> Self {
380 self.message.content = content.into();
381 self
382 }
383
384 pub fn name(mut self, name: impl Into<String>) -> Self {
386 self.message.name = Some(name.into());
387 self
388 }
389
390 pub fn tool_calls(mut self, tool_calls: Vec<ToolCall>) -> Self {
392 self.message.tool_calls = Some(tool_calls);
393 self
394 }
395
396 pub fn tool_call_id(mut self, tool_call_id: impl Into<String>) -> Self {
398 self.message.tool_call_id = Some(tool_call_id.into());
399 self
400 }
401
402 pub fn metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
404 self.message.metadata.insert(key.into(), value);
405 self
406 }
407
408 pub fn build(self) -> ChatMessage {
410 self.message
411 }
412}
413
414impl Default for MessageBuilder {
415 fn default() -> Self {
416 Self::new(MessageRole::User)
417 }
418}
419
420#[cfg(test)]
421mod tests {
422 use super::*;
423
424 #[test]
425 fn test_message_creation() {
426 let msg = ChatMessage::user("Hello, world!");
427 assert_eq!(msg.role, MessageRole::User);
428 assert_eq!(msg.text(), Some("Hello, world!"));
429 assert!(!msg.is_empty());
430 }
431
432 #[test]
433 fn test_message_builder() {
434 let msg = MessageBuilder::new(MessageRole::Assistant)
435 .content("Hello there!")
436 .name("Assistant")
437 .metadata("source", serde_json::Value::String("test".to_string()))
438 .build();
439
440 assert_eq!(msg.role, MessageRole::Assistant);
441 assert_eq!(msg.text(), Some("Hello there!"));
442 assert_eq!(msg.name, Some("Assistant".to_string()));
443 assert!(msg.metadata.contains_key("source"));
444 }
445
446 #[test]
447 fn test_multi_modal_content() {
448 let content = MessageContent::multi_modal("Check this image").with_attachment(
449 ContentAttachment::image_url("https://example.com/image.jpg"),
450 );
451
452 assert_eq!(content.text_content(), Some("Check this image"));
453 assert_eq!(content.attachments().len(), 1);
454 }
455}