rexis_llm/
message.rs

1//! # RSLLM Message Types
2//!
3//! Chat message types and content handling for RSLLM.
4//! Supports text, multi-modal content, and role-based messaging.
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9/// Role of a message participant in a conversation
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
11#[serde(rename_all = "lowercase")]
12pub enum MessageRole {
13    /// System message (instructions, context)
14    System,
15    /// User message (human input)
16    User,
17    /// Assistant message (AI response)
18    Assistant,
19    /// Function/tool call message
20    Tool,
21}
22
23impl MessageRole {
24    /// Check if this role can initiate a conversation
25    pub fn can_initiate(&self) -> bool {
26        matches!(self, MessageRole::System | MessageRole::User)
27    }
28
29    /// Check if this role can respond to messages
30    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/// Content of a chat message - supports text and multi-modal content
47#[derive(Debug, Clone, Serialize, Deserialize)]
48#[serde(untagged)]
49pub enum MessageContent {
50    /// Simple text content
51    Text(String),
52
53    /// Multi-modal content (text + images/attachments)
54    MultiModal {
55        text: Option<String>,
56        attachments: Vec<ContentAttachment>,
57    },
58}
59
60impl MessageContent {
61    /// Create text content
62    pub fn text(content: impl Into<String>) -> Self {
63        Self::Text(content.into())
64    }
65
66    /// Create multi-modal content with text
67    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    /// Add an attachment to multi-modal content
75    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    /// Get the text content, if any
92    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    /// Get attachments, if any
100    pub fn attachments(&self) -> &[ContentAttachment] {
101        match self {
102            Self::Text(_) => &[],
103            Self::MultiModal { attachments, .. } => attachments,
104        }
105    }
106
107    /// Check if content is empty
108    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/// Attachment within message content
131#[derive(Debug, Clone, Serialize, Deserialize)]
132pub struct ContentAttachment {
133    /// Type of attachment
134    pub attachment_type: AttachmentType,
135
136    /// Content of the attachment
137    pub content: AttachmentContent,
138
139    /// Optional metadata
140    pub metadata: Option<HashMap<String, serde_json::Value>>,
141}
142
143impl ContentAttachment {
144    /// Create an image attachment from base64 data
145    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    /// Create an image attachment from URL
157    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    /// Add metadata to the attachment
166    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/// Type of content attachment
175#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
176#[serde(rename_all = "lowercase")]
177pub enum AttachmentType {
178    /// Image attachment
179    Image,
180    /// Audio attachment
181    Audio,
182    /// Video attachment
183    Video,
184    /// Document attachment
185    Document,
186    /// Other/custom attachment type
187    Other,
188}
189
190/// Content of an attachment
191#[derive(Debug, Clone, Serialize, Deserialize)]
192#[serde(tag = "type", rename_all = "lowercase")]
193pub enum AttachmentContent {
194    /// Base64-encoded content
195    Base64 { mime_type: String, data: String },
196
197    /// URL reference
198    Url { url: String },
199
200    /// Raw bytes (for internal use)
201    #[serde(skip)]
202    Bytes { mime_type: String, data: Vec<u8> },
203}
204
205/// A chat message in a conversation
206#[derive(Debug, Clone, Serialize, Deserialize)]
207pub struct ChatMessage {
208    /// Role of the message sender
209    pub role: MessageRole,
210
211    /// Content of the message
212    pub content: MessageContent,
213
214    /// Optional name of the sender (for user/assistant disambiguation)
215    pub name: Option<String>,
216
217    /// Tool call information (for assistant messages)
218    pub tool_calls: Option<Vec<ToolCall>>,
219
220    /// Tool call ID (for tool response messages)
221    pub tool_call_id: Option<String>,
222
223    /// Message metadata
224    pub metadata: HashMap<String, serde_json::Value>,
225
226    /// Message timestamp
227    #[serde(with = "chrono::serde::ts_seconds_option")]
228    pub timestamp: Option<chrono::DateTime<chrono::Utc>>,
229}
230
231impl ChatMessage {
232    /// Create a new chat message
233    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    /// Create a system message
246    pub fn system(content: impl Into<MessageContent>) -> Self {
247        Self::new(MessageRole::System, content)
248    }
249
250    /// Create a user message
251    pub fn user(content: impl Into<MessageContent>) -> Self {
252        Self::new(MessageRole::User, content)
253    }
254
255    /// Create an assistant message
256    pub fn assistant(content: impl Into<MessageContent>) -> Self {
257        Self::new(MessageRole::Assistant, content)
258    }
259
260    /// Create a tool response message
261    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    /// Set the sender name
274    pub fn with_name(mut self, name: impl Into<String>) -> Self {
275        self.name = Some(name.into());
276        self
277    }
278
279    /// Add tool calls to the message
280    pub fn with_tool_calls(mut self, tool_calls: Vec<ToolCall>) -> Self {
281        self.tool_calls = Some(tool_calls);
282        self
283    }
284
285    /// Add metadata to the message
286    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    /// Get the text content of the message
292    pub fn text(&self) -> Option<&str> {
293        self.content.text_content()
294    }
295
296    /// Check if message is empty
297    pub fn is_empty(&self) -> bool {
298        self.content.is_empty()
299    }
300
301    /// Get message length in characters
302    pub fn len(&self) -> usize {
303        self.text().map_or(0, |t| t.len())
304    }
305}
306
307/// Tool call information
308#[derive(Debug, Clone, Serialize, Deserialize)]
309pub struct ToolCall {
310    /// Unique identifier for this tool call
311    pub id: String,
312
313    /// Type of tool call
314    #[serde(rename = "type")]
315    pub call_type: ToolCallType,
316
317    /// Tool function details
318    pub function: ToolFunction,
319}
320
321impl ToolCall {
322    /// Create a new function tool call
323    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/// Type of tool call
340#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
341#[serde(rename_all = "lowercase")]
342pub enum ToolCallType {
343    /// Function call
344    Function,
345}
346
347/// Tool function call details
348#[derive(Debug, Clone, Serialize, Deserialize)]
349pub struct ToolFunction {
350    /// Name of the function to call
351    pub name: String,
352
353    /// Arguments to pass to the function (as JSON)
354    pub arguments: serde_json::Value,
355}
356
357/// Message builder for fluent message construction
358pub struct MessageBuilder {
359    message: ChatMessage,
360}
361
362impl MessageBuilder {
363    /// Start building a new message
364    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    /// Set the content
379    pub fn content(mut self, content: impl Into<MessageContent>) -> Self {
380        self.message.content = content.into();
381        self
382    }
383
384    /// Set the sender name
385    pub fn name(mut self, name: impl Into<String>) -> Self {
386        self.message.name = Some(name.into());
387        self
388    }
389
390    /// Add tool calls
391    pub fn tool_calls(mut self, tool_calls: Vec<ToolCall>) -> Self {
392        self.message.tool_calls = Some(tool_calls);
393        self
394    }
395
396    /// Set tool call ID
397    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    /// Add metadata
403    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    /// Build the message
409    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}