mod completer;
mod highlighter;
mod prompt;
use self::completer::ReplCompleter;
use self::highlighter::ReplHighlighter;
use self::prompt::ReplPrompt;
use crate::client::send_stream;
use crate::config::{GlobalConfig, Input, InputContext, State};
use crate::function::need_send_call_results;
use crate::render::render_error;
use crate::utils::{create_abort_signal, set_text, AbortSignal};
use anyhow::{bail, Context, Result};
use async_recursion::async_recursion;
use fancy_regex::Regex;
use lazy_static::lazy_static;
use nu_ansi_term::Color;
use reedline::{
default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
ColumnarMenu, EditCommand, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline,
ReedlineEvent, ReedlineMenu, ValidationResult, Validator, Vi,
};
use reedline::{MenuBuilder, Signal};
use std::{env, process};
lazy_static! {
static ref SPLIT_FILES_TEXT_ARGS_RE: Regex =
Regex::new(r"(?m) (-- |--\n|--\r\n|--\r|--$)").unwrap();
}
const MENU_NAME: &str = "completion_menu";
lazy_static! {
static ref REPL_COMMANDS: [ReplCommand; 16] = [
ReplCommand::new(".help", "Show this help message", State::all()),
ReplCommand::new(".info", "View system info", State::all()),
ReplCommand::new(".model", "Change the current LLM", State::all()),
ReplCommand::new(
".prompt",
"Create a temporary role using a prompt",
State::able_change_role()
),
ReplCommand::new(
".role",
"Switch to a specific role",
State::able_change_role()
),
ReplCommand::new(".info role", "View role info", State::in_role(),),
ReplCommand::new(".exit role", "Leave the role", State::in_role(),),
ReplCommand::new(".session", "Begin a chat session", State::not_in_session(),),
ReplCommand::new(".info session", "View session info", State::in_session(),),
ReplCommand::new(
".save session",
"Save the chat to file",
State::in_session(),
),
ReplCommand::new(
".clear messages",
"Erase messages in the current session",
State::unable_change_role()
),
ReplCommand::new(
".exit session",
"End the current session",
State::in_session(),
),
ReplCommand::new(".file", "Include files with the message", State::all()),
ReplCommand::new(".set", "Adjust settings", State::all()),
ReplCommand::new(".copy", "Copy the last response", State::all()),
ReplCommand::new(".exit", "Exit the REPL", State::all()),
];
static ref COMMAND_RE: Regex = Regex::new(r"^\s*(\.\S*)\s*").unwrap();
static ref MULTILINE_RE: Regex = Regex::new(r"(?s)^\s*:::\s*(.*)\s*:::\s*$").unwrap();
}
pub struct Repl {
config: GlobalConfig,
editor: Reedline,
prompt: ReplPrompt,
abort: AbortSignal,
}
impl Repl {
pub fn init(config: &GlobalConfig) -> Result<Self> {
let editor = Self::create_editor(config)?;
let prompt = ReplPrompt::new(config);
let abort = create_abort_signal();
Ok(Self {
config: config.clone(),
editor,
prompt,
abort,
})
}
pub async fn run(&mut self) -> Result<()> {
self.banner();
loop {
if self.abort.aborted_ctrld() {
break;
}
let sig = self.editor.read_line(&self.prompt);
match sig {
Ok(Signal::Success(line)) => {
self.abort.reset();
match self.handle(&line).await {
Ok(exit) => {
if exit {
break;
}
}
Err(err) => {
render_error(err, self.config.read().highlight);
println!()
}
}
}
Ok(Signal::CtrlC) => {
self.abort.set_ctrlc();
println!("(To exit, press Ctrl+D or enter \".exit\")\n");
}
Ok(Signal::CtrlD) => {
self.abort.set_ctrld();
break;
}
_ => {}
}
}
self.handle(".exit session").await?;
Ok(())
}
async fn handle(&self, mut line: &str) -> Result<bool> {
if let Ok(Some(captures)) = MULTILINE_RE.captures(line) {
if let Some(text_match) = captures.get(1) {
line = text_match.as_str();
}
}
match parse_command(line) {
Some((cmd, args)) => match cmd {
".help" => {
dump_repl_help();
}
".info" => match args {
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 => {
let output = self.config.read().system_info()?;
println!("{}", output);
}
},
".model" => match args {
Some(name) => {
self.config.write().set_model(name)?;
if self.config.read().state().is_normal() {
self.config.write().set_model_id();
}
}
None => println!("Usage: .model <name>"),
},
".prompt" => match args {
Some(text) => {
self.config.write().set_prompt(text)?;
}
None => println!("Usage: .prompt <text>..."),
},
".role" => match args {
Some(args) => match args.split_once(|c| c == '\n' || c == ' ') {
Some((name, text)) => {
let role = self.config.read().retrieve_role(name.trim())?;
let input = Input::from_str(
&self.config,
text.trim(),
Some(InputContext::role(role)),
);
ask(&self.config, self.abort.clone(), input).await?;
}
None => {
self.config.write().set_role(args)?;
}
},
None => println!(r#"Usage: .role <name> [text]..."#),
},
".session" => {
self.config.write().start_session(args)?;
}
".save" => {
match args.map(|v| match v.split_once(' ') {
Some((subcmd, args)) => (subcmd, args.trim()),
None => (v, ""),
}) {
Some(("session", name)) => {
self.config.write().save_session(name)?;
}
_ => {
println!(r#"Usage: .save session [name]"#)
}
}
}
".set" => match args {
Some(args) => {
self.config.write().update(args)?;
}
_ => {
println!("Usage: .set <key> <value>...")
}
},
".copy" => {
let config = self.config.read();
self.copy(config.last_reply())
.with_context(|| "Failed to copy the last output")?;
}
".file" => match args {
Some(args) => {
let (files, text) = split_files_text(args);
let files = shell_words::split(files).with_context(|| "Invalid args")?;
let input = Input::new(&self.config, text, files, None)?;
ask(&self.config, self.abort.clone(), input).await?;
}
None => println!("Usage: .file <files>... [-- <text>...]"),
},
".exit" => match args {
Some("role") => {
self.config.write().clear_role()?;
}
Some("session") => {
self.config.write().end_session()?;
}
Some(_) => unknown_command()?,
None => {
return Ok(true);
}
},
".clear" => match args {
Some("messages") => {
self.config.write().clear_session_messages()?;
}
_ => unknown_command()?,
},
_ => unknown_command()?,
},
None => {
let input = Input::from_str(&self.config, line, None);
ask(&self.config, self.abort.clone(), input).await?;
}
}
println!();
Ok(false)
}
fn banner(&self) {
let version = env!("CARGO_PKG_VERSION");
print!(
r#"Welcome to aichat {version}
Type ".help" for additional help.
"#
)
}
fn create_editor(config: &GlobalConfig) -> Result<Reedline> {
let completer = ReplCompleter::new(config);
let highlighter = ReplHighlighter::new(config);
let menu = Self::create_menu();
let edit_mode = Self::create_edit_mode(config);
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)
.use_bracketed_paste(true)
.with_validator(Box::new(ReplValidator))
.with_ansi_colors(true);
if let Some(cmd) = config.read().buffer_editor() {
let temp_file =
env::temp_dir().join(format!("aichat-{}.txt", chrono::Utc::now().timestamp()));
let command = process::Command::new(cmd);
editor = editor.with_buffer_editor(command, temp_file);
}
Ok(editor)
}
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::SHIFT,
KeyCode::BackTab,
ReedlineEvent::MenuPrevious,
);
keybindings.add_binding(
KeyModifiers::CONTROL,
KeyCode::Enter,
ReedlineEvent::Edit(vec![EditCommand::InsertNewline]),
);
}
fn create_edit_mode(config: &GlobalConfig) -> Box<dyn EditMode> {
let edit_mode: Box<dyn EditMode> = 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))
};
edit_mode
}
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!("Empty text")
}
set_text(text)?;
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct ReplCommand {
name: &'static str,
description: &'static str,
valid_states: Vec<State>,
}
impl ReplCommand {
fn new(name: &'static str, desc: &'static str, valid_states: Vec<State>) -> Self {
Self {
name,
description: desc,
valid_states,
}
}
fn is_valid(&self, state: &State) -> bool {
self.valid_states.contains(state)
}
}
struct ReplValidator;
impl Validator for ReplValidator {
fn validate(&self, line: &str) -> ValidationResult {
let line = line.trim();
if line.starts_with(r#":::"#) && !line[3..].ends_with(r#":::"#) {
ValidationResult::Incomplete
} else {
ValidationResult::Complete
}
}
}
#[async_recursion]
async fn ask(config: &GlobalConfig, abort: AbortSignal, input: Input) -> Result<()> {
if input.is_empty() {
return Ok(());
}
while config.read().is_compressing_session() {
std::thread::sleep(std::time::Duration::from_millis(100));
}
let client = input.create_client()?;
let (output, tool_call_results) =
send_stream(&input, client.as_ref(), config, abort.clone()).await?;
config
.write()
.save_message(&input, &output, &tool_call_results)?;
config.read().maybe_copy(&output);
if config.write().should_compress_session() {
let config = config.clone();
let color = if config.read().light_theme {
Color::LightGray
} else {
Color::DarkGray
};
print!(
"\n📢 {}{}{}\n",
color.normal().paint(
"Session compression is being activated because the current tokens exceed `"
),
color.italic().paint("compress_threshold"),
color.normal().paint("`."),
);
tokio::spawn(async move {
let _ = compress_session(&config).await;
config.write().end_compressing_session();
});
}
if need_send_call_results(&tool_call_results) {
ask(
config,
abort,
input.merge_tool_call(output, tool_call_results),
)
.await
} else {
Ok(())
}
}
fn unknown_command() -> Result<()> {
bail!(r#"Unknown command. Type ".help" for additional help."#);
}
fn dump_repl_help() {
let head = REPL_COMMANDS
.iter()
.map(|cmd| format!("{:<24} {}", cmd.name, cmd.description))
.collect::<Vec<String>>()
.join("\n");
println!(
r###"{head}
Type ::: to start multi-line editing, type ::: to finish it.
Press Ctrl+O to open an editor for editing the input buffer.
Press Ctrl+C to cancel the response, Ctrl+D to exit the REPL."###,
);
}
fn parse_command(line: &str) -> Option<(&str, Option<&str>)> {
match COMMAND_RE.captures(line) {
Ok(Some(captures)) => {
let cmd = captures.get(1)?.as_str();
let args = line[captures[0].len()..].trim();
let args = if args.is_empty() { None } else { Some(args) };
Some((cmd, args))
}
_ => None,
}
}
async fn compress_session(config: &GlobalConfig) -> Result<()> {
let input = Input::from_str(config, config.read().summarize_prompt(), None);
let client = input.create_client()?;
let summary = client.send_message(input).await?.text;
config.write().compress_session(&summary);
Ok(())
}
fn split_files_text(args: &str) -> (&str, &str) {
match SPLIT_FILES_TEXT_ARGS_RE.find(args).ok().flatten() {
Some(mat) => {
let files = &args[0..mat.start()];
let text = if mat.end() < args.len() {
&args[mat.end()..]
} else {
""
};
(files, text)
}
None => (args, ""),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_process_command_line() {
assert_eq!(parse_command(" ."), Some((".", None)));
assert_eq!(parse_command(" .role"), Some((".role", None)));
assert_eq!(parse_command(" .role "), Some((".role", None)));
assert_eq!(
parse_command(" .set dry_run true"),
Some((".set", Some("dry_run true")))
);
assert_eq!(
parse_command(" .set dry_run true "),
Some((".set", Some("dry_run true")))
);
assert_eq!(
parse_command(".prompt \nabc\n"),
Some((".prompt", Some("abc")))
);
}
#[test]
fn test_split_files_text() {
assert_eq!(split_files_text("file.txt"), ("file.txt", ""));
assert_eq!(split_files_text("file.txt --"), ("file.txt", ""));
assert_eq!(split_files_text("file.txt -- hello"), ("file.txt", "hello"));
assert_eq!(
split_files_text("file.txt --\nhello"),
("file.txt", "hello")
);
assert_eq!(
split_files_text("file.txt --\r\nhello"),
("file.txt", "hello")
);
assert_eq!(
split_files_text("file.txt --\rhello"),
("file.txt", "hello")
);
}
}