Skip to main content

orchestra_rs/messages/
mod.rs

1use serde::{Deserialize, Serialize};
2
3/// Represents different types of messages in a conversation
4#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
5pub enum Message {
6    /// Message from a human user
7    Human(HumanMessage),
8    /// Message from an AI assistant
9    Assistant(AssistantMessage),
10    /// System instruction or context message
11    System(SystemMessage),
12}
13
14/// Message from a human user
15#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
16pub struct HumanMessage {
17    pub content: MessageContent,
18}
19
20/// Message from an AI assistant
21#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
22pub struct AssistantMessage {
23    pub content: MessageContent,
24}
25
26/// System instruction or context message
27#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
28pub struct SystemMessage {
29    pub content: String,
30}
31
32/// Content of a message, which can be text or include tool calls
33#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
34pub enum MessageContent {
35    /// Simple text content
36    Text(String),
37    /// Mixed content with text and tool calls
38    Mixed {
39        text: Option<String>,
40        tool_calls: Vec<ToolCall>,
41    },
42}
43
44impl MessageContent {
45    /// Create a new text content
46    pub fn text<S: Into<String>>(text: S) -> Self {
47        Self::Text(text.into())
48    }
49
50    /// Create mixed content with text and tool calls
51    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    /// Get the text content, if any
59    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    /// Get the text content as a string, combining all text parts
67    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    /// Check if this content has tool calls
75    pub fn has_tool_calls(&self) -> bool {
76        matches!(self, Self::Mixed { tool_calls, .. } if !tool_calls.is_empty())
77    }
78
79    /// Get tool calls, if any
80    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    /// Create a new human message with text content
102    pub fn human<S: Into<String>>(content: S) -> Self {
103        Self::Human(HumanMessage {
104            content: MessageContent::text(content),
105        })
106    }
107
108    /// Create a new assistant message with text content
109    pub fn assistant<S: Into<String>>(content: S) -> Self {
110        Self::Assistant(AssistantMessage {
111            content: MessageContent::text(content),
112        })
113    }
114
115    /// Create a new system message
116    pub fn system<S: Into<String>>(content: S) -> Self {
117        Self::System(SystemMessage {
118            content: content.into(),
119        })
120    }
121
122    /// Get the role of this message as a string
123    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    /// Get the text content of this message
132    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    /// Create a new human message with text content
143    pub fn new<S: Into<String>>(content: S) -> Self {
144        Self {
145            content: MessageContent::text(content),
146        }
147    }
148
149    /// Create a new human message with mixed content
150    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    /// Create a new assistant message with text content
159    pub fn new<S: Into<String>>(content: S) -> Self {
160        Self {
161            content: MessageContent::text(content),
162        }
163    }
164
165    /// Create a new assistant message with mixed content
166    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    /// Create a new system message
175    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/// Basic text content.
295#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
296pub struct Text {
297    pub text: String,
298}
299
300/// Describes a tool call with an id and function to call.
301#[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/// Describes a tool function to call with a name and arguments.
309#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
310pub struct ToolFunction {
311    pub name: String,
312    pub arguments: serde_json::Value,
313}