mirror of https://github.com/sigoden/aichat
feat: command mode supports stream out (#31)
* feat: command mode supports stream out * update clipull/32/head
sigoden
1 year ago
committed by
GitHub
parent
c7fcdb1744
commit
360264121c
@ -0,0 +1,37 @@
|
|||||||
|
use super::MarkdownRender;
|
||||||
|
use crate::repl::{ReplyStreamEvent, SharedAbortSignal};
|
||||||
|
use crate::utils::dump;
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use crossbeam::channel::Receiver;
|
||||||
|
|
||||||
|
pub fn cmd_render_stream(rx: Receiver<ReplyStreamEvent>, abort: SharedAbortSignal) -> Result<()> {
|
||||||
|
let mut buffer = String::new();
|
||||||
|
let mut markdown_render = MarkdownRender::new();
|
||||||
|
loop {
|
||||||
|
if abort.aborted() {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
if let Ok(evt) = rx.try_recv() {
|
||||||
|
match evt {
|
||||||
|
ReplyStreamEvent::Text(text) => {
|
||||||
|
if text.contains('\n') {
|
||||||
|
let text = format!("{buffer}{text}");
|
||||||
|
let mut lines: Vec<&str> = text.split('\n').collect();
|
||||||
|
buffer = lines.pop().unwrap_or_default().to_string();
|
||||||
|
let output = lines.join("\n");
|
||||||
|
dump(markdown_render.render(&output), 1);
|
||||||
|
} else {
|
||||||
|
buffer = format!("{buffer}{text}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ReplyStreamEvent::Done => {
|
||||||
|
let output = markdown_render.render(&buffer);
|
||||||
|
dump(output, 2);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@ -1,125 +1,44 @@
|
|||||||
|
mod cmd;
|
||||||
mod markdown;
|
mod markdown;
|
||||||
|
mod repl;
|
||||||
|
|
||||||
|
use self::cmd::cmd_render_stream;
|
||||||
pub use self::markdown::MarkdownRender;
|
pub use self::markdown::MarkdownRender;
|
||||||
use crate::repl::{ReplyStreamEvent, SharedAbortSignal};
|
use self::repl::repl_render_stream;
|
||||||
|
|
||||||
|
use crate::client::ChatGptClient;
|
||||||
|
use crate::repl::{ReplyStreamHandler, SharedAbortSignal};
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use crossbeam::channel::Receiver;
|
use crossbeam::channel::unbounded;
|
||||||
use crossterm::{
|
use crossbeam::sync::WaitGroup;
|
||||||
cursor,
|
use std::thread::spawn;
|
||||||
event::{self, Event, KeyCode, KeyModifiers},
|
|
||||||
queue, style,
|
pub fn render_stream(
|
||||||
terminal::{self, disable_raw_mode, enable_raw_mode},
|
input: &str,
|
||||||
};
|
prompt: Option<String>,
|
||||||
use std::{
|
client: &ChatGptClient,
|
||||||
io::{self, Stdout, Write},
|
highlight: bool,
|
||||||
time::{Duration, Instant},
|
repl: bool,
|
||||||
};
|
|
||||||
use unicode_width::UnicodeWidthStr;
|
|
||||||
|
|
||||||
pub fn render_stream(rx: Receiver<ReplyStreamEvent>, abort: SharedAbortSignal) -> Result<()> {
|
|
||||||
enable_raw_mode()?;
|
|
||||||
let mut stdout = io::stdout();
|
|
||||||
queue!(stdout, event::DisableMouseCapture)?;
|
|
||||||
|
|
||||||
let ret = render_stream_inner(rx, abort, &mut stdout);
|
|
||||||
|
|
||||||
queue!(stdout, event::DisableMouseCapture)?;
|
|
||||||
disable_raw_mode()?;
|
|
||||||
|
|
||||||
ret
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn render_stream_inner(
|
|
||||||
rx: Receiver<ReplyStreamEvent>,
|
|
||||||
abort: SharedAbortSignal,
|
abort: SharedAbortSignal,
|
||||||
writer: &mut Stdout,
|
wg: WaitGroup,
|
||||||
) -> Result<()> {
|
) -> Result<String> {
|
||||||
let mut last_tick = Instant::now();
|
let mut stream_handler = if highlight {
|
||||||
let tick_rate = Duration::from_millis(100);
|
let (tx, rx) = unbounded();
|
||||||
let mut buffer = String::new();
|
let abort_clone = abort.clone();
|
||||||
let mut markdown_render = MarkdownRender::new();
|
spawn(move || {
|
||||||
let terminal_columns = terminal::size()?.0;
|
let _ = if repl {
|
||||||
loop {
|
repl_render_stream(rx, abort)
|
||||||
if abort.aborted() {
|
} else {
|
||||||
return Ok(());
|
cmd_render_stream(rx, abort)
|
||||||
}
|
};
|
||||||
|
drop(wg);
|
||||||
if let Ok(evt) = rx.try_recv() {
|
});
|
||||||
recover_cursor(writer, terminal_columns, &buffer)?;
|
ReplyStreamHandler::new(Some(tx), abort_clone)
|
||||||
|
|
||||||
match evt {
|
|
||||||
ReplyStreamEvent::Text(text) => {
|
|
||||||
if text.contains('\n') {
|
|
||||||
let text = format!("{buffer}{text}");
|
|
||||||
let mut lines: Vec<&str> = text.split('\n').collect();
|
|
||||||
buffer = lines.pop().unwrap_or_default().to_string();
|
|
||||||
let output = markdown_render.render(&lines.join("\n"));
|
|
||||||
for line in output.split('\n') {
|
|
||||||
queue!(
|
|
||||||
writer,
|
|
||||||
style::Print(line),
|
|
||||||
style::Print("\n"),
|
|
||||||
cursor::MoveLeft(terminal_columns),
|
|
||||||
)?;
|
|
||||||
}
|
|
||||||
queue!(writer, style::Print(&buffer),)?;
|
|
||||||
} else {
|
|
||||||
buffer = format!("{buffer}{text}");
|
|
||||||
let output = markdown_render.render_line_stateless(&buffer);
|
|
||||||
queue!(writer, style::Print(&output))?;
|
|
||||||
}
|
|
||||||
writer.flush()?;
|
|
||||||
}
|
|
||||||
ReplyStreamEvent::Done => {
|
|
||||||
let output = markdown_render.render_line_stateless(&buffer);
|
|
||||||
queue!(writer, style::Print(output), style::Print("\n"))?;
|
|
||||||
writer.flush()?;
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
let timeout = tick_rate
|
|
||||||
.checked_sub(last_tick.elapsed())
|
|
||||||
.unwrap_or_else(|| Duration::from_secs(0));
|
|
||||||
if crossterm::event::poll(timeout)? {
|
|
||||||
if let Event::Key(key) = event::read()? {
|
|
||||||
match key.code {
|
|
||||||
KeyCode::Char('c') if key.modifiers == KeyModifiers::CONTROL => {
|
|
||||||
abort.set_ctrlc();
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
KeyCode::Char('d') if key.modifiers == KeyModifiers::CONTROL => {
|
|
||||||
abort.set_ctrld();
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
_ => {}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if last_tick.elapsed() >= tick_rate {
|
|
||||||
last_tick = Instant::now();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn recover_cursor(writer: &mut Stdout, terminal_columns: u16, buffer: &str) -> Result<()> {
|
|
||||||
let buffer_rows = (buffer.width() as u16 + terminal_columns - 1) / terminal_columns;
|
|
||||||
let (_, row) = cursor::position()?;
|
|
||||||
if buffer_rows == 0 {
|
|
||||||
queue!(writer, cursor::MoveTo(0, row))?;
|
|
||||||
} else if row + 1 >= buffer_rows {
|
|
||||||
queue!(writer, cursor::MoveTo(0, row + 1 - buffer_rows))?;
|
|
||||||
} else {
|
} else {
|
||||||
queue!(
|
drop(wg);
|
||||||
writer,
|
ReplyStreamHandler::new(None, abort)
|
||||||
terminal::ScrollUp(buffer_rows - 1 - row),
|
};
|
||||||
cursor::MoveTo(0, 0)
|
client.send_message_streaming(input, prompt, &mut stream_handler)?;
|
||||||
)?;
|
let buffer = stream_handler.get_buffer();
|
||||||
}
|
Ok(buffer.to_string())
|
||||||
Ok(())
|
|
||||||
}
|
}
|
||||||
|
|||||||
@ -0,0 +1,123 @@
|
|||||||
|
use super::MarkdownRender;
|
||||||
|
use crate::repl::{ReplyStreamEvent, SharedAbortSignal};
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use crossbeam::channel::Receiver;
|
||||||
|
use crossterm::{
|
||||||
|
cursor,
|
||||||
|
event::{self, Event, KeyCode, KeyModifiers},
|
||||||
|
queue, style,
|
||||||
|
terminal::{self, disable_raw_mode, enable_raw_mode},
|
||||||
|
};
|
||||||
|
use std::{
|
||||||
|
io::{self, Stdout, Write},
|
||||||
|
time::{Duration, Instant},
|
||||||
|
};
|
||||||
|
use unicode_width::UnicodeWidthStr;
|
||||||
|
|
||||||
|
pub fn repl_render_stream(rx: Receiver<ReplyStreamEvent>, abort: SharedAbortSignal) -> Result<()> {
|
||||||
|
enable_raw_mode()?;
|
||||||
|
let mut stdout = io::stdout();
|
||||||
|
queue!(stdout, event::DisableMouseCapture)?;
|
||||||
|
|
||||||
|
let ret = repl_render_stream_inner(rx, abort, &mut stdout);
|
||||||
|
|
||||||
|
queue!(stdout, event::DisableMouseCapture)?;
|
||||||
|
disable_raw_mode()?;
|
||||||
|
|
||||||
|
ret
|
||||||
|
}
|
||||||
|
|
||||||
|
fn repl_render_stream_inner(
|
||||||
|
rx: Receiver<ReplyStreamEvent>,
|
||||||
|
abort: SharedAbortSignal,
|
||||||
|
writer: &mut Stdout,
|
||||||
|
) -> Result<()> {
|
||||||
|
let mut last_tick = Instant::now();
|
||||||
|
let tick_rate = Duration::from_millis(100);
|
||||||
|
let mut buffer = String::new();
|
||||||
|
let mut markdown_render = MarkdownRender::new();
|
||||||
|
let terminal_columns = terminal::size()?.0;
|
||||||
|
loop {
|
||||||
|
if abort.aborted() {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Ok(evt) = rx.try_recv() {
|
||||||
|
recover_cursor(writer, terminal_columns, &buffer)?;
|
||||||
|
|
||||||
|
match evt {
|
||||||
|
ReplyStreamEvent::Text(text) => {
|
||||||
|
if text.contains('\n') {
|
||||||
|
let text = format!("{buffer}{text}");
|
||||||
|
let mut lines: Vec<&str> = text.split('\n').collect();
|
||||||
|
buffer = lines.pop().unwrap_or_default().to_string();
|
||||||
|
let output = markdown_render.render(&lines.join("\n"));
|
||||||
|
for line in output.split('\n') {
|
||||||
|
queue!(
|
||||||
|
writer,
|
||||||
|
style::Print(line),
|
||||||
|
style::Print("\n"),
|
||||||
|
cursor::MoveLeft(terminal_columns),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
queue!(writer, style::Print(&buffer),)?;
|
||||||
|
} else {
|
||||||
|
buffer = format!("{buffer}{text}");
|
||||||
|
let output = markdown_render.render_line_stateless(&buffer);
|
||||||
|
queue!(writer, style::Print(&output))?;
|
||||||
|
}
|
||||||
|
writer.flush()?;
|
||||||
|
}
|
||||||
|
ReplyStreamEvent::Done => {
|
||||||
|
let output = markdown_render.render_line_stateless(&buffer);
|
||||||
|
queue!(writer, style::Print(output), style::Print("\n"))?;
|
||||||
|
writer.flush()?;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let timeout = tick_rate
|
||||||
|
.checked_sub(last_tick.elapsed())
|
||||||
|
.unwrap_or_else(|| Duration::from_secs(0));
|
||||||
|
if crossterm::event::poll(timeout)? {
|
||||||
|
if let Event::Key(key) = event::read()? {
|
||||||
|
match key.code {
|
||||||
|
KeyCode::Char('c') if key.modifiers == KeyModifiers::CONTROL => {
|
||||||
|
abort.set_ctrlc();
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
KeyCode::Char('d') if key.modifiers == KeyModifiers::CONTROL => {
|
||||||
|
abort.set_ctrld();
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if last_tick.elapsed() >= tick_rate {
|
||||||
|
last_tick = Instant::now();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn recover_cursor(writer: &mut Stdout, terminal_columns: u16, buffer: &str) -> Result<()> {
|
||||||
|
let buffer_rows = (buffer.width() as u16 + terminal_columns - 1) / terminal_columns;
|
||||||
|
let (_, row) = cursor::position()?;
|
||||||
|
if buffer_rows == 0 {
|
||||||
|
queue!(writer, cursor::MoveTo(0, row))?;
|
||||||
|
} else if row + 1 >= buffer_rows {
|
||||||
|
queue!(writer, cursor::MoveTo(0, row + 1 - buffer_rows))?;
|
||||||
|
} else {
|
||||||
|
queue!(
|
||||||
|
writer,
|
||||||
|
terminal::ScrollUp(buffer_rows - 1 - row),
|
||||||
|
cursor::MoveTo(0, 0)
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue