diff options
Diffstat (limited to 'src/config')
| -rw-r--r-- | src/config/conversation.rs | 83 | ||||
| -rw-r--r-- | src/config/mod.rs | 429 |
2 files changed, 512 insertions, 0 deletions
diff --git a/src/config/conversation.rs b/src/config/conversation.rs new file mode 100644 index 0000000..ca50233 --- /dev/null +++ b/src/config/conversation.rs @@ -0,0 +1,83 @@ +use anyhow::Result; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct Session { + pub tokens: usize, + pub messages: Vec<Message>, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct Message { + pub role: MessageRole, + pub content: String, +} + +impl Session { + pub fn new() -> Self { + Self { + tokens: 0, + messages: vec![], + } + } + + pub fn add_conversatoin(&mut self, input: &str, output: &str) -> Result<()> { + self.messages.push(Message { + role: MessageRole::User, + content: input.to_string(), + }); + self.messages.push(Message { + role: MessageRole::Assistant, + content: output.to_string(), + }); + Ok(()) + } + + /// Readline prompt + pub fn add_prompt(&mut self, prompt: &str) { + self.messages.push(Message { + role: MessageRole::System, + content: prompt.into(), + }); + } + + pub fn echo_messages(&self, content: &str) -> String { + let mut messages = self.messages.to_vec(); + messages.push(Message { + role: MessageRole::User, + content: content.into(), + }); + serde_yaml::to_string(&messages).unwrap_or("Unable to echo message".into()) + } + + pub fn build_emssages(&self, content: &str) -> Value { + let mut messages: Vec<Value> = self.messages.iter().map(msg_to_value).collect(); + messages.push(msg_to_value(&Message { + role: MessageRole::User, + content: content.into(), + })); + json!(messages) + } +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub enum MessageRole { + System, + Assistant, + User, +} + +impl MessageRole { + pub fn name(&self) -> &'static str { + match self { + MessageRole::System => "system", + MessageRole::Assistant => "assistant", + MessageRole::User => "user", + } + } +} + +fn msg_to_value(msg: &Message) -> Value { + json!({ "role": msg.role.name(), "content": msg.content }) +} diff --git a/src/config/mod.rs b/src/config/mod.rs new file mode 100644 index 0000000..6555b14 --- /dev/null +++ b/src/config/mod.rs @@ -0,0 +1,429 @@ +mod conversation; + +use self::conversation::Session; + +use crate::utils::{emphasis, now}; + +use anyhow::{anyhow, bail, Context, Result}; +use inquire::{Confirm, Text}; +use parking_lot::Mutex; +use serde::Deserialize; +use serde_json::{json, Value}; +use std::{ + env, + fs::{create_dir_all, read_to_string, File, OpenOptions}, + io::Write, + path::{Path, PathBuf}, + process::exit, + sync::Arc, +}; + +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 TEMP_ROLE_NAME: &str = "%PROMPT%"; +const SET_COMPLETIONS: [&str; 9] = [ + ".set api_key", + ".set temperature", + ".set save true", + ".set save false", + ".set highlight true", + ".set highlight false", + ".set proxy", + ".set dry_run true", + ".set dry_run false", +]; + +#[derive(Debug, Clone, Deserialize)] +pub struct Config { + /// Openai api key + pub api_key: String, + /// What sampling temperature to use, between 0 and 2 + pub temperature: Option<f64>, + /// Whether to persistently save chat messages + #[serde(default)] + pub save: bool, + /// Whether to disable highlight + #[serde(default = "highlight_value")] + pub highlight: bool, + /// Set proxy + pub proxy: Option<String>, + /// Used only for debugging + #[serde(default)] + pub dry_run: bool, + /// Predefined roles + #[serde(default, skip)] + pub roles: Vec<Role>, + /// Current selected role + #[serde(default, skip)] + pub role: Option<Role>, + /// Current conversation + #[serde(default, skip)] + pub conversation: Option<Session>, +} + +pub type SharedConfig = Arc<Mutex<Config>>; + +impl Config { + pub fn init(is_interactive: bool) -> Result<Config> { + let config_path = Config::config_file()?; + if is_interactive && !config_path.exists() { + create_config_file(&config_path)?; + } + let content = read_to_string(&config_path) + .with_context(|| format!("Failed to load config at {}", config_path.display()))?; + let mut config: Config = serde_yaml::from_str(&content) + .with_context(|| format!("Invalid config at {}", config_path.display()))?; + config.load_roles()?; + Ok(config) + } + + pub fn find_role(&self, name: &str) -> Option<Role> { + self.roles.iter().find(|v| v.name == name).cloned() + } + + pub fn config_dir() -> Result<PathBuf> { + let env_name = format!( + "{}_CONFIG_DIR", + env!("CARGO_CRATE_NAME").to_ascii_uppercase() + ); + let path = match env::var(env_name) { + Ok(v) => PathBuf::from(v), + Err(_) => { + let mut dir = dirs::config_dir().ok_or_else(|| anyhow!("Not found config dir"))?; + dir.push(env!("CARGO_CRATE_NAME")); + dir + } + }; + if !path.exists() { + create_dir_all(&path).map_err(|err| { + anyhow!("Failed to create config dir at {}, {err}", path.display()) + })?; + } + Ok(path) + } + + pub fn local_file(name: &str) -> Result<PathBuf> { + let mut path = Self::config_dir()?; + path.push(name); + Ok(path) + } + + pub fn save_message(&self, input: &str, output: &str) -> Result<()> { + if !self.save { + return Ok(()); + } + let mut file = self.open_message_file()?; + if output.is_empty() || !self.save { + return Ok(()); + } + let timestamp = now(); + let output = match self.role.as_ref() { + None => { + format!("# CHAT:[{timestamp}]\n{input}\n--------\n{output}\n--------\n\n",) + } + Some(v) => { + if v.name == TEMP_ROLE_NAME { + format!( + "# CHAT:[{timestamp}]\n{}\n{input}\n--------\n{output}\n--------\n\n", + v.prompt + ) + } else { + format!( + "# CHAT:[{timestamp}] ({})\n{input}\n--------\n{output}\n--------\n\n", + v.name, + ) + } + } + }; + file.write_all(output.as_bytes()) + .with_context(|| "Failed to save message") + } + + pub fn config_file() -> Result<PathBuf> { + Self::local_file(CONFIG_FILE_NAME) + } + + pub fn roles_file() -> Result<PathBuf> { + Self::local_file(ROLES_FILE_NAME) + } + + pub fn history_file() -> Result<PathBuf> { + Self::local_file(HISTORY_FILE_NAME) + } + + pub fn messages_file() -> Result<PathBuf> { + Self::local_file(MESSAGE_FILE_NAME) + } + + pub fn change_role(&mut self, name: &str) -> Result<String> { + self.ensure_no_conversation()?; + if self.conversation.is_some() { + bail!("") + } + match self.find_role(name) { + Some(role) => { + let temperature = match role.temperature { + Some(v) => format!("{v}"), + None => "null".into(), + }; + let output = format!( + "{}: {}\n{}: {}\n{}: {}", + emphasis("name"), + role.name, + emphasis("prompt"), + role.prompt.trim(), + emphasis("temperature"), + temperature + ); + self.role = Some(role); + Ok(output) + } + None => bail!("Error: Unknown role"), + } + } + + pub fn create_temp_role(&mut self, prompt: &str) -> Result<()> { + self.ensure_no_conversation()?; + self.role = Some(Role { + name: TEMP_ROLE_NAME.into(), + prompt: prompt.into(), + temperature: self.temperature, + }); + Ok(()) + } + + pub fn get_temperature(&self) -> Option<f64> { + self.role + .as_ref() + .and_then(|v| v.temperature) + .or(self.temperature) + } + + pub fn echo_messages(&self, content: &str) -> String { + if let Some(conversation) = self.conversation.as_ref() { + conversation.echo_messages(content) + } else if let Some(role) = self.role.as_ref() { + format!("{}\n{content}", role.prompt.trim()) + } else { + content.to_string() + } + } + + pub fn build_messages(&self, content: &str) -> Value { + let user_message = json!({ "role": "user", "content": content }); + if let Some(conversation) = self.conversation.as_ref() { + conversation.build_emssages(content) + } else if let Some(role) = self.role.as_ref() { + let system_message = json!({ "role": "system", "content": role.prompt.trim() }); + json!([system_message, user_message]) + } else { + json!([user_message]) + } + } + + pub fn info(&self) -> Result<String> { + let file_info = |path: &Path| { + let state = if path.exists() { "" } else { " ⚠️" }; + format!("{}{state}", path.display()) + }; + let proxy = self + .proxy + .as_ref() + .map(|v| v.to_string()) + .unwrap_or("-".into()); + let temperature = self + .temperature + .map(|v| v.to_string()) + .unwrap_or("-".into()); + let role_name = self + .role + .as_ref() + .map(|v| v.name.to_string()) + .unwrap_or("-".into()); + let items = vec![ + ("config_file", file_info(&Config::config_file()?)), + ("roles_file", file_info(&Config::roles_file()?)), + ("messages_file", file_info(&Config::messages_file()?)), + ("role", role_name), + ("api_key", self.api_key.clone()), + ("temperature", temperature), + ("save", self.save.to_string()), + ("highlight", self.highlight.to_string()), + ("proxy", proxy), + ("dry_run", self.dry_run.to_string()), + ]; + let mut output = String::new(); + for (name, value) in items { + output.push_str(&format!("{name:<20}{value}\n")); + } + Ok(output) + } + + pub fn repl_completions(&self) -> Vec<String> { + let mut completion: Vec<String> = self + .roles + .iter() + .map(|v| format!(".role {}", v.name)) + .collect(); + + completion.extend(SET_COMPLETIONS.map(|v| v.to_string())); + completion + } + + pub fn update(&mut self, data: &str) -> Result<()> { + let parts: Vec<&str> = data.split_whitespace().collect(); + if parts.len() != 2 { + bail!("Usage: .set <key> <value>. If value is null, unset key."); + } + let key = parts[0]; + let value = parts[1]; + let unset = value == "null"; + match key { + "api_key" => { + if unset { + bail!("Error: Not allowed"); + } else { + self.api_key = value.to_string(); + } + } + "temperature" => { + if unset { + self.temperature = None; + } else { + let value = value.parse().with_context(|| "Invalid value")?; + self.temperature = Some(value); + } + } + "save" => { + let value = value.parse().with_context(|| "Invalid value")?; + self.save = value; + } + "highlight" => { + 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; + } + _ => bail!("Error: Unknown key `{key}`"), + } + Ok(()) + } + + pub fn start_conversation(&mut self) -> Result<()> { + if self.conversation.is_some() { + let ans = Confirm::new("Already in a conversation, start a new one?") + .with_default(true) + .prompt()?; + if !ans { + return Ok(()); + } + } + let mut conversation = Session::new(); + if let Some(role) = self.role.as_ref() { + conversation.add_prompt(&role.prompt); + } + self.conversation = Some(conversation); + Ok(()) + } + + pub fn end_conversation(&mut self) { + self.conversation = None; + } + + pub fn record_conversation(&mut self, input: &str, output: &str) -> Result<()> { + if let Some(conversation) = self.conversation.as_mut() { + conversation.add_conversatoin(input, output)?; + } + Ok(()) + } + + fn open_message_file(&self) -> Result<File> { + let path = Config::messages_file()?; + OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .with_context(|| format!("Failed to create/append {}", path.display())) + } + + fn ensure_no_conversation(&self) -> Result<()> { + if self.conversation.is_some() { + bail!("Error: Cannot perform this action in a conversation"); + } + Ok(()) + } + + fn load_roles(&mut self) -> Result<()> { + let path = Self::roles_file()?; + if !path.exists() { + return Ok(()); + } + let content = read_to_string(&path) + .with_context(|| format!("Failed to load roles at {}", path.display()))?; + let roles: Vec<Role> = + serde_yaml::from_str(&content).with_context(|| "Invalid roles config")?; + self.roles = roles; + Ok(()) + } +} + +#[derive(Debug, Clone, Deserialize)] +pub struct Role { + /// Role name + pub name: String, + /// Prompt text send to ai for setting up a role + pub prompt: String, + /// What sampling temperature to use, between 0 and 2 + pub temperature: Option<f64>, +} + +fn create_config_file(config_path: &Path) -> Result<()> { + let confirm_map_err = |_| anyhow!("Error with 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)?; + 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) + .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")); + } + + let ans = Confirm::new("Save chat messages") + .with_default(false) + .prompt() + .map_err(confirm_map_err)?; + if ans { + raw_config.push_str("save: true\n"); + } + + std::fs::write(config_path, raw_config).with_context(|| "Failed to write to config file")?; + Ok(()) +} + +fn highlight_value() -> bool { + true +} |