chatty_rs/models/
message.rs

1#[derive(Debug, Clone)]
2pub enum Issuer {
3    System(String),
4    User(String),
5}
6
7#[derive(Debug, Clone)]
8pub struct Message {
9    id: String,
10    issuer: Issuer,
11    text: String,
12    token_count: usize,
13    created_at: chrono::DateTime<chrono::Utc>,
14    /// Indicates if the message is part of a context
15    context: bool,
16}
17
18impl Message {
19    pub fn new(issuer: Issuer, text: impl Into<String>) -> Self {
20        Self {
21            id: uuid::Uuid::new_v4().to_string(),
22            issuer,
23            text: text.into(),
24            token_count: 0,
25            created_at: chrono::Utc::now(),
26            context: false,
27        }
28    }
29
30    pub fn new_system(system: &str, text: impl Into<String>) -> Self {
31        Self::new(Issuer::System(system.to_string()), text)
32    }
33
34    pub fn new_user(user: &str, text: impl Into<String>) -> Self {
35        Self::new(Issuer::User(user.to_string()), text)
36    }
37
38    pub fn with_context(mut self, context: bool) -> Self {
39        if matches!(self.issuer, Issuer::System(_)) {
40            self.context = context;
41        }
42        self
43    }
44
45    /// Indicates if the message is part of a context. Only avaiable for
46    /// system messages. If backend receives a message is marked as context
47    ///, it should be treated as a context message. E.g: OpenAI's role should
48    /// be "system" instead of "assistant".
49    pub fn is_context(&self) -> bool {
50        self.context
51    }
52
53    pub fn with_id(mut self, id: impl Into<String>) -> Self {
54        self.id = id.into();
55        self
56    }
57
58    pub fn with_created_at(mut self, timestamp: chrono::DateTime<chrono::Utc>) -> Self {
59        self.created_at = timestamp;
60        self
61    }
62
63    pub fn with_text(mut self, text: impl Into<String>) -> Self {
64        self.text = text.into();
65        self
66    }
67
68    pub fn with_token_count(mut self, token_count: usize) -> Self {
69        self.set_token_count(token_count);
70        self
71    }
72
73    pub fn set_token_count(&mut self, token_count: usize) {
74        self.token_count = token_count;
75    }
76
77    pub fn token_count(&self) -> usize {
78        self.token_count
79    }
80
81    pub fn id(&self) -> &str {
82        &self.id
83    }
84
85    pub fn is_system(&self) -> bool {
86        matches!(self.issuer, Issuer::System(_))
87    }
88
89    pub fn issuer(&self) -> &Issuer {
90        &self.issuer
91    }
92
93    pub fn text(&self) -> &str {
94        &self.text
95    }
96
97    pub fn created_at(&self) -> chrono::DateTime<chrono::Utc> {
98        self.created_at
99    }
100
101    pub fn issuer_str(&self) -> &str {
102        match &self.issuer {
103            Issuer::System(s) => s,
104            Issuer::User(u) => u,
105        }
106    }
107
108    pub fn append(&mut self, text: impl Into<String>) {
109        let text = text.into();
110        self.text += &text.replace('\t', "  ");
111    }
112}
113
114impl Issuer {
115    pub fn user() -> Self {
116        Self::User("".to_string())
117    }
118
119    pub fn user_with_name(name: impl Into<String>) -> Self {
120        Self::User(name.into())
121    }
122
123    pub fn system() -> Self {
124        Self::System("".to_string())
125    }
126
127    pub fn system_with_name(name: impl Into<String>) -> Self {
128        Self::System(name.into())
129    }
130}
chatty_rs/models/message.rs

chatty_rs/models/
message.rs
