summaryrefslogtreecommitdiffstats
path: root/src/repl
diff options
context:
space:
mode:
Diffstat (limited to 'src/repl')
-rw-r--r--src/repl/handler.rs161
-rw-r--r--src/repl/init.rs121
-rw-r--r--src/repl/mod.rs173
3 files changed, 455 insertions, 0 deletions
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<ReplCmdHandlerState>,
+ ctrlc: Arc<AtomicBool>,
+ render: Arc<MarkdownRender>,
+}
+
+pub struct ReplCmdHandlerState {
+ reply: String,
+ save_file: Option<File>,
+}
+
+impl ReplCmdHandler {
+ pub fn init(client: ChatGptClient, config: SharedConfig) -> Result<Self> {
+ 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<AtomicBool> {
+ self.ctrlc.clone()
+ }
+}
+
+pub struct ReplyReceiver {
+ output: String,
+ sender: Option<Sender<RenderStreamEvent>>,
+}
+
+impl ReplyReceiver {
+ pub fn new(sender: Option<Sender<RenderStreamEvent>>) -> 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<Self> {
+ 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<String> = 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<Box<FileBackedHistory>> {
+ 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<char> = 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<ReplCmdHandler>, line: String) -> Result<bool> {
+ 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 <name>", 2),
+ },
+ ".info" => {
+ handler.handle(ReplCmd::Info)?;
+ }
+ ".multiline" => {
+ let mut text = args.unwrap_or_default().to_string();
+ if text.is_empty() {
+ dump("Usage: .multiline { <your multiline content> }", 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 { <your multiline content> }.", 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::<Vec<String>>()
+ .join("\n");
+ dump(
+ format!("{head}\n\nPress Ctrl+C to abort session, Ctrl+D to exit the REPL"),
+ 2,
+ );
+}
generated by cgit v1.2.3 (git 2.39.1) at 2026-10-03 06:28:03 +0000