Skip to main content

cake_core/models/
chat.rs

1use serde::{Deserialize, Serialize};
2
3/// The role of a message in a chat.
4#[derive(Debug, Serialize, Deserialize)]
5#[serde(rename_all = "lowercase")]
6pub enum MessageRole {
7    /// System prompt.
8    System,
9    /// User prompt.
10    User,
11    /// Assistant response.
12    Assistant,
13}
14
15impl std::fmt::Display for MessageRole {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        write!(
18            f,
19            "{}",
20            match self {
21                MessageRole::System => "system",
22                MessageRole::User => "user",
23                MessageRole::Assistant => "assistant",
24            }
25        )
26    }
27}
28
29/// A chat message.
30#[derive(Debug, Serialize, Deserialize)]
31pub struct Message {
32    /// Message role.
33    pub role: MessageRole,
34    /// Messagae content.
35    pub content: String,
36}
37
38impl Message {
39    /// Create a system message.
40    pub fn system(content: String) -> Self {
41        Self {
42            role: MessageRole::System,
43            content,
44        }
45    }
46
47    /// Create a user message.
48    pub fn user(content: String) -> Self {
49        Self {
50            role: MessageRole::User,
51            content,
52        }
53    }
54
55    /// Create an assistant message.
56    pub fn assistant(content: String) -> Self {
57        Self {
58            role: MessageRole::Assistant,
59            content,
60        }
61    }
62}