From f6da06dad9b2a76016209a7d58ad923c2c72f150 Mon Sep 17 00:00:00 2001 From: sigoden Date: Wed, 1 Nov 2023 22:15:55 +0800 Subject: 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 --- src/config/mod.rs | 64 ++++++++++++++++++++++++++----------------------------- 1 file changed, 30 insertions(+), 34 deletions(-) (limited to 'src/config/mod.rs') diff --git a/src/config/mod.rs b/src/config/mod.rs index 18f1731..eff9bcf 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -12,7 +12,6 @@ use crate::client::{ all_models, create_client_config, list_client_types, ClientConfig, ExtraConfig, OpenAIClient, SendData, }; -use crate::config::message::num_tokens_from_messages; use crate::render::RenderOptions; use crate::utils::{get_env_name, light_theme_from_colorfgbg, now, prompt_op_err}; @@ -274,7 +273,7 @@ impl Config { pub fn set_temperature(&mut self, value: Option) -> Result<()> { self.temperature = value; if let Some(session) = self.session.as_mut() { - session.temperature = value; + session.set_temperature(value); } Ok(()) } @@ -298,13 +297,6 @@ impl Config { let message = Message::new(content); vec![message] }; - let tokens = num_tokens_from_messages(&messages); - if let Some(max_tokens) = self.model_info.max_tokens { - if tokens >= max_tokens { - bail!("Exceed max tokens limit") - } - } - Ok(messages) } @@ -326,7 +318,7 @@ impl Config { let models = all_models(self); let mut model_info = None; if value.contains(':') { - if let Some(model) = models.iter().find(|v| v.stringify() == value) { + if let Some(model) = models.iter().find(|v| v.full_name() == value) { model_info = Some(model.clone()); } } else if let Some(model) = models.iter().find(|v| v.client == value) { @@ -336,7 +328,7 @@ impl Config { None => bail!("Unknown model '{}'", value), Some(model_info) => { if let Some(session) = self.session.as_mut() { - session.set_model(&model_info.stringify())?; + session.set_model(model_info.clone())?; } self.model_info = model_info; Ok(()) @@ -361,7 +353,7 @@ impl Config { ("roles_file", path_info(&Self::roles_file()?)), ("messages_file", path_info(&Self::messages_file()?)), ("sessions_dir", path_info(&Self::sessions_dir()?)), - ("model", self.model_info.stringify()), + ("model", self.model_info.full_name()), ("temperature", temperature), ("save", self.save.to_string()), ("highlight", self.highlight.to_string()), @@ -389,7 +381,7 @@ impl Config { completion.extend( all_models(self) .iter() - .map(|v| format!(".model {}", v.stringify())), + .map(|v| format!(".model {}", v.full_name())), ); let sessions = self.list_sessions().unwrap_or_default(); completion.extend(sessions.iter().map(|v| format!(".session {}", v))); @@ -444,7 +436,7 @@ impl Config { } self.session = Some(Session::new( TEMP_SESSION_NAME, - &self.model_info.stringify(), + self.model_info.clone(), self.role.clone(), )); } @@ -453,13 +445,13 @@ impl Config { if !session_path.exists() { self.session = Some(Session::new( name, - &self.model_info.stringify(), + self.model_info.clone(), self.role.clone(), )); } else { let session = Session::load(name, &session_path)?; - let model = session.model.clone(); - self.temperature = session.temperature; + let model = session.model().to_string(); + self.temperature = session.temperature(); self.session = Some(session); self.set_model(&model)?; } @@ -472,7 +464,8 @@ impl Config { "Start a session that incorporates the last question and answer?", ) .with_default(false) - .prompt()?; + .prompt() + .map_err(prompt_op_err)?; if ans { session.add_message(input, output)?; } @@ -487,13 +480,19 @@ impl Config { self.last_message = None; self.temperature = self.default_temperature; if session.should_save() { - let ans = Confirm::new("Save session?").with_default(true).prompt()?; + let ans = Confirm::new("Save session?") + .with_default(false) + .prompt() + .map_err(prompt_op_err)?; if !ans { return Ok(()); } - let mut name = session.name.clone(); + let mut name = session.name().to_string(); if session.is_temp() { - name = Text::new("Session name:").with_default(&name).prompt()?; + name = Text::new("Session name:") + .with_default(&name) + .prompt() + .map_err(prompt_op_err)?; } let session_path = Self::session_file(&name)?; let sessions_dir = session_path.parent().ok_or_else(|| { @@ -558,17 +557,13 @@ impl Config { pub fn render_prompt_right(&self) -> String { if let Some(session) = &self.session { - let tokens = session.tokens; - // 10000(%32) - match self.model_info.max_tokens { - Some(max_tokens) => { - let ratio = tokens as f32 / max_tokens as f32; - let percent = ratio * 100.0; - let percent = (percent * 100.0).round() / 100.0; - format!("{tokens}({percent}%)") - } - None => format!("{tokens}"), - } + let (tokens, percent) = session.tokens_and_percent(); + let percent = if percent == 0.0 { + String::new() + } else { + format!("({percent}%)") + }; + format!("{tokens}{percent}") } else { String::new() } @@ -576,6 +571,7 @@ impl Config { pub fn prepare_send_data(&self, content: &str, stream: bool) -> Result { let messages = self.build_messages(content)?; + self.model_info.max_tokens_limit(&messages)?; Ok(SendData { messages, temperature: self.get_temperature(), @@ -586,7 +582,7 @@ impl Config { pub fn maybe_print_send_tokens(&self, input: &str) { if self.dry_run { if let Ok(messages) = self.build_messages(input) { - let tokens = num_tokens_from_messages(&messages); + let tokens = self.model_info.totatl_tokens(&messages); println!(">>> This message consumes {tokens} tokens. <<<"); } } @@ -642,7 +638,7 @@ impl Config { bail!("No available model"); } - models[0].stringify() + models[0].full_name() } }; self.set_model(&model)?; -- cgit v1.2.3