summaryrefslogtreecommitdiffstats
path: root/src/config/message.rs
diff options
context:
space:
mode:
authorsigoden <sigoden@gmail.com>2023-11-01 22:15:55 +0800
committerGitHub <noreply@github.com>2023-11-01 22:15:55 +0800
commitf6da06dad9b2a76016209a7d58ad923c2c72f150 (patch)
tree1c74f306c194d3289fe06ed31a39a00c27ad0ca9 /src/config/message.rs
parentda3c541b681b140feb52b4297167ca233535fa90 (diff)
downloadaichat-f6da06dad9b2a76016209a7d58ad923c2c72f150.tar.gz
refactor: improve code quanity (#194)
- extends ModelInfo for tokens calculating - refactor config/session.rs, improve export, render, getter/setter - modify main.rs, allow --model override session.model
Diffstat (limited to 'src/config/message.rs')
-rw-r--r--src/config/message.rs19
1 files changed, 7 insertions, 12 deletions
diff --git a/src/config/message.rs b/src/config/message.rs
index 5882337..55b2663 100644
--- a/src/config/message.rs
+++ b/src/config/message.rs
@@ -1,5 +1,3 @@
-use crate::utils::count_tokens;
-
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Deserialize, Serialize)]
@@ -25,22 +23,19 @@ pub enum MessageRole {
User,
}
+#[allow(dead_code)]
impl MessageRole {
- #[allow(dead_code)]
pub fn is_system(&self) -> bool {
matches!(self, MessageRole::System)
}
-}
-pub fn num_tokens_from_messages(messages: &[Message]) -> usize {
- let mut num_tokens = 0;
- for message in messages.iter() {
- num_tokens += 4;
- num_tokens += count_tokens(&message.content);
- num_tokens += 1; // role always take 1 token
+ pub fn is_user(&self) -> bool {
+ matches!(self, MessageRole::User)
+ }
+
+ pub fn is_assistant(&self) -> bool {
+ matches!(self, MessageRole::Assistant)
}
- num_tokens += 2;
- num_tokens
}
#[cfg(test)]