diff options
Diffstat (limited to 'src/config/mod.rs')
| -rw-r--r-- | src/config/mod.rs | 192 |
1 files changed, 82 insertions, 110 deletions
diff --git a/src/config/mod.rs b/src/config/mod.rs index 3fc3f5b..4a0a126 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -6,14 +6,15 @@ use self::conversation::Conversation; use self::message::Message; use self::role::Role; +use crate::client::openai::{OpenAIClient, OpenAIConfig}; +use crate::client::{all_clients, create_client_config, list_models, ClientConfig, ModelInfo}; use crate::config::message::num_tokens_from_messages; -use crate::utils::{mask_text, now}; +use crate::utils::{get_env_name, now}; use anyhow::{anyhow, bail, Context, Result}; -use inquire::{Confirm, Text}; +use inquire::{Confirm, Select}; use parking_lot::RwLock; use serde::Deserialize; -use std::time::Duration; use std::{ env, fs::{create_dir_all, read_to_string, File, OpenOptions}, @@ -23,24 +24,16 @@ use std::{ sync::Arc, }; -pub const MODELS: [(&str, usize); 4] = [ - ("gpt-4", 8192), - ("gpt-4-32k", 32768), - ("gpt-3.5-turbo", 4096), - ("gpt-3.5-turbo-16k", 16384), -]; - const CONFIG_FILE_NAME: &str = "config.yaml"; const ROLES_FILE_NAME: &str = "roles.yaml"; const HISTORY_FILE_NAME: &str = "history.txt"; const MESSAGE_FILE_NAME: &str = "messages.md"; -const SET_COMPLETIONS: [&str; 8] = [ +const SET_COMPLETIONS: [&str; 7] = [ ".set temperature", ".set save true", ".set save false", ".set highlight true", ".set highlight false", - ".set proxy", ".set dry_run true", ".set dry_run false", ]; @@ -49,33 +42,26 @@ const SET_COMPLETIONS: [&str; 8] = [ #[derive(Debug, Clone, Deserialize)] #[serde(default)] pub struct Config { - /// OpenAI api key - pub api_key: Option<String>, - /// OpenAI organization id - pub organization_id: Option<String>, - /// OpenAI model - #[serde(rename(serialize = "model", deserialize = "model"))] - pub model_name: Option<String>, + /// LLM model + pub model: Option<String>, /// What sampling temperature to use, between 0 and 2 pub temperature: Option<f64>, /// Whether to persistently save chat messages pub save: bool, /// Whether to disable highlight pub highlight: bool, - /// Set proxy - pub proxy: Option<String>, /// Used only for debugging pub dry_run: bool, /// If set ture, start a conversation immediately upon repl pub conversation_first: bool, /// Is ligth theme pub light_theme: bool, - /// Set a timeout in seconds for connect to gpt - pub connect_timeout: usize, /// Automatically copy the last output to the clipboard pub auto_copy: bool, /// Use vi keybindings, overriding the default Emacs keybindings pub vi_keybindings: bool, + /// LLM clients + pub clients: Vec<ClientConfig>, /// Predefined roles #[serde(skip)] pub roles: Vec<Role>, @@ -86,29 +72,26 @@ pub struct Config { #[serde(skip)] pub conversation: Option<Conversation>, #[serde(skip)] - pub model: (String, usize), + pub model_info: ModelInfo, } impl Default for Config { fn default() -> Self { Self { - api_key: None, - organization_id: None, - model_name: None, + model: None, temperature: None, save: false, highlight: true, - proxy: None, dry_run: false, conversation_first: false, light_theme: false, - connect_timeout: 10, auto_copy: false, vi_keybindings: false, roles: vec![], + clients: vec![ClientConfig::OpenAI(OpenAIConfig::default())], role: None, conversation: None, - model: ("gpt-3.5-turbo".into(), 4096), + model_info: Default::default(), } } } @@ -118,27 +101,29 @@ pub type SharedConfig = Arc<RwLock<Config>>; impl Config { pub fn init(is_interactive: bool) -> Result<Self> { - let api_key = env::var(get_env_name("api_key")).ok(); let config_path = Self::config_file()?; - if is_interactive && api_key.is_none() && !config_path.exists() { + + let api_key = env::var(get_env_name("api_key")).ok(); + + let exist_config_path = config_path.exists(); + if is_interactive && api_key.is_none() && !exist_config_path { create_config_file(&config_path)?; } - let mut config = if api_key.is_some() && !config_path.exists() { + let mut config = if api_key.is_some() && !exist_config_path { Self::default() } else { Self::load_config(&config_path)? }; - if api_key.is_some() { - config.api_key = api_key; - } - if config.api_key.is_none() { - bail!("api_key not set"); + + // Compatible with old configuration files + if exist_config_path { + config.compat_old_config(&config_path)?; } - if let Some(name) = config.model_name.clone() { + + if let Some(name) = config.model.clone() { config.set_model(&name)?; } config.merge_env_vars(); - config.maybe_proxy(); config.load_roles()?; Ok(config) @@ -211,12 +196,6 @@ impl Config { Self::local_file(CONFIG_FILE_NAME) } - pub fn get_api_key(&self) -> (String, Option<String>) { - let api_key = self.api_key.as_ref().expect("api_key not set"); - let organization_id = self.organization_id.as_ref(); - (api_key.into(), organization_id.cloned()) - } - pub fn roles_file() -> Result<PathBuf> { let env_name = get_env_name("roles_file"); env::var(env_name).map_or_else( @@ -283,14 +262,6 @@ impl Config { } } - pub const fn get_connect_timeout(&self) -> Duration { - Duration::from_secs(self.connect_timeout as u64) - } - - pub fn get_model(&self) -> (String, usize) { - self.model.clone() - } - pub fn build_messages(&self, content: &str) -> Result<Vec<Message>> { #[allow(clippy::option_if_let_else)] let messages = if let Some(conversation) = self.conversation.as_ref() { @@ -302,24 +273,29 @@ impl Config { vec![message] }; let tokens = num_tokens_from_messages(&messages); - if tokens >= self.model.1 { + if tokens >= self.model_info.max_tokens { bail!("Exceed max tokens limit") } Ok(messages) } - pub fn set_model(&mut self, name: &str) -> Result<()> { - if let Some(token) = MODELS.iter().find(|(v, _)| *v == name).map(|(_, v)| *v) { - self.model = (name.to_string(), token); - } else { - bail!("Invalid model") + pub fn set_model(&mut self, value: &str) -> Result<()> { + let models = list_models(self); + if value.contains(':') { + if let Some(model) = models.iter().find(|v| v.stringify() == value) { + self.model_info = model.clone(); + return Ok(()); + } + } else if let Some(model) = models.iter().find(|v| v.client == value) { + self.model_info = model.clone(); + return Ok(()); } - Ok(()) + bail!("Invalid model") } pub const fn get_reamind_tokens(&self) -> usize { - let mut tokens = self.model.1; + let mut tokens = self.model_info.max_tokens; if let Some(conversation) = self.conversation.as_ref() { tokens = tokens.saturating_sub(conversation.tokens); } @@ -331,30 +307,19 @@ impl Config { let state = if path.exists() { "" } else { " ⚠️" }; format!("{}{state}", path.display()) }; - let proxy = self - .proxy - .as_ref() - .map_or_else(|| String::from("-"), std::string::ToString::to_string); let temperature = self .temperature .map_or_else(|| String::from("-"), |v| v.to_string()); - let (api_key, organization_id) = self.get_api_key(); - let api_key = mask_text(&api_key, 3, 4); - let organization_id = organization_id.map_or_else(|| "-".into(), |v| mask_text(&v, 3, 4)); let items = vec![ ("config_file", file_info(&Self::config_file()?)), ("roles_file", file_info(&Self::roles_file()?)), ("messages_file", file_info(&Self::messages_file()?)), - ("api_key", api_key), - ("organization_id", organization_id), - ("model", self.model.0.to_string()), + ("model", self.model_info.stringify()), ("temperature", temperature), ("save", self.save.to_string()), ("highlight", self.highlight.to_string()), - ("proxy", proxy), ("conversation_first", self.conversation_first.to_string()), ("light_theme", self.light_theme.to_string()), - ("connect_timeout", self.connect_timeout.to_string()), ("dry_run", self.dry_run.to_string()), ("vi_keybindings", self.vi_keybindings.to_string()), ]; @@ -373,7 +338,11 @@ impl Config { .collect(); completion.extend(SET_COMPLETIONS.map(std::string::ToString::to_string)); - completion.extend(MODELS.map(|(v, _)| format!(".model {v}"))); + completion.extend( + list_models(self) + .iter() + .map(|v| format!(".model {}", v.stringify())), + ); completion } @@ -402,13 +371,6 @@ impl Config { let value = value.parse().with_context(|| "Invalid value")?; self.highlight = value; } - "proxy" => { - if unset { - self.proxy = None; - } else { - self.proxy = Some(value.to_string()); - } - } "dry_run" => { let value = value.parse().with_context(|| "Invalid value")?; self.dry_run = value; @@ -501,44 +463,62 @@ impl Config { } } - fn maybe_proxy(&mut self) { - if self.proxy.is_some() { - return; + fn compat_old_config(&mut self, config_path: &PathBuf) -> Result<()> { + let content = read_to_string(config_path)?; + let value: serde_json::Value = serde_yaml::from_str(&content)?; + if value.get("client").is_some() { + return Ok(()); } - if let Ok(value) = env::var("HTTPS_PROXY").or_else(|_| env::var("ALL_PROXY")) { - self.proxy = Some(value); + + if let Some(model_name) = value.get("model").and_then(|v| v.as_str()) { + if model_name.starts_with("gpt") { + self.model = Some(format!("{}:{}", OpenAIClient::name(), model_name)); + } } + + if let Some(ClientConfig::OpenAI(client_config)) = self.clients.get_mut(0) { + if let Some(api_key) = value.get("api_key").and_then(|v| v.as_str()) { + client_config.api_key = Some(api_key.to_string()) + } + + if let Some(organization_id) = value.get("organization_id").and_then(|v| v.as_str()) { + client_config.organization_id = Some(organization_id.to_string()) + } + + if let Some(proxy) = value.get("proxy").and_then(|v| v.as_str()) { + client_config.proxy = Some(proxy.to_string()) + } + + if let Some(connect_timeout) = value.get("connect_timeout").and_then(|v| v.as_i64()) { + client_config.connect_timeout = Some(connect_timeout as _) + } + } + Ok(()) } } fn create_config_file(config_path: &Path) -> Result<()> { - let confirm_map_err = |_| anyhow!("Not finish questionnaire, try again later."); - let text_map_err = |_| anyhow!("An error happened when asking for your key, try again later."); let ans = Confirm::new("No config file, create a new one?") .with_default(true) .prompt() - .map_err(confirm_map_err)?; + .map_err(|_| anyhow!("Not finish questionnaire, try again later."))?; if !ans { exit(0); } - let api_key = Text::new("OpenAI API Key:") - .prompt() - .map_err(text_map_err)?; - let mut raw_config = format!("api_key: {api_key}\n"); - let ans = Confirm::new("Use proxy?") - .with_default(false) + let client = Select::new("Choose bots?", all_clients()) .prompt() - .map_err(confirm_map_err)?; - if ans { - let proxy = Text::new("Set proxy:").prompt().map_err(text_map_err)?; - raw_config.push_str(&format!("proxy: {proxy}\n")); - } + .map_err(|_| anyhow!("An error happened when selecting bots, try again later."))?; + + let mut raw_config = create_client_config(client)?; + + raw_config.push_str(&format!("model: {client}\n")); let ans = Confirm::new("Save chat messages") .with_default(true) .prompt() - .map_err(confirm_map_err)?; + .map_err(|_| anyhow!("Not finish questionnaire, try again later."))?; + if ans { raw_config.push_str("save: true\n"); } @@ -571,14 +551,6 @@ fn ensure_parent_exists(path: &Path) -> Result<()> { Ok(()) } -fn get_env_name(key: &str) -> String { - format!( - "{}_{}", - env!("CARGO_CRATE_NAME").to_ascii_uppercase(), - key.to_ascii_uppercase(), - ) -} - fn set_bool(target: &mut bool, value: &str) { match value { "1" | "true" => *target = true, |