cortexai_llm_client/
message.rs1use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
7#[serde(rename_all = "lowercase")]
8pub enum Role {
9 System,
11 User,
13 Assistant,
15 Tool,
17}
18
19impl Role {
20 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#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct Message {
34 pub role: Role,
36 pub content: String,
38 #[serde(skip_serializing_if = "Option::is_none")]
40 pub tool_call_id: Option<String>,
41}
42
43impl Message {
44 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 pub fn system(content: impl Into<String>) -> Self {
55 Self::new(Role::System, content)
56 }
57
58 pub fn user(content: impl Into<String>) -> Self {
60 Self::new(Role::User, content)
61 }
62
63 pub fn assistant(content: impl Into<String>) -> Self {
65 Self::new(Role::Assistant, content)
66 }
67
68 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}