diff options
| author | sigoden <sigoden@gmail.com> | 2023-11-01 22:15:55 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-11-01 22:15:55 +0800 |
| commit | f6da06dad9b2a76016209a7d58ad923c2c72f150 (patch) | |
| tree | 1c74f306c194d3289fe06ed31a39a00c27ad0ca9 /src/config/message.rs | |
| parent | da3c541b681b140feb52b4297167ca233535fa90 (diff) | |
| download | aichat-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.rs | 19 |
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)] |
