Archives/aichat
2
0
Fork 0
mirror of https://github.com/sigoden/aichat synced 2024-11-18 09:28:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: opitmize code, move MAX_TOKENS related to config/message.rs

Browse Source
This commit is contained in:
sigoden 2023-03-10 07:25:33 +08:00
parent 899c4af9ea
commit 978d4c54c0
3 changed files with 15 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/config/conversation.rs
View File

@ -1,6 +1,5 @@
use super::message::{num_tokens_from_messages, Message, MessageRole}; use super::message::{num_tokens_from_messages, Message, MessageRole, MAX_TOKENS};
use super::role::Role; use super::role::Role;
use super::MAX_TOKENS;
use anyhow::{bail, Result}; use anyhow::{bail, Result};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};

13
src/config/message.rs
View File

@ -1,6 +1,9 @@
use crate::utils::count_tokens;
use anyhow::{bail, Result};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::utils::count_tokens; pub const MAX_TOKENS: usize = 4096;
#[derive(Debug, Clone, Deserialize, Serialize)] #[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Message { pub struct Message {
@ -25,6 +28,14 @@ pub enum MessageRole {
User, User,
} }
pub fn within_max_tokens_limit(messages: &[Message]) -> Result<()> {
let tokens = num_tokens_from_messages(messages);
if tokens >= MAX_TOKENS {
bail!("Exceed max tokens limit")
}
Ok(())
}
pub fn num_tokens_from_messages(messages: &[Message]) -> usize { pub fn num_tokens_from_messages(messages: &[Message]) -> usize {
let mut num_tokens = 0; let mut num_tokens = 0;
for message in messages.iter() { for message in messages.iter() {

8
src/config/mod.rs
View File

@ -4,7 +4,7 @@ mod role;
use self::message::Message; use self::message::Message;
use self::role::Role; use self::role::Role;
use self::{conversation::Conversation, message::num_tokens_from_messages}; use self::{conversation::Conversation, message::within_max_tokens_limit};
use crate::utils::now; use crate::utils::now;
@ -21,7 +21,6 @@ use std::{
sync::Arc, sync::Arc,
}; };
const MAX_TOKENS: usize = 4096;
const CONFIG_FILE_NAME: &str = "config.yaml"; const CONFIG_FILE_NAME: &str = "config.yaml";
const ROLES_FILE_NAME: &str = "roles.yaml"; const ROLES_FILE_NAME: &str = "roles.yaml";
const HISTORY_FILE_NAME: &str = "history.txt"; const HISTORY_FILE_NAME: &str = "history.txt";
@ -229,10 +228,7 @@ impl Config {
let message = Message::new(content); let message = Message::new(content);
vec![message] vec![message]
}; };
let tokens = num_tokens_from_messages(&messages); within_max_tokens_limit(&messages)?;
if tokens >= MAX_TOKENS {
bail!("Exceed max tokens limit")
}
Ok(messages) Ok(messages)
} }