From 444f4ebe9de8aa68afdf8d39b734a8543b9b5c4b Mon Sep 17 00:00:00 2001 From: sigoden Date: Thu, 2 Nov 2023 09:53:54 +0800 Subject: refactor: improve code quanity (#196) - rewrite Repl, remove ReplHandler - move ReplyStreamHandler to repl/ and rename it to ReplyHandler - deprecate utils::print_now - refactor session info --- src/repl/mod.rs | 277 ++++++++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 221 insertions(+), 56 deletions(-) (limited to 'src/repl/mod.rs') diff --git a/src/repl/mod.rs b/src/repl/mod.rs index eb99fda..c9c3c73 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -1,24 +1,35 @@ mod abort; -mod handler; mod highlighter; -mod init; mod prompt; mod validator; -pub use self::abort::*; -pub use self::handler::*; -pub use self::init::Repl; +pub use self::abort::{create_abort_signal, AbortSignal}; +use self::highlighter::ReplHighlighter; +use self::prompt::ReplPrompt; +use self::validator::ReplValidator; + +use crate::client::init_client; use crate::config::SharedConfig; -use crate::print_now; +use crate::render::{render_error, render_stream}; -use anyhow::Result; +use anyhow::{bail, Context, Result}; +use arboard::Clipboard; +use crossbeam::sync::WaitGroup; use fancy_regex::Regex; use lazy_static::lazy_static; use reedline::Signal; -use std::rc::Rc; +use reedline::{ + default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings, + ColumnarMenu, DefaultCompleter, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline, + ReedlineEvent, ReedlineMenu, Vi, +}; +use std::cell::RefCell; +use std::io::Read; + +const MENU_NAME: &str = "completion_menu"; -pub const REPL_COMMANDS: [(&str, &str); 14] = [ +const REPL_COMMANDS: [(&str, &str); 14] = [ (".help", "Print this help message"), (".info", "Print system info"), (".edit", "Multi-line editing (CTRL+S to finish)"), @@ -40,128 +51,284 @@ lazy_static! { static ref EDIT_RE: Regex = Regex::new(r"^\s*\.edit\s*").unwrap(); } +pub struct Repl { + config: SharedConfig, + editor: Reedline, + prompt: ReplPrompt, + abort: AbortSignal, + clipboard: std::result::Result, arboard::Error>, +} + impl Repl { - pub fn run(&mut self, config: SharedConfig) -> Result<()> { - let abort = AbortSignal::new(); - let handler = ReplCmdHandler::init(config, abort.clone())?; - print_now!("Welcome to aichat {}\n", env!("CARGO_PKG_VERSION")); - print_now!("Type \".help\" for more information.\n"); + pub fn init(config: SharedConfig) -> Result { + let commands: Vec = REPL_COMMANDS + .into_iter() + .map(|(v, _)| v.to_string()) + .collect(); + + let completer = Self::create_completer(&config, &commands); + let highlighter = ReplHighlighter::new(commands, config.clone()); + let menu = Self::create_menu(); + let edit_mode: Box = if config.read().keybindings.is_vi() { + let mut normal_keybindings = default_vi_normal_keybindings(); + let mut insert_keybindings = default_vi_insert_keybindings(); + Self::extra_keybindings(&mut normal_keybindings); + Self::extra_keybindings(&mut insert_keybindings); + Box::new(Vi::new(insert_keybindings, normal_keybindings)) + } else { + let mut keybindings = default_emacs_keybindings(); + Self::extra_keybindings(&mut keybindings); + Box::new(Emacs::new(keybindings)) + }; + let mut editor = Reedline::create() + .with_completer(Box::new(completer)) + .with_highlighter(Box::new(highlighter)) + .with_menu(menu) + .with_edit_mode(edit_mode) + .with_quick_completions(true) + .with_partial_completions(true) + .with_validator(Box::new(ReplValidator)) + .with_ansi_colors(true); + + editor.enable_bracketed_paste()?; + + let prompt = ReplPrompt::new(config.clone()); + + let abort = create_abort_signal(); + + let clipboard = Clipboard::new().map(RefCell::new); + + Ok(Self { + config, + editor, + prompt, + clipboard, + abort, + }) + } + + pub fn run(&mut self) -> Result<()> { + self.banner(); + let mut already_ctrlc = false; - let handler = Rc::new(handler); + loop { - if abort.aborted_ctrld() { + if self.abort.aborted_ctrld() { break; } - if abort.aborted_ctrlc() && !already_ctrlc { + if self.abort.aborted_ctrlc() && !already_ctrlc { already_ctrlc = true; } let sig = self.editor.read_line(&self.prompt); match sig { Ok(Signal::Success(line)) => { already_ctrlc = false; - abort.reset(); - match self.handle_line(&handler, &line) { + self.abort.reset(); + match self.handle(&line) { Ok(quit) => { if quit { break; } } Err(err) => { - let err = format!("{err:?}"); - print_now!("Error: {}\n\n", err.trim()); + render_error(err, self.config.read().highlight); } } } Ok(Signal::CtrlC) => { - abort.set_ctrlc(); + self.abort.set_ctrlc(); if already_ctrlc { break; } already_ctrlc = true; - print_now!("(To exit, press Ctrl+C again or Ctrl+D or type .exit)\n\n"); + println!("(To exit, press Ctrl+C again or Ctrl+D or type .exit)\n"); } Ok(Signal::CtrlD) => { - abort.set_ctrld(); + self.abort.set_ctrld(); break; } _ => {} } } - handler.handle(ReplCmd::ExitSession)?; + self.handle(".exit session")?; Ok(()) } - fn handle_line(&mut self, handler: &Rc, line: &str) -> Result { + fn handle(&self, line: &str) -> Result { match parse_command(line) { Some((cmd, args)) => match cmd { ".help" => { dump_repl_help(); } ".info" => match args { - Some("role") => handler.handle(ReplCmd::RoleInfo)?, - Some("session") => handler.handle(ReplCmd::SessionInfo)?, - Some(_) => unknown_command(), + Some("role") => { + let info = self.config.read().role_info()?; + println!("{}", info); + } + Some("session") => { + let info = self.config.read().session_info()?; + println!("{}", info); + } + Some(_) => unknown_command()?, None => { - handler.handle(ReplCmd::Info)?; + let output = self.config.read().sys_info()?; + println!("{}", output); } }, ".edit" => { if let Some(text) = args { - handler.handle(ReplCmd::Submit(text.to_string()))?; + self.ask(text)?; } } ".model" => match args { - Some(name) => handler.handle(ReplCmd::SetModel(name.to_string()))?, - None => print_now!("Usage: .model \n\n"), + Some(name) => { + self.config.write().set_model(name)?; + } + None => println!("Usage: .model "), }, ".role" => match args { - Some(name) => handler.handle(ReplCmd::SetRole(name.to_string()))?, - None => print_now!("Usage: .role \n\n"), + Some(name) => { + self.config.write().set_role(name)?; + } + None => println!("Usage: .role "), }, ".session" => { - handler.handle(ReplCmd::StartSession(args.map(|v| v.to_string())))?; + self.config.write().start_session(args)?; } ".set" => { - handler.handle(ReplCmd::Set(args.unwrap_or_default().to_string()))?; + if let Some(args) = args { + self.config.write().update(args)?; + } } ".copy" => { - handler.handle(ReplCmd::Copy)?; + let config = self.config.read(); + self.copy(config.last_reply()) + .with_context(|| "Failed to copy the last output")?; } ".read" => match args { - Some(file) => handler.handle(ReplCmd::ReadFile(file.to_string()))?, - None => print_now!("Usage: .read \n\n"), + Some(file) => { + let mut content = String::new(); + let mut file = + std::fs::File::open(file).with_context(|| "Unable to open file")?; + file.read_to_string(&mut content) + .with_context(|| "Unable to read file")?; + self.ask(&content)?; + } + None => println!("Usage: .read "), }, ".exit" => match args { - Some("role") => handler.handle(ReplCmd::ExitRole)?, - Some("session") => handler.handle(ReplCmd::ExitSession)?, - Some(_) => unknown_command(), + Some("role") => { + self.config.write().clear_role()?; + } + Some("session") => { + self.config.write().end_session()?; + } + Some(_) => unknown_command()?, None => { return Ok(true); } }, - // deprecated + // deprecated this command ".clear" => match args { Some("role") => { - print_now!("Deprecated. Use '.exit role' instead.\n\n"); + println!(r#"Deprecated. Use ".exit role" instead."#); } - Some("session") => { - print_now!("Deprecated. Use '.exit session' instead.\n\n"); + Some("conversation") => { + println!(r#"Deprecated. Use ".exit session" instead."#); } - _ => unknown_command(), + _ => unknown_command()?, }, - _ => unknown_command(), + _ => unknown_command()?, }, None => { - handler.handle(ReplCmd::Submit(line.to_string()))?; + self.ask(line)?; } } + println!(); + Ok(false) } + + fn ask(&self, input: &str) -> Result<()> { + if input.is_empty() { + return Ok(()); + } + self.config.read().maybe_print_send_tokens(input); + let wg = WaitGroup::new(); + let client = init_client(self.config.clone())?; + let ret = render_stream( + input, + client.as_ref(), + &self.config, + true, + self.abort.clone(), + wg.clone(), + ); + wg.wait(); + let buffer = ret?; + self.config.write().save_message(input, &buffer)?; + if self.config.read().auto_copy { + let _ = self.copy(&buffer); + } + Ok(()) + } + + fn banner(&self) { + let version = env!("CARGO_PKG_VERSION"); + print!( + r#"Welcome to aichat {version} +Type ".help" for more information. +"# + ) + } + + fn create_completer(config: &SharedConfig, commands: &[String]) -> DefaultCompleter { + let mut completion = commands.to_vec(); + completion.extend(config.read().repl_completions()); + let mut completer = + DefaultCompleter::with_inclusions(&['.', '-', '_', ':']).set_min_word_len(2); + completer.insert(completion.clone()); + completer + } + + fn extra_keybindings(keybindings: &mut 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('s'), + ReedlineEvent::Submit, + ); + } + + fn create_menu() -> ReedlineMenu { + let completion_menu = ColumnarMenu::default().with_name(MENU_NAME); + ReedlineMenu::EngineCompleter(Box::new(completion_menu)) + } + + fn copy(&self, text: &str) -> Result<()> { + if text.is_empty() { + bail!("No text") + } + match self.clipboard.as_ref() { + Err(err) => bail!("{}", err), + Ok(clip) => { + clip.borrow_mut().set_text(text)?; + Ok(()) + } + } + } } -fn unknown_command() { - print_now!("Unknown command. Try `.help`.\n\n"); +fn unknown_command() -> Result<()> { + bail!(r#"Unknown command. Type ".help" for more information."#); } fn dump_repl_help() { @@ -170,12 +337,10 @@ fn dump_repl_help() { .map(|(name, desc)| format!("{name:<24} {desc}")) .collect::>() .join("\n"); - print_now!( + println!( r###"{head} -Press Ctrl+C to abort readline, Ctrl+D to exit the REPL - -"###, +Press Ctrl+C to abort readline, Ctrl+D to exit the REPL"###, ); } -- cgit v1.2.3