use crate::client::ChatGptClient; use crate::config::SharedConfig; use crate::print_now; use crate::render::render_stream; use super::abort::SharedAbortSignal; use anyhow::{Context, Result}; use crossbeam::channel::Sender; use crossbeam::sync::WaitGroup; use std::cell::RefCell; pub enum ReplCmd { Submit(String), SetRole(String), UpdateConfig(String), Prompt(String), ClearRole, ViewInfo, StartConversation, EndConversatoin, } pub struct ReplCmdHandler { client: ChatGptClient, config: SharedConfig, reply: RefCell, abort: SharedAbortSignal, } impl ReplCmdHandler { pub fn init( client: ChatGptClient, config: SharedConfig, abort: SharedAbortSignal, ) -> Result { let reply = RefCell::new(String::new()); Ok(Self { client, config, reply, abort, }) } pub fn handle(&self, cmd: ReplCmd) -> Result<()> { match cmd { ReplCmd::Submit(input) => { if input.is_empty() { self.reply.borrow_mut().clear(); return Ok(()); } let highlight = self.config.lock().highlight; let light_theme = self.config.lock().light_theme; let wg = WaitGroup::new(); let ret = render_stream( &input, &self.client, highlight, light_theme, true, self.abort.clone(), wg.clone(), ); wg.wait(); let buffer = ret?; self.config.lock().save_message(&input, &buffer)?; self.config.lock().save_conversation(&input, &buffer)?; *self.reply.borrow_mut() = buffer; } ReplCmd::SetRole(name) => { let output = self.config.lock().change_role(&name)?; print_now!("{}\n\n", output.trim_end()); } ReplCmd::ClearRole => { self.config.lock().clear_role()?; print_now!("\n"); } ReplCmd::Prompt(prompt) => { self.config.lock().add_prompt(&prompt)?; print_now!("\n"); } ReplCmd::ViewInfo => { let output = self.config.lock().info()?; print_now!("{}\n\n", output.trim_end()); } ReplCmd::UpdateConfig(input) => { self.config.lock().update(&input)?; print_now!("\n"); } ReplCmd::StartConversation => { self.config.lock().start_conversation()?; print_now!("\n"); } ReplCmd::EndConversatoin => { self.config.lock().end_conversation(); print_now!("\n"); } } Ok(()) } } pub struct ReplyStreamHandler { sender: Option>, buffer: String, abort: SharedAbortSignal, repl: bool, } impl ReplyStreamHandler { pub fn new( sender: Option>, repl: bool, abort: SharedAbortSignal, ) -> Self { Self { sender, abort, buffer: String::new(), repl, } } pub fn text(&mut self, text: &str) -> Result<()> { if self.buffer.is_empty() && text == "\n\n" { return Ok(()); } self.buffer.push_str(text); match self.sender.as_ref() { Some(tx) => { let ret = tx .send(ReplyStreamEvent::Text(text.to_string())) .with_context(|| "Failed to send StreamEvent:Text"); self.safe_ret(ret)?; } None => { print_now!("{}", text); } } Ok(()) } pub fn done(&mut self) -> Result<()> { match self.sender.as_ref() { Some(tx) => { let ret = tx .send(ReplyStreamEvent::Done) .with_context(|| "Failed to send StreamEvent:Done"); self.safe_ret(ret)?; } None => { if !self.buffer.ends_with('\n') { print_now!("\n") } if self.repl { print_now!("\n"); if cfg!(macos) { print_now!("\n") } } } } Ok(()) } pub fn get_buffer(&self) -> &str { &self.buffer } pub fn get_abort(&self) -> SharedAbortSignal { self.abort.clone() } fn safe_ret(&self, ret: Result<()>) -> Result<()> { if ret.is_err() && self.abort.aborted() { return Ok(()); } ret } } pub enum ReplyStreamEvent { Text(String), Done, }