From 2152667fc6551129fd0bb14ca9e403728855aea3 Mon Sep 17 00:00:00 2001 From: sigoden Date: Mon, 6 Mar 2023 06:45:00 +0800 Subject: refactor: split src/repl.rs to submodules (#22) --- src/config.rs | 100 +++++++----- src/main.rs | 3 +- src/repl.rs | 432 ---------------------------------------------------- src/repl/handler.rs | 161 ++++++++++++++++++++ src/repl/init.rs | 121 +++++++++++++++ src/repl/mod.rs | 173 +++++++++++++++++++++ 6 files changed, 517 insertions(+), 473 deletions(-) delete mode 100644 src/repl.rs create mode 100644 src/repl/handler.rs create mode 100644 src/repl/init.rs create mode 100644 src/repl/mod.rs (limited to 'src') diff --git a/src/config.rs b/src/config.rs index 48b386b..c3d338c 100644 --- a/src/config.rs +++ b/src/config.rs @@ -18,7 +18,26 @@ 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: &str = "%TEMP%"; +const TEMP_ROLE_NAME: &str = "%TEMP%"; +const UPDATE_KEYS: [&str; 6] = [ + "api_key", + "temperature", + "save", + "highlight", + "proxy", + "dry_run", +]; +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 { @@ -48,14 +67,6 @@ pub struct Config { pub type SharedConfig = Arc>; impl Config { - pub const UPDATE_KEYS: [&str; 6] = [ - "api_key", - "temperature", - "save", - "highlight", - "proxy", - "dry_run", - ]; pub fn init(is_interactive: bool) -> Result { let config_path = Config::config_file()?; if is_interactive && !config_path.exists() { @@ -115,40 +126,40 @@ impl Config { Ok(file) } - pub fn save_message(&self, file: Option<&mut File>, input: &str, output: &str) { - if output.is_empty() || !self.save { - return; + pub fn save_message(&self, file: Option<&mut File>, input: &str, output: &str) -> Result<()> { + if output.is_empty() || !self.save || file.is_none() { + return Ok(()); } - if let Some(file) = file { - let timestamp = now(); - let output = match self.role.as_ref() { - None => { + let file = file.unwrap(); + let timestamp = now(); + let output = match self.role.as_ref() { + None => { + format!( + "# CHAT:[{timestamp}]\n{}\n\n--------\n{}\n--------\n\n", + input.trim(), + output.trim(), + ) + } + Some(v) => { + if v.name == TEMP_ROLE_NAME { format!( - "# CHAT:[{timestamp}]\n{}\n\n--------\n{}\n--------\n\n", + "# CHAT:[{timestamp}]\n{}\n{}\n\n--------\n{}\n--------\n\n", + v.prompt, + input.trim(), + output.trim(), + ) + } else { + format!( + "# CHAT:[{timestamp}] ({})\n{}\n\n--------\n{}\n--------\n\n", + v.name, input.trim(), output.trim(), ) } - Some(v) => { - if v.name == TEMP_ROLE { - format!( - "# CHAT:[{timestamp}]\n{}\n{}\n\n--------\n{}\n--------\n\n", - v.prompt, - input.trim(), - output.trim(), - ) - } else { - format!( - "# CHAT:[{timestamp}] ({})\n{}\n\n--------\n{}\n--------\n\n", - v.name, - input.trim(), - output.trim(), - ) - } - } - }; - let _ = file.write_all(output.as_bytes()); - } + } + }; + file.write_all(output.as_bytes()) + .with_context(|| "Failed to save message") } pub fn config_file() -> Result { @@ -180,7 +191,7 @@ impl Config { pub fn create_temp_role(&mut self, prompt: &str) -> String { self.role = Some(Role { - name: TEMP_ROLE.into(), + name: TEMP_ROLE_NAME.into(), prompt: prompt.into(), }); "Done".into() @@ -234,6 +245,17 @@ impl Config { Ok(output) } + pub fn repl_completions(&self) -> Vec { + let mut completion: Vec = 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 { @@ -280,7 +302,7 @@ impl Config { _ => { return Ok(format!( "Unknown key, valid keys are {}", - Config::UPDATE_KEYS.join(", ") + UPDATE_KEYS.join(", ") )) } } diff --git a/src/main.rs b/src/main.rs index 39a49e4..ec127f5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -68,8 +68,7 @@ fn start_directive(client: ChatGptClient, config: SharedConfig, input: &str) -> println!("{output}"); } - config.borrow().save_message(file.as_mut(), input, output); - Ok(()) + config.borrow().save_message(file.as_mut(), input, output) } fn start_interactive(client: ChatGptClient, config: SharedConfig) -> Result<()> { diff --git a/src/repl.rs b/src/repl.rs deleted file mode 100644 index 79b7bdb..0000000 --- a/src/repl.rs +++ /dev/null @@ -1,432 +0,0 @@ -use crate::client::ChatGptClient; -use crate::config::{Config, SharedConfig}; -use crate::render::{self, MarkdownRender}; -use crate::term; -use crate::utils::{copy, dump}; -use anyhow::{Context, Result}; -use crossbeam::channel::{unbounded, Sender}; -use crossbeam::sync::WaitGroup; -use reedline::{ - default_emacs_keybindings, ColumnarMenu, DefaultCompleter, DefaultPrompt, DefaultPromptSegment, - Emacs, FileBackedHistory, KeyCode, KeyModifiers, Keybindings, Reedline, ReedlineEvent, - ReedlineMenu, Signal, ValidationResult, Validator, -}; -use std::cell::RefCell; -use std::fs::File; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; -use std::thread::spawn; - -const REPL_COMMANDS: [(&str, &str); 12] = [ - (".role", "Specifies the role the AI will play"), - (".clear role", "Clear the currently selected role"), - (".prompt", "Add prompt, aka create a temporary role"), - (".history", "Print the history"), - (".clear history", "Clear the history"), - (".multiline", "Enter multiline editor mode"), - (".copy", "Copy last reply message"), - (".info", "Print the information"), - (".set", "Modify the configuration temporarily"), - (".help", "Print this help message"), - (".exit", "Exit the REPL"), - (".clear screen", "Clear the screen"), -]; - -const MENU_NAME: &str = "completion_menu"; - -pub struct Repl { - editor: Reedline, - prompt: DefaultPrompt, -} - -impl Repl { - pub fn init(config: SharedConfig) -> Result { - let completer = Self::create_completer(config); - let keybindings = Self::create_keybindings(); - let history = Self::create_history()?; - let menu = Self::create_menu(); - let edit_mode = Box::new(Emacs::new(keybindings)); - let editor = Reedline::create() - .with_completer(Box::new(completer)) - .with_history(history) - .with_menu(menu) - .with_edit_mode(edit_mode) - .with_quick_completions(true) - .with_partial_completions(true) - .with_validator(Box::new(ReplValidator { - multiline_cmds: [".multiline", ".prompt"].to_vec(), - })) - .with_ansi_colors(true); - let prompt = Self::create_prompt(); - Ok(Self { editor, prompt }) - } - - pub fn run(&mut self, handler: ReplCmdHandler) -> Result<()> { - dump( - format!("Welcome to aichat {}", env!("CARGO_PKG_VERSION")), - 1, - ); - dump("Type \".help\" for more information.", 1); - let mut current_ctrlc = false; - let handler = Arc::new(handler); - loop { - if handler.ctrlc.load(Ordering::SeqCst) { - handler.ctrlc.store(false, Ordering::SeqCst); - current_ctrlc = true - } - match self.editor.read_line(&self.prompt) { - Ok(Signal::Success(line)) => { - current_ctrlc = false; - match self.handle_line(handler.clone(), line) { - Ok(quit) => { - if quit { - break; - } - } - Err(err) => { - let err = format!("{err:?}"); - dump(err.trim(), 2); - } - } - } - Ok(Signal::CtrlC) => { - if !current_ctrlc { - current_ctrlc = true; - dump("(To exit, press Ctrl+C again or Ctrl+D or type .exit)", 2); - } else { - break; - } - } - Ok(Signal::CtrlD) => { - break; - } - _ => {} - } - } - Ok(()) - } - - fn handle_line(&mut self, handler: Arc, line: String) -> Result { - if line.starts_with('.') { - let (cmd, args) = match line.split_once(' ') { - Some((head, tail)) => (head, Some(tail.trim())), - None => (line.as_str(), None), - }; - match cmd { - ".exit" => { - return Ok(true); - } - ".help" => { - dump_repl_help(); - } - ".clear" => match args { - Some("screen") => term::clear_screen(0)?, - Some("history") => { - let history = Box::new(self.editor.history_mut()); - history.clear().with_context(|| "Failed to clear history")?; - dump("", 1); - } - Some("role") => handler.handle(ReplCmd::ClearRole)?, - _ => dump_unknown_command(), - }, - ".history" => { - self.editor.print_history()?; - dump("", 1); - } - ".role" => match args { - Some(name) => handler.handle(ReplCmd::SetRole(name.to_string()))?, - None => dump("Usage: .role ", 2), - }, - ".info" => { - handler.handle(ReplCmd::Info)?; - } - ".multiline" => { - let mut text = args.unwrap_or_default().to_string(); - if text.is_empty() { - dump("Usage: .multiline { }", 2); - } else { - if text.starts_with('{') && text.ends_with('}') { - text = text[1..text.len() - 1].to_string() - } - handler.handle(ReplCmd::Submit(text))?; - } - } - ".copy" => { - let reply = handler.get_reply(); - if reply.is_empty() { - dump("No reply messages that can be copied", 1) - } else { - copy(&reply)?; - dump("Copied", 1); - } - } - ".set" => { - handler.handle(ReplCmd::UpdateConfig(args.unwrap_or_default().to_string()))? - } - ".prompt" => { - let mut text = args.unwrap_or_default().to_string(); - if text.is_empty() { - dump("Usage: .prompt { }.", 2); - } else { - if text.starts_with('{') && text.ends_with('}') { - text = text[1..text.len() - 1].to_string() - } - handler.handle(ReplCmd::Prompt(text))?; - } - } - _ => dump_unknown_command(), - } - } else { - handler.handle(ReplCmd::Submit(line))?; - } - - Ok(false) - } - - fn create_prompt() -> DefaultPrompt { - DefaultPrompt::new(DefaultPromptSegment::Empty, DefaultPromptSegment::Empty) - } - - fn create_completer(config: SharedConfig) -> DefaultCompleter { - let mut commands: Vec = REPL_COMMANDS - .into_iter() - .map(|(v, _)| v.to_string()) - .collect(); - commands.extend( - config - .as_ref() - .borrow() - .roles - .iter() - .map(|v| format!(".role {}", v.name)), - ); - commands.extend(Config::UPDATE_KEYS.map(|v| format!(".set {v}"))); - let mut completer = DefaultCompleter::with_inclusions(&['.', '-', '_']).set_min_word_len(2); - completer.insert(commands.clone()); - completer - } - - fn create_keybindings() -> Keybindings { - let mut keybindings = default_emacs_keybindings(); - keybindings.add_binding( - KeyModifiers::NONE, - KeyCode::Tab, - ReedlineEvent::UntilFound(vec![ - ReedlineEvent::Menu(MENU_NAME.to_string()), - ReedlineEvent::MenuNext, - ]), - ); - keybindings.add_binding( - KeyModifiers::CONTROL, - KeyCode::Char('l'), - ReedlineEvent::ExecuteHostCommand(".clear screen".into()), - ); - keybindings - } - - fn create_menu() -> ReedlineMenu { - let completion_menu = ColumnarMenu::default().with_name(MENU_NAME); - ReedlineMenu::EngineCompleter(Box::new(completion_menu)) - } - - fn create_history() -> Result> { - Ok(Box::new( - FileBackedHistory::with_file(1000, Config::history_file()?) - .with_context(|| "Failed to setup history file")?, - )) - } -} -pub struct ReplValidator { - multiline_cmds: Vec<&'static str>, -} - -impl Validator for ReplValidator { - fn validate(&self, line: &str) -> ValidationResult { - if line.split('"').count() % 2 == 0 || incomplete_brackets(line, &self.multiline_cmds) { - ValidationResult::Incomplete - } else { - ValidationResult::Complete - } - } -} - -fn incomplete_brackets(line: &str, multiline_cmds: &[&str]) -> bool { - let mut balance: Vec = Vec::new(); - let line = line.trim_start(); - if !multiline_cmds.iter().any(|v| line.starts_with(v)) { - return false; - } - - for c in line.chars() { - if c == '{' { - balance.push('}'); - } else if c == '}' { - if let Some(last) = balance.last() { - if last == &c { - balance.pop(); - } - } - } - } - - !balance.is_empty() -} - -pub struct ReplCmdHandler { - client: ChatGptClient, - config: SharedConfig, - state: RefCell, - ctrlc: Arc, - render: Arc, -} - -struct ReplCmdHandlerState { - reply: String, - save_file: Option, -} - -impl ReplCmdHandler { - pub fn init(client: ChatGptClient, config: SharedConfig) -> Result { - let render = Arc::new(MarkdownRender::init()?); - let save_file = config.as_ref().borrow().open_message_file()?; - let ctrlc = Arc::new(AtomicBool::new(false)); - let state = RefCell::new(ReplCmdHandlerState { - save_file, - reply: String::new(), - }); - Ok(Self { - client, - config, - state, - ctrlc, - render, - }) - } - fn handle(&self, cmd: ReplCmd) -> Result<()> { - match cmd { - ReplCmd::Submit(input) => { - if input.is_empty() { - self.state.borrow_mut().reply.clear(); - return Ok(()); - } - let prompt = self.config.borrow().get_prompt(); - let wg = WaitGroup::new(); - let highlight = self.config.borrow().highlight; - let mut receiver = if highlight { - let (tx, rx) = unbounded(); - let ctrlc = self.ctrlc.clone(); - let wg = wg.clone(); - let render = self.render.clone(); - spawn(move || { - let _ = render::render_stream(rx, ctrlc, render); - drop(wg); - }); - ReplyReceiver::new(Some(tx)) - } else { - ReplyReceiver::new(None) - }; - self.client - .acquire_stream(&input, prompt, &mut receiver, self.ctrlc.clone())?; - self.config.borrow().save_message( - self.state.borrow_mut().save_file.as_mut(), - &input, - &receiver.output, - ); - wg.wait(); - self.state.borrow_mut().reply = receiver.output; - } - ReplCmd::SetRole(name) => { - let output = self.config.borrow_mut().change_role(&name); - dump(output.trim(), 2); - } - ReplCmd::ClearRole => { - self.config.borrow_mut().role = None; - dump("Done", 2); - } - ReplCmd::Prompt(prompt) => { - let output = self.config.borrow_mut().create_temp_role(&prompt); - dump(output.trim(), 2); - } - ReplCmd::Info => { - let output = self.config.borrow().info()?; - dump(output.trim(), 2); - } - ReplCmd::UpdateConfig(input) => { - let output = self.config.borrow_mut().update(&input)?; - dump(output.trim(), 2); - } - } - Ok(()) - } - - fn get_reply(&self) -> String { - self.state.borrow().reply.to_string() - } -} - -pub struct ReplyReceiver { - output: String, - sender: Option>, -} - -impl ReplyReceiver { - pub fn new(sender: Option>) -> Self { - Self { - output: String::new(), - sender, - } - } - - pub fn text(&mut self, text: &str) { - match self.sender.as_ref() { - Some(tx) => { - let _ = tx.send(RenderStreamEvent::Text(text.to_string())); - } - None => { - dump(text, 0); - } - } - self.output.push_str(text); - } - - pub fn done(&mut self) { - match self.sender.as_ref() { - Some(tx) => { - let _ = tx.send(RenderStreamEvent::Done); - } - None => { - dump("", 2); - } - } - } -} - -pub enum RenderStreamEvent { - Text(String), - Done, -} - -enum ReplCmd { - Submit(String), - SetRole(String), - UpdateConfig(String), - Prompt(String), - ClearRole, - Info, -} - -fn dump_repl_help() { - let head = REPL_COMMANDS - .iter() - .map(|(name, desc)| format!("{name:<15} {desc}")) - .collect::>() - .join("\n"); - dump( - format!("{head}\n\nPress Ctrl+C to abort session, Ctrl+D to exit the REPL"), - 2, - ); -} - -fn dump_unknown_command() { - dump("Unknown command. Type \".help\" for more information.", 2); -} diff --git a/src/repl/handler.rs b/src/repl/handler.rs new file mode 100644 index 0000000..23fa609 --- /dev/null +++ b/src/repl/handler.rs @@ -0,0 +1,161 @@ +use crate::client::ChatGptClient; +use crate::config::SharedConfig; +use crate::render::{self, MarkdownRender}; +use crate::utils::dump; + +use anyhow::Result; +use crossbeam::channel::{unbounded, Sender}; +use crossbeam::sync::WaitGroup; +use std::cell::RefCell; +use std::fs::File; +use std::sync::atomic::AtomicBool; +use std::sync::Arc; +use std::thread::spawn; + +pub enum ReplCmd { + Submit(String), + SetRole(String), + UpdateConfig(String), + Prompt(String), + ClearRole, + Info, +} + +pub struct ReplCmdHandler { + client: ChatGptClient, + config: SharedConfig, + state: RefCell, + ctrlc: Arc, + render: Arc, +} + +pub struct ReplCmdHandlerState { + reply: String, + save_file: Option, +} + +impl ReplCmdHandler { + pub fn init(client: ChatGptClient, config: SharedConfig) -> Result { + let render = Arc::new(MarkdownRender::init()?); + let save_file = config.as_ref().borrow().open_message_file()?; + let ctrlc = Arc::new(AtomicBool::new(false)); + let state = RefCell::new(ReplCmdHandlerState { + save_file, + reply: String::new(), + }); + Ok(Self { + client, + config, + state, + ctrlc, + render, + }) + } + + pub fn handle(&self, cmd: ReplCmd) -> Result<()> { + match cmd { + ReplCmd::Submit(input) => { + if input.is_empty() { + self.state.borrow_mut().reply.clear(); + return Ok(()); + } + let prompt = self.config.borrow().get_prompt(); + let wg = WaitGroup::new(); + let highlight = self.config.borrow().highlight; + let mut receiver = if highlight { + let (tx, rx) = unbounded(); + let ctrlc = self.ctrlc.clone(); + let wg = wg.clone(); + let render = self.render.clone(); + spawn(move || { + let _ = render::render_stream(rx, ctrlc, render); + drop(wg); + }); + ReplyReceiver::new(Some(tx)) + } else { + ReplyReceiver::new(None) + }; + self.client + .acquire_stream(&input, prompt, &mut receiver, self.ctrlc.clone())?; + self.config.borrow().save_message( + self.state.borrow_mut().save_file.as_mut(), + &input, + &receiver.output, + )?; + wg.wait(); + self.state.borrow_mut().reply = receiver.output; + } + ReplCmd::SetRole(name) => { + let output = self.config.borrow_mut().change_role(&name); + dump(output.trim(), 2); + } + ReplCmd::ClearRole => { + self.config.borrow_mut().role = None; + dump("Done", 2); + } + ReplCmd::Prompt(prompt) => { + let output = self.config.borrow_mut().create_temp_role(&prompt); + dump(output.trim(), 2); + } + ReplCmd::Info => { + let output = self.config.borrow().info()?; + dump(output.trim(), 2); + } + ReplCmd::UpdateConfig(input) => { + let output = self.config.borrow_mut().update(&input)?; + dump(output.trim(), 2); + } + } + Ok(()) + } + + pub fn get_reply(&self) -> String { + self.state.borrow().reply.to_string() + } + + pub fn get_ctrlc(&self) -> Arc { + self.ctrlc.clone() + } +} + +pub struct ReplyReceiver { + output: String, + sender: Option>, +} + +impl ReplyReceiver { + pub fn new(sender: Option>) -> Self { + Self { + output: String::new(), + sender, + } + } + + pub fn text(&mut self, text: &str) { + match self.sender.as_ref() { + Some(tx) => { + let _ = tx.send(RenderStreamEvent::Text(text.to_string())); + } + None => { + dump(text, 0); + } + } + self.output.push_str(text); + } + + pub fn done(&mut self) { + match self.sender.as_ref() { + Some(tx) => { + let _ = tx.send(RenderStreamEvent::Done); + } + None => { + dump("", 2); + } + } + } +} + +pub enum RenderStreamEvent { + Text(String), + Done, +} diff --git a/src/repl/init.rs b/src/repl/init.rs new file mode 100644 index 0000000..f640deb --- /dev/null +++ b/src/repl/init.rs @@ -0,0 +1,121 @@ +use super::Repl; + +use crate::config::{Config, SharedConfig}; + +use anyhow::{Context, Result}; +use reedline::{ + default_emacs_keybindings, ColumnarMenu, DefaultCompleter, DefaultPrompt, DefaultPromptSegment, + Emacs, FileBackedHistory, KeyCode, KeyModifiers, Keybindings, Reedline, ReedlineEvent, + ReedlineMenu, ValidationResult, Validator, +}; + +use super::REPL_COMMANDS; + +const MENU_NAME: &str = "completion_menu"; + +impl Repl { + pub fn init(config: SharedConfig) -> Result { + let multiline_commands: Vec<&'static str> = REPL_COMMANDS + .iter() + .filter(|(_, _, v)| *v) + .map(|(v, _, _)| *v) + .collect(); + let completer = Self::create_completer(config); + let keybindings = Self::create_keybindings(); + let history = Self::create_history()?; + let menu = Self::create_menu(); + let edit_mode = Box::new(Emacs::new(keybindings)); + let editor = Reedline::create() + .with_completer(Box::new(completer)) + .with_history(history) + .with_menu(menu) + .with_edit_mode(edit_mode) + .with_quick_completions(true) + .with_partial_completions(true) + .with_validator(Box::new(ReplValidator { multiline_commands })) + .with_ansi_colors(true); + let prompt = Self::create_prompt(); + Ok(Self { editor, prompt }) + } + + fn create_prompt() -> DefaultPrompt { + DefaultPrompt::new(DefaultPromptSegment::Empty, DefaultPromptSegment::Empty) + } + + fn create_completer(config: SharedConfig) -> DefaultCompleter { + let mut completion: Vec = REPL_COMMANDS + .into_iter() + .map(|(v, _, _)| v.to_string()) + .collect(); + completion.extend(config.borrow().repl_completions()); + let mut completer = DefaultCompleter::with_inclusions(&['.', '-', '_']).set_min_word_len(2); + completer.insert(completion.clone()); + completer + } + + fn create_keybindings() -> Keybindings { + let mut keybindings = default_emacs_keybindings(); + keybindings.add_binding( + KeyModifiers::NONE, + KeyCode::Tab, + ReedlineEvent::UntilFound(vec![ + ReedlineEvent::Menu(MENU_NAME.to_string()), + ReedlineEvent::MenuNext, + ]), + ); + keybindings.add_binding( + KeyModifiers::CONTROL, + KeyCode::Char('l'), + ReedlineEvent::ExecuteHostCommand(".clear screen".into()), + ); + keybindings + } + + fn create_menu() -> ReedlineMenu { + let completion_menu = ColumnarMenu::default().with_name(MENU_NAME); + ReedlineMenu::EngineCompleter(Box::new(completion_menu)) + } + + fn create_history() -> Result> { + Ok(Box::new( + FileBackedHistory::with_file(1000, Config::history_file()?) + .with_context(|| "Failed to setup history file")?, + )) + } +} + +struct ReplValidator { + multiline_commands: Vec<&'static str>, +} + +impl Validator for ReplValidator { + fn validate(&self, line: &str) -> ValidationResult { + if line.split('"').count() % 2 == 0 || incomplete_brackets(line, &self.multiline_commands) { + ValidationResult::Incomplete + } else { + ValidationResult::Complete + } + } +} + +fn incomplete_brackets(line: &str, multiline_commands: &[&str]) -> bool { + let mut balance: Vec = Vec::new(); + let line = line.trim_start(); + if !multiline_commands.iter().any(|v| line.starts_with(v)) { + return false; + } + + for c in line.chars() { + if c == '{' { + balance.push('}'); + } else if c == '}' { + if let Some(last) = balance.last() { + if last == &c { + balance.pop(); + } + } + } + } + + !balance.is_empty() +} diff --git a/src/repl/mod.rs b/src/repl/mod.rs new file mode 100644 index 0000000..6aeb233 --- /dev/null +++ b/src/repl/mod.rs @@ -0,0 +1,173 @@ +mod handler; +mod init; + +use crate::term; +use crate::utils::{copy, dump}; + +use anyhow::{Context, Result}; +use reedline::{DefaultPrompt, Reedline, Signal}; +use std::sync::atomic::Ordering; +use std::sync::Arc; + +pub use self::handler::*; + +pub const REPL_COMMANDS: [(&str, &str, bool); 12] = [ + (".info", "Print the information", false), + (".set", "Modify the configuration temporarily", false), + (".role", "Specifies the role the AI will play", false), + (".clear role", "Clear the currently selected role", false), + (".prompt", "Add prompt, aka create a temporary role", true), + (".history", "Print the history", false), + (".clear history", "Clear the history", false), + (".clear screen", "Clear the screen", false), + (".multiline", "Enter multiline editor mode", true), + (".copy", "Copy last reply message", false), + (".help", "Print this help message", false), + (".exit", "Exit the REPL", false), +]; + +pub struct Repl { + editor: Reedline, + prompt: DefaultPrompt, +} + +impl Repl { + pub fn run(&mut self, handler: ReplCmdHandler) -> Result<()> { + dump( + format!("Welcome to aichat {}", env!("CARGO_PKG_VERSION")), + 1, + ); + dump("Type \".help\" for more information.", 1); + let mut current_ctrlc = false; + let handler = Arc::new(handler); + loop { + let handler_ctrlc = handler.get_ctrlc(); + if handler_ctrlc.load(Ordering::SeqCst) { + handler_ctrlc.store(false, Ordering::SeqCst); + current_ctrlc = true + } + match self.editor.read_line(&self.prompt) { + Ok(Signal::Success(line)) => { + current_ctrlc = false; + match self.handle_line(handler.clone(), line) { + Ok(quit) => { + if quit { + break; + } + } + Err(err) => { + let err = format!("{err:?}"); + dump(err.trim(), 2); + } + } + } + Ok(Signal::CtrlC) => { + if !current_ctrlc { + current_ctrlc = true; + dump("(To exit, press Ctrl+C again or Ctrl+D or type .exit)", 2); + } else { + break; + } + } + Ok(Signal::CtrlD) => { + break; + } + _ => {} + } + } + Ok(()) + } + + fn handle_line(&mut self, handler: Arc, line: String) -> Result { + if line.starts_with('.') { + let (cmd, args) = match line.split_once(' ') { + Some((head, tail)) => (head, Some(tail.trim())), + None => (line.as_str(), None), + }; + match cmd { + ".exit" => { + return Ok(true); + } + ".help" => { + dump_repl_help(); + } + ".clear" => match args { + Some("screen") => term::clear_screen(0)?, + Some("history") => { + let history = Box::new(self.editor.history_mut()); + history.clear().with_context(|| "Failed to clear history")?; + dump("", 1); + } + Some("role") => handler.handle(ReplCmd::ClearRole)?, + _ => dump_unknown_command(), + }, + ".history" => { + self.editor.print_history()?; + dump("", 1); + } + ".role" => match args { + Some(name) => handler.handle(ReplCmd::SetRole(name.to_string()))?, + None => dump("Usage: .role ", 2), + }, + ".info" => { + handler.handle(ReplCmd::Info)?; + } + ".multiline" => { + let mut text = args.unwrap_or_default().to_string(); + if text.is_empty() { + dump("Usage: .multiline { }", 2); + } else { + if text.starts_with('{') && text.ends_with('}') { + text = text[1..text.len() - 1].to_string() + } + handler.handle(ReplCmd::Submit(text))?; + } + } + ".copy" => { + let reply = handler.get_reply(); + if reply.is_empty() { + dump("No reply messages that can be copied", 1) + } else { + copy(&reply)?; + dump("Copied", 1); + } + } + ".set" => { + handler.handle(ReplCmd::UpdateConfig(args.unwrap_or_default().to_string()))? + } + ".prompt" => { + let mut text = args.unwrap_or_default().to_string(); + if text.is_empty() { + dump("Usage: .prompt { }.", 2); + } else { + if text.starts_with('{') && text.ends_with('}') { + text = text[1..text.len() - 1].to_string() + } + handler.handle(ReplCmd::Prompt(text))?; + } + } + _ => dump_unknown_command(), + } + } else { + handler.handle(ReplCmd::Submit(line))?; + } + + Ok(false) + } +} + +fn dump_unknown_command() { + dump("Unknown command. Type \".help\" for more information.", 2); +} + +fn dump_repl_help() { + let head = REPL_COMMANDS + .iter() + .map(|(name, desc, _)| format!("{name:<15} {desc}")) + .collect::>() + .join("\n"); + dump( + format!("{head}\n\nPress Ctrl+C to abort session, Ctrl+D to exit the REPL"), + 2, + ); +} -- cgit v1.2.3