Skip to main content

cortexai_llm_client/
message.rs

1//! Message types for conversations.
2
3use serde::{Deserialize, Serialize};
4
5/// Role in a conversation.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
7#[serde(rename_all = "lowercase")]
8pub enum Role {
9    /// System message (instructions)
10    System,
11    /// User message
12    User,
13    /// Assistant message
14    Assistant,
15    /// Tool/function result
16    Tool,
17}
18
19impl Role {
20    /// Convert to string for API requests.
21    pub fn as_str(&self) -> &'static str {
22        match self {
23            Self::System => "system",
24            Self::User => "user",
25            Self::Assistant => "assistant",
26            Self::Tool => "tool",
27        }
28    }
29}
30
31/// A message in a conversation.
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct Message {
34    /// Role of the message sender
35    pub role: Role,
36    /// Content of the message
37    pub content: String,
38    /// Tool call ID (for tool responses)
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub tool_call_id: Option<String>,
41}
42
43impl Message {
44    /// Create a new message.
45    pub fn new(role: Role, content: impl Into<String>) -> Self {
46        Self {
47            role,
48            content: content.into(),
49            tool_call_id: None,
50        }
51    }
52
53    /// Create a system message.
54    pub fn system(content: impl Into<String>) -> Self {
55        Self::new(Role::System, content)
56    }
57
58    /// Create a user message.
59    pub fn user(content: impl Into<String>) -> Self {
60        Self::new(Role::User, content)
61    }
62
63    /// Create an assistant message.
64    pub fn assistant(content: impl Into<String>) -> Self {
65        Self::new(Role::Assistant, content)
66    }
67
68    /// Create a tool result message.
69    pub fn tool(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
70        Self {
71            role: Role::Tool,
72            content: content.into(),
73            tool_call_id: Some(tool_call_id.into()),
74        }
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    #[test]
83    fn test_message_creation() {
84        let msg = Message::user("Hello");
85        assert_eq!(msg.role, Role::User);
86        assert_eq!(msg.content, "Hello");
87    }
88
89    #[test]
90    fn test_message_serialization() {
91        let msg = Message::system("You are helpful");
92        let json = serde_json::to_string(&msg).unwrap();
93        assert!(json.contains("system"));
94        assert!(json.contains("You are helpful"));
95    }
96
97    #[test]
98    fn test_tool_message() {
99        let msg = Message::tool("call_123", r#"{"result": "success"}"#);
100        assert_eq!(msg.role, Role::Tool);
101        assert_eq!(msg.tool_call_id, Some("call_123".to_string()));
102    }
103}