Skip to main content

simple_agent_type/
message.rs

1//! Message types for LLM interactions.
2//!
3//! Provides role-based messages compatible with OpenAI's message format.
4
5use serde::{Deserialize, Serialize};
6
7/// Role of a message in a conversation.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
9#[serde(rename_all = "lowercase")]
10pub enum Role {
11    /// User message
12    User,
13    /// Assistant (LLM) message
14    Assistant,
15    /// System instruction message
16    System,
17    /// Tool/function call result
18    #[serde(rename = "tool")]
19    Tool,
20}
21
22/// A message in a conversation.
23#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
24pub struct Message {
25    /// Role of the message sender
26    pub role: Role,
27    /// Content of the message
28    pub content: String,
29    /// Optional name (for multi-user conversations or tool calls)
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub name: Option<String>,
32    /// Tool call ID (for tool role messages)
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub tool_call_id: Option<String>,
35}
36
37impl Message {
38    /// Create a user message.
39    ///
40    /// # Example
41    /// ```
42    /// use simple_agent_type::message::{Message, Role};
43    ///
44    /// let msg = Message::user("Hello!");
45    /// assert_eq!(msg.role, Role::User);
46    /// assert_eq!(msg.content, "Hello!");
47    /// ```
48    pub fn user(content: impl Into<String>) -> Self {
49        Self {
50            role: Role::User,
51            content: content.into(),
52            name: None,
53            tool_call_id: None,
54        }
55    }
56
57    /// Create an assistant message.
58    ///
59    /// # Example
60    /// ```
61    /// use simple_agent_type::message::{Message, Role};
62    ///
63    /// let msg = Message::assistant("Hi there!");
64    /// assert_eq!(msg.role, Role::Assistant);
65    /// ```
66    pub fn assistant(content: impl Into<String>) -> Self {
67        Self {
68            role: Role::Assistant,
69            content: content.into(),
70            name: None,
71            tool_call_id: None,
72        }
73    }
74
75    /// Create a system message.
76    ///
77    /// # Example
78    /// ```
79    /// use simple_agent_type::message::{Message, Role};
80    ///
81    /// let msg = Message::system("You are a helpful assistant.");
82    /// assert_eq!(msg.role, Role::System);
83    /// ```
84    pub fn system(content: impl Into<String>) -> Self {
85        Self {
86            role: Role::System,
87            content: content.into(),
88            name: None,
89            tool_call_id: None,
90        }
91    }
92
93    /// Create a tool message.
94    ///
95    /// # Example
96    /// ```
97    /// use simple_agent_type::message::{Message, Role};
98    ///
99    /// let msg = Message::tool("result", "call_123");
100    /// assert_eq!(msg.role, Role::Tool);
101    /// assert_eq!(msg.tool_call_id, Some("call_123".to_string()));
102    /// ```
103    pub fn tool(content: impl Into<String>, tool_call_id: impl Into<String>) -> Self {
104        Self {
105            role: Role::Tool,
106            content: content.into(),
107            name: None,
108            tool_call_id: Some(tool_call_id.into()),
109        }
110    }
111
112    /// Set the name field (builder pattern).
113    ///
114    /// # Example
115    /// ```
116    /// use simple_agent_type::message::Message;
117    ///
118    /// let msg = Message::user("Hello").with_name("Alice");
119    /// assert_eq!(msg.name, Some("Alice".to_string()));
120    /// ```
121    pub fn with_name(mut self, name: impl Into<String>) -> Self {
122        self.name = Some(name.into());
123        self
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn test_message_user() {
133        let msg = Message::user("test");
134        assert_eq!(msg.role, Role::User);
135        assert_eq!(msg.content, "test");
136        assert_eq!(msg.name, None);
137        assert_eq!(msg.tool_call_id, None);
138    }
139
140    #[test]
141    fn test_message_assistant() {
142        let msg = Message::assistant("response");
143        assert_eq!(msg.role, Role::Assistant);
144        assert_eq!(msg.content, "response");
145    }
146
147    #[test]
148    fn test_message_system() {
149        let msg = Message::system("instruction");
150        assert_eq!(msg.role, Role::System);
151        assert_eq!(msg.content, "instruction");
152    }
153
154    #[test]
155    fn test_message_tool() {
156        let msg = Message::tool("result", "call_123");
157        assert_eq!(msg.role, Role::Tool);
158        assert_eq!(msg.content, "result");
159        assert_eq!(msg.tool_call_id, Some("call_123".to_string()));
160    }
161
162    #[test]
163    fn test_message_with_name() {
164        let msg = Message::user("test").with_name("Alice");
165        assert_eq!(msg.name, Some("Alice".to_string()));
166    }
167
168    #[test]
169    fn test_role_serialization() {
170        let json = serde_json::to_string(&Role::User).unwrap();
171        assert_eq!(json, "\"user\"");
172
173        let json = serde_json::to_string(&Role::Assistant).unwrap();
174        assert_eq!(json, "\"assistant\"");
175
176        let json = serde_json::to_string(&Role::System).unwrap();
177        assert_eq!(json, "\"system\"");
178
179        let json = serde_json::to_string(&Role::Tool).unwrap();
180        assert_eq!(json, "\"tool\"");
181    }
182
183    #[test]
184    fn test_message_serialization() {
185        let msg = Message::user("Hello");
186        let json = serde_json::to_string(&msg).unwrap();
187        let parsed: Message = serde_json::from_str(&json).unwrap();
188        assert_eq!(msg, parsed);
189    }
190
191    #[test]
192    fn test_message_optional_fields_not_serialized() {
193        let msg = Message::user("test");
194        let json = serde_json::to_value(&msg).unwrap();
195        assert!(json.get("name").is_none());
196        assert!(json.get("tool_call_id").is_none());
197    }
198
199    #[test]
200    fn test_message_with_name_serialized() {
201        let msg = Message::user("test").with_name("Alice");
202        let json = serde_json::to_value(&msg).unwrap();
203        assert_eq!(json.get("name").and_then(|v| v.as_str()), Some("Alice"));
204    }
205}