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/repl/handler.rs | 161 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 161 insertions(+) create mode 100644 src/repl/handler.rs (limited to 'src/repl/handler.rs') 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, +} -- cgit v1.2.3