summaryrefslogtreecommitdiffstats
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/client/azure_openai.rs4
-rw-r--r--src/client/common.rs15
-rw-r--r--src/client/localai.rs4
-rw-r--r--src/client/mod.rs6
-rw-r--r--src/client/openai.rs11
-rw-r--r--src/config/mod.rs62
-rw-r--r--src/config/model_info.rs2
-rw-r--r--src/config/role.rs2
-rw-r--r--src/config/session.rs75
-rw-r--r--src/main.rs22
-rw-r--r--src/render/cmd.rs19
-rw-r--r--src/render/mod.rs84
-rw-r--r--src/render/repl.rs16
-rw-r--r--src/repl/abort.rs12
-rw-r--r--src/repl/handler.rs205
-rw-r--r--src/repl/init.rs85
-rw-r--r--src/repl/mod.rs277
-rw-r--r--src/utils/mod.rs19
18 files changed, 443 insertions, 477 deletions
diff --git a/src/client/azure_openai.rs b/src/client/azure_openai.rs
index 5603666..fe3ec0f 100644
--- a/src/client/azure_openai.rs
+++ b/src/client/azure_openai.rs
@@ -1,5 +1,7 @@
use super::openai::{openai_build_body, OPENAI_TOKENS_COUNT_FACTORS};
-use super::{AzureOpenAIClient, ExtraConfig, ModelInfo, PromptKind, PromptType, SendData};
+use super::{AzureOpenAIClient, ExtraConfig, PromptType, SendData};
+
+use crate::{config::ModelInfo, utils::PromptKind};
use anyhow::{anyhow, Result};
use async_trait::async_trait;
diff --git a/src/client/common.rs b/src/client/common.rs
index 7963c76..0d0c0e2 100644
--- a/src/client/common.rs
+++ b/src/client/common.rs
@@ -1,6 +1,7 @@
use crate::{
config::{Message, SharedConfig},
- repl::{ReplyStreamHandler, SharedAbortSignal},
+ render::ReplyHandler,
+ repl::AbortSignal,
utils::{init_tokio_runtime, prompt_input_integer, prompt_input_string, tokenize, PromptKind},
};
@@ -139,7 +140,7 @@ macro_rules! openai_compatible_client {
async fn send_message_streaming_inner(
&self,
client: &reqwest::Client,
- handler: &mut $crate::repl::ReplyStreamHandler,
+ handler: &mut $crate::render::ReplyHandler,
data: $crate::client::SendData,
) -> Result<()> {
let builder = self.request_builder(client, data)?;
@@ -201,12 +202,8 @@ pub trait Client {
})
}
- fn send_message_streaming(
- &self,
- content: &str,
- handler: &mut ReplyStreamHandler,
- ) -> Result<()> {
- async fn watch_abort(abort: SharedAbortSignal) {
+ fn send_message_streaming(&self, content: &str, handler: &mut ReplyHandler) -> Result<()> {
+ async fn watch_abort(abort: AbortSignal) {
loop {
if abort.aborted() {
break;
@@ -252,7 +249,7 @@ pub trait Client {
async fn send_message_streaming_inner(
&self,
client: &ReqwestClient,
- handler: &mut ReplyStreamHandler,
+ handler: &mut ReplyHandler,
data: SendData,
) -> Result<()>;
}
diff --git a/src/client/localai.rs b/src/client/localai.rs
index d438388..796b574 100644
--- a/src/client/localai.rs
+++ b/src/client/localai.rs
@@ -1,5 +1,7 @@
use super::openai::{openai_build_body, OPENAI_TOKENS_COUNT_FACTORS};
-use super::{ExtraConfig, LocalAIClient, ModelInfo, PromptKind, PromptType, SendData};
+use super::{ExtraConfig, LocalAIClient, PromptType, SendData};
+
+use crate::{config::ModelInfo, utils::PromptKind};
use anyhow::Result;
use async_trait::async_trait;
diff --git a/src/client/mod.rs b/src/client/mod.rs
index 5fa6146..e55055d 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -3,12 +3,6 @@ mod common;
pub use common::*;
-use crate::{
- config::{ModelInfo, TokensCountFactors},
- repl::ReplyStreamHandler,
- utils::PromptKind,
-};
-
register_client!(
(openai, "openai", OpenAI, OpenAIConfig, OpenAIClient),
(localai, "localai", LocalAI, LocalAIConfig, LocalAIClient),
diff --git a/src/client/openai.rs b/src/client/openai.rs
index a82ee9a..98de5a8 100644
--- a/src/client/openai.rs
+++ b/src/client/openai.rs
@@ -1,6 +1,11 @@
use super::{
- ExtraConfig, ModelInfo, OpenAIClient, PromptKind, PromptType, ReplyStreamHandler, SendData,
- TokensCountFactors,
+ ExtraConfig, OpenAIClient, PromptType, SendData,
+};
+
+use crate::{
+ config::{ModelInfo, TokensCountFactors},
+ render::ReplyHandler,
+ utils::PromptKind,
};
use anyhow::{anyhow, bail, Result};
@@ -88,7 +93,7 @@ pub async fn openai_send_message(builder: RequestBuilder) -> Result<String> {
pub async fn openai_send_message_streaming(
builder: RequestBuilder,
- handler: &mut ReplyStreamHandler,
+ handler: &mut ReplyHandler,
) -> Result<()> {
let mut es = builder.eventsource()?;
while let Some(event) = es.next().await {
diff --git a/src/config/mod.rs b/src/config/mod.rs
index ea45897..731cc10 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -12,7 +12,7 @@ use crate::client::{
all_models, create_client_config, list_client_types, ClientConfig, ExtraConfig, OpenAIClient,
SendData,
};
-use crate::render::RenderOptions;
+use crate::render::{MarkdownRender, RenderOptions};
use crate::utils::{get_env_name, light_theme_from_colorfgbg, now, prompt_op_err};
use anyhow::{anyhow, bail, Context, Result};
@@ -336,7 +336,7 @@ impl Config {
}
}
- pub fn info(&self) -> Result<String> {
+ pub fn sys_info(&self) -> Result<String> {
let path_info = |path: &Path| {
let state = if path.exists() { "" } else { " ⚠️" };
format!("{}{state}", path.display())
@@ -349,27 +349,63 @@ impl Config {
.clone()
.map_or_else(|| String::from("no"), |v| v.to_string());
let items = vec![
- ("config_file", path_info(&Self::config_file()?)),
- ("roles_file", path_info(&Self::roles_file()?)),
- ("messages_file", path_info(&Self::messages_file()?)),
- ("sessions_dir", path_info(&Self::sessions_dir()?)),
("model", self.model_info.full_name()),
("temperature", temperature),
+ ("dry_run", self.dry_run.to_string()),
("save", self.save.to_string()),
("highlight", self.highlight.to_string()),
- ("light_theme", self.light_theme.to_string()),
("wrap", wrap),
("wrap_code", self.wrap_code.to_string()),
- ("dry_run", self.dry_run.to_string()),
+ ("light_theme", self.light_theme.to_string()),
("keybindings", self.keybindings.stringify().into()),
+ ("config_file", path_info(&Self::config_file()?)),
+ ("roles_file", path_info(&Self::roles_file()?)),
+ ("messages_file", path_info(&Self::messages_file()?)),
+ ("sessions_dir", path_info(&Self::sessions_dir()?)),
];
- let mut output = String::new();
- for (name, value) in items {
- output.push_str(&format!("{name:<20}{value}\n"));
- }
+ let output = items
+ .iter()
+ .map(|(name, value)| format!("{name:<20}{value}"))
+ .collect::<Vec<String>>()
+ .join("\n");
Ok(output)
}
+ pub fn role_info(&self) -> Result<String> {
+ if let Some(role) = &self.role {
+ role.info()
+ } else {
+ bail!("No role")
+ }
+ }
+
+ pub fn session_info(&self) -> Result<String> {
+ if let Some(session) = &self.session {
+ let render_options = self.get_render_options()?;
+ let mut markdown_render = MarkdownRender::init(render_options)?;
+ session.render(&mut markdown_render)
+ } else {
+ bail!("No session")
+ }
+ }
+
+ pub fn info(&self) -> Result<String> {
+ if let Some(session) = &self.session {
+ session.export()
+ } else if let Some(role) = &self.role {
+ role.info()
+ } else {
+ self.sys_info()
+ }
+ }
+
+ pub fn last_reply(&self) -> &str {
+ self.last_message
+ .as_ref()
+ .map(|(_, reply)| reply.as_str())
+ .unwrap_or_default()
+ }
+
pub fn repl_completions(&self) -> Vec<String> {
let mut completion: Vec<String> = self
.roles
@@ -423,7 +459,7 @@ impl Config {
Ok(())
}
- pub fn start_session(&mut self, session: &Option<String>) -> Result<()> {
+ pub fn start_session(&mut self, session: Option<&str>) -> Result<()> {
if self.session.is_some() {
bail!("Already in a session, please use '.clear session' to exit the session first?");
}
diff --git a/src/config/model_info.rs b/src/config/model_info.rs
index 793c014..7a52e63 100644
--- a/src/config/model_info.rs
+++ b/src/config/model_info.rs
@@ -1,4 +1,4 @@
-use super::Message;
+use super::message::Message;
use crate::utils::count_tokens;
diff --git a/src/config/role.rs b/src/config/role.rs
index 16a6400..819cc12 100644
--- a/src/config/role.rs
+++ b/src/config/role.rs
@@ -19,7 +19,7 @@ impl Role {
pub fn info(&self) -> Result<String> {
let output = serde_yaml::to_string(&self)
.with_context(|| format!("Unable to show info about role {}", &self.name))?;
- Ok(output)
+ Ok(output.trim_end().to_string())
}
pub fn embedded(&self) -> bool {
diff --git a/src/config/session.rs b/src/config/session.rs
index 62c1b09..7ed5714 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -82,12 +82,12 @@ impl Session {
if let Some(temperature) = self.temperature() {
data["temperature"] = temperature.into();
}
- data["total-tokens"] = tokens.into();
+ data["total_tokens"] = tokens.into();
if let Some(max_tokens) = self.model_info.max_tokens {
- data["max-tokens"] = max_tokens.into();
+ data["max_tokens"] = max_tokens.into();
}
if percent != 0.0 {
- data["total/max-tokens"] = format!("{}%", percent).into();
+ data["total/max"] = format!("{}%", percent).into();
}
data["messages"] = json!(self.messages);
@@ -97,43 +97,46 @@ impl Session {
}
pub fn render(&self, render: &mut MarkdownRender) -> Result<String> {
- let path = self.path.clone().unwrap_or_else(|| "-".to_string());
-
- let temperature = self
- .temperature()
- .map_or_else(|| String::from("-"), |v| v.to_string());
-
- let max_tokens = self
- .model_info
- .max_tokens
- .map(|v| v.to_string())
- .unwrap_or_else(|| '-'.to_string());
-
- let items = vec![
- ("path", path),
- ("model", self.model().to_string()),
- ("temperature", temperature),
- ("max_tokens", max_tokens),
- ];
- let mut lines = vec![];
- for (name, value) in items {
- lines.push(format!("{name:<20}{value}"));
+ let mut items = vec![];
+
+ if let Some(path) = &self.path {
+ items.push(("path", path.to_string()));
}
- lines.push("".into());
- for message in &self.messages {
- match message.role {
- MessageRole::System => {
- continue;
- }
- MessageRole::Assistant => {
- lines.push(render.render(&message.content));
- lines.push("".into());
- }
- MessageRole::User => {
- lines.push(format!("{}）{}", self.name, message.content));
+
+ items.push(("model", self.model_info.full_name()));
+
+ if let Some(temperature) = self.temperature() {
+ items.push(("temperature", temperature.to_string()));
+ }
+
+ if let Some(max_tokens) = self.model_info.max_tokens {
+ items.push(("max_tokens", max_tokens.to_string()));
+ }
+
+ let mut lines: Vec<String> = items
+ .iter()
+ .map(|(name, value)| format!("{name:<20}{value}"))
+ .collect();
+
+ if !self.is_empty() {
+ lines.push("".into());
+
+ for message in &self.messages {
+ match message.role {
+ MessageRole::System => {
+ continue;
+ }
+ MessageRole::Assistant => {
+ lines.push(render.render(&message.content));
+ lines.push("".into());
+ }
+ MessageRole::User => {
+ lines.push(format!("{}）{}", self.name, message.content));
+ }
}
}
}
+
let output = lines.join("\n");
Ok(output)
}
diff --git a/src/main.rs b/src/main.rs
index 6492b4a..b801e47 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -17,7 +17,7 @@ use crossbeam::sync::WaitGroup;
use is_terminal::IsTerminal;
use parking_lot::RwLock;
use render::{render_stream, MarkdownRender};
-use repl::{AbortSignal, Repl};
+use repl::{create_abort_signal, Repl};
use std::io::{stdin, Read};
use std::sync::Arc;
use std::{io::stdout, process::exit};
@@ -59,7 +59,9 @@ fn main() -> Result<()> {
config.write().set_role(name)?;
}
if let Some(session) = &cli.session {
- config.write().start_session(session)?;
+ config
+ .write()
+ .start_session(session.as_ref().map(|v| v.as_str()))?;
}
if let Some(model) = &cli.model {
config.write().set_model(model)?;
@@ -68,14 +70,8 @@ fn main() -> Result<()> {
config.write().highlight = false;
}
if cli.info {
- let info = if let Some(session) = &config.read().session {
- session.export()?
- } else if let Some(role) = &config.read().role {
- role.info()?
- } else {
- config.read().info()?
- };
- println!("{info}");
+ let info = config.read().info()?;
+ println!("{}", info);
exit(0);
}
let no_stream = cli.no_stream;
@@ -116,12 +112,12 @@ fn start_directive(
output
} else {
let wg = WaitGroup::new();
- let abort = AbortSignal::new();
+ let abort = create_abort_signal();
let abort_clone = abort.clone();
ctrlc::set_handler(move || {
abort_clone.set_ctrlc();
})
- .expect("Error setting Ctrl-C handler");
+ .expect("Failed to setting Ctrl-C handler");
let output = render_stream(input, client, config, false, abort, wg.clone())?;
wg.wait();
output
@@ -132,5 +128,5 @@ fn start_directive(
fn start_interactive(config: SharedConfig) -> Result<()> {
cl100k_base_singleton();
let mut repl: Repl = Repl::init(config.clone())?;
- repl.run(config)
+ repl.run()
}
diff --git a/src/render/cmd.rs b/src/render/cmd.rs
index 115170b..e09bd71 100644
--- a/src/render/cmd.rs
+++ b/src/render/cmd.rs
@@ -1,7 +1,6 @@
-use super::MarkdownRender;
+use super::{MarkdownRender, ReplyEvent};
-use crate::print_now;
-use crate::repl::{ReplyStreamEvent, SharedAbortSignal};
+use crate::repl::AbortSignal;
use crate::utils::{spaces, split_line_sematic, split_line_tail};
use anyhow::Result;
@@ -9,9 +8,9 @@ use crossbeam::channel::Receiver;
use textwrap::core::display_width;
pub fn cmd_render_stream(
- rx: &Receiver<ReplyStreamEvent>,
+ rx: &Receiver<ReplyEvent>,
render: &mut MarkdownRender,
- abort: &SharedAbortSignal,
+ abort: &AbortSignal,
) -> Result<()> {
let mut buffer = String::new();
let mut col = 0;
@@ -21,14 +20,14 @@ pub fn cmd_render_stream(
}
if let Ok(evt) = rx.try_recv() {
match evt {
- ReplyStreamEvent::Text(text) => {
+ ReplyEvent::Text(text) => {
if text.contains('\n') {
let text = format!("{buffer}{text}");
let (head, tail) = split_line_tail(&text);
buffer = tail.to_string();
let input = format!("{}{head}", spaces(col));
let output = render.render(&input);
- print_now!("{}\n", &output[col..]);
+ println!("{}", &output[col..]);
col = 0;
} else {
buffer = format!("{buffer}{text}");
@@ -51,14 +50,14 @@ pub fn cmd_render_stream(
col += display_width(output);
}
}
- print_now!("{}", output);
+ print!("{}", output);
}
}
}
}
- ReplyStreamEvent::Done => {
+ ReplyEvent::Done => {
let input = format!("{}{buffer}", spaces(col));
- print_now!("{}\n", render.render(&input));
+ println!("{}", render.render(&input));
break;
}
}
diff --git a/src/render/mod.rs b/src/render/mod.rs
index 03b67ea..27f9e81 100644
--- a/src/render/mod.rs
+++ b/src/render/mod.rs
@@ -8,12 +8,12 @@ use self::repl::repl_render_stream;
use crate::client::Client;
use crate::config::SharedConfig;
-use crate::print_now;
-use crate::repl::{ReplyStreamHandler, SharedAbortSignal};
+use crate::repl::AbortSignal;
-use anyhow::Result;
-use crossbeam::channel::unbounded;
+use anyhow::{Context, Result};
+use crossbeam::channel::{unbounded, Sender};
use crossbeam::sync::WaitGroup;
+use nu_ansi_term::{Color, Style};
use std::thread::spawn;
pub fn render_stream(
@@ -21,13 +21,14 @@ pub fn render_stream(
client: &dyn Client,
config: &SharedConfig,
repl: bool,
- abort: SharedAbortSignal,
+ abort: AbortSignal,
wg: WaitGroup,
) -> Result<String> {
let render_options = config.read().get_render_options()?;
let mut stream_handler = {
let (tx, rx) = unbounded();
let abort_clone = abort.clone();
+ let highlight = config.read().highlight;
spawn(move || {
let run = move || {
if repl {
@@ -39,14 +40,81 @@ pub fn render_stream(
}
};
if let Err(err) = run() {
- let err = format!("{err:?}");
- print_now!("\n{}\n\n", err.trim());
+ render_error(err, highlight);
}
drop(wg);
});
- ReplyStreamHandler::new(tx, abort_clone)
+ ReplyHandler::new(tx, abort_clone)
};
client.send_message_streaming(input, &mut stream_handler)?;
let buffer = stream_handler.get_buffer();
Ok(buffer.to_string())
}
+
+pub fn render_error(err: anyhow::Error, highlight: bool) {
+ let err = format!("{err:?}\n");
+ if highlight {
+ let style = Style::new().fg(Color::Red);
+ println!("{}", style.paint(err.trim()));
+ } else {
+ println!("{}", err.trim());
+ }
+}
+
+pub struct ReplyHandler {
+ sender: Sender<ReplyEvent>,
+ buffer: String,
+ abort: AbortSignal,
+}
+
+impl ReplyHandler {
+ pub fn new(sender: Sender<ReplyEvent>, abort: AbortSignal) -> Self {
+ Self {
+ sender,
+ abort,
+ buffer: String::new(),
+ }
+ }
+
+ pub fn text(&mut self, text: &str) -> Result<()> {
+ if self.buffer.is_empty() && text == "\n\n" {
+ return Ok(());
+ }
+ self.buffer.push_str(text);
+ let ret = self
+ .sender
+ .send(ReplyEvent::Text(text.to_string()))
+ .with_context(|| "Failed to send ReplyEvent:Text");
+ self.safe_ret(ret)?;
+ Ok(())
+ }
+
+ pub fn done(&mut self) -> Result<()> {
+ let ret = self
+ .sender
+ .send(ReplyEvent::Done)
+ .with_context(|| "Failed to send ReplyEvent::Done");
+ self.safe_ret(ret)?;
+ Ok(())
+ }
+
+ pub fn get_buffer(&self) -> &str {
+ &self.buffer
+ }
+
+ pub fn get_abort(&self) -> AbortSignal {
+ self.abort.clone()
+ }
+
+ fn safe_ret(&self, ret: Result<()>) -> Result<()> {
+ if ret.is_err() && self.abort.aborted() {
+ return Ok(());
+ }
+ ret
+ }
+}
+
+pub enum ReplyEvent {
+ Text(String),
+ Done,
+}
diff --git a/src/render/repl.rs b/src/render/repl.rs
index 7363923..efad0e7 100644
--- a/src/render/repl.rs
+++ b/src/render/repl.rs
@@ -1,6 +1,6 @@
-use super::MarkdownRender;
+use super::{MarkdownRender, ReplyEvent};
-use crate::repl::{ReplyStreamEvent, SharedAbortSignal};
+use crate::repl::AbortSignal;
use crate::utils::split_line_tail;
use anyhow::Result;
@@ -18,9 +18,9 @@ use std::{
use textwrap::core::display_width;
pub fn repl_render_stream(
- rx: &Receiver<ReplyStreamEvent>,
+ rx: &Receiver<ReplyEvent>,
render: &mut MarkdownRender,
- abort: &SharedAbortSignal,
+ abort: &AbortSignal,
) -> Result<()> {
enable_raw_mode()?;
let mut stdout = io::stdout();
@@ -33,9 +33,9 @@ pub fn repl_render_stream(
}
fn repl_render_stream_inner(
- rx: &Receiver<ReplyStreamEvent>,
+ rx: &Receiver<ReplyEvent>,
render: &mut MarkdownRender,
- abort: &SharedAbortSignal,
+ abort: &AbortSignal,
writer: &mut Stdout,
) -> Result<()> {
let mut last_tick = Instant::now();
@@ -51,7 +51,7 @@ fn repl_render_stream_inner(
if let Ok(evt) = rx.try_recv() {
match evt {
- ReplyStreamEvent::Text(text) => {
+ ReplyEvent::Text(text) => {
let (col, mut row) = cursor::position()?;
// fix unexpected duplicate lines on kitty, see https://github.com/sigoden/aichat/issues/105
@@ -95,7 +95,7 @@ fn repl_render_stream_inner(
writer.flush()?;
}
- ReplyStreamEvent::Done => {
+ ReplyEvent::Done => {
#[cfg(target_os = "windows")]
let eol = "\n\n";
#[cfg(not(target_os = "windows"))]
diff --git a/src/repl/abort.rs b/src/repl/abort.rs
index f76abb6..af58b35 100644
--- a/src/repl/abort.rs
+++ b/src/repl/abort.rs
@@ -3,15 +3,19 @@ use std::sync::{
Arc,
};
-pub type SharedAbortSignal = Arc<AbortSignal>;
+pub type AbortSignal = Arc<AbortSignalInner>;
-pub struct AbortSignal {
+pub struct AbortSignalInner {
ctrlc: AtomicBool,
ctrld: AtomicBool,
}
-impl AbortSignal {
- pub fn new() -> SharedAbortSignal {
+pub fn create_abort_signal() -> AbortSignal {
+ AbortSignalInner::new()
+}
+
+impl AbortSignalInner {
+ pub fn new() -> AbortSignal {
Arc::new(Self {
ctrlc: AtomicBool::new(false),
ctrld: AtomicBool::new(false),
diff --git a/src/repl/handler.rs b/src/repl/handler.rs
deleted file mode 100644
index 94f91a1..0000000
--- a/src/repl/handler.rs
+++ /dev/null
@@ -1,205 +0,0 @@
-use crate::client::init_client;
-use crate::config::SharedConfig;
-use crate::print_now;
-use crate::render::{render_stream, MarkdownRender};
-use std::fs;
-use std::io::Read;
-
-use super::abort::SharedAbortSignal;
-
-use anyhow::{bail, Context, Result};
-use arboard::Clipboard;
-use crossbeam::channel::Sender;
-use crossbeam::sync::WaitGroup;
-use std::cell::RefCell;
-
-pub enum ReplCmd {
- Submit(String),
- Info,
- RoleInfo,
- SessionInfo,
- SetModel(String),
- SetRole(String),
- ExitRole,
- StartSession(Option<String>),
- ExitSession,
- Set(String),
- Copy,
- ReadFile(String),
-}
-
-pub struct ReplCmdHandler {
- config: SharedConfig,
- abort: SharedAbortSignal,
- clipboard: std::result::Result<RefCell<Clipboard>, arboard::Error>,
-}
-
-impl ReplCmdHandler {
- pub fn init(config: SharedConfig, abort: SharedAbortSignal) -> Result<Self> {
- let clipboard = Clipboard::new().map(RefCell::new);
- Ok(Self {
- config,
- abort,
- clipboard,
- })
- }
-
- pub fn handle(&self, cmd: ReplCmd) -> Result<()> {
- match cmd {
- ReplCmd::Submit(input) => {
- if input.is_empty() {
- return Ok(());
- }
- self.config.read().maybe_print_send_tokens(&input);
- let wg = WaitGroup::new();
- let client = init_client(self.config.clone())?;
- let ret = render_stream(
- &input,
- client.as_ref(),
- &self.config,
- true,
- self.abort.clone(),
- wg.clone(),
- );
- wg.wait();
- let buffer = ret?;
- self.config.write().save_message(&input, &buffer)?;
- if self.config.read().auto_copy {
- let _ = self.copy(&buffer);
- }
- }
- ReplCmd::Info => {
- let output = self.config.read().info()?;
- print_now!("{}\n\n", output.trim_end());
- }
- ReplCmd::SetModel(name) => {
- self.config.write().set_model(&name)?;
- print_now!("\n");
- }
- ReplCmd::SetRole(name) => {
- self.config.write().set_role(&name)?;
- print_now!("\n");
- }
- ReplCmd::RoleInfo => {
- if let Some(role) = &self.config.read().role {
- print_now!("{}\n\n", role.info()?);
- } else {
- bail!("No role")
- }
- }
- ReplCmd::ExitRole => {
- self.config.write().clear_role()?;
- print_now!("\n");
- }
- ReplCmd::StartSession(name) => {
- self.config.write().start_session(&name)?;
- print_now!("\n");
- }
- ReplCmd::SessionInfo => {
- if let Some(session) = &self.config.read().session {
- let render_options = self.config.read().get_render_options()?;
- let mut markdown_render = MarkdownRender::init(render_options)?;
- print_now!("{}\n\n", session.render(&mut markdown_render)?);
- } else {
- bail!("No session")
- }
- }
- ReplCmd::ExitSession => {
- self.config.write().end_session()?;
- print_now!("\n");
- }
- ReplCmd::Set(input) => {
- self.config.write().update(&input)?;
- print_now!("\n");
- }
- ReplCmd::Copy => {
- let reply = self
- .config
- .read()
- .last_message
- .as_ref()
- .map(|v| v.1.clone())
- .unwrap_or_default();
- self.copy(&reply)
- .with_context(|| "Failed to copy the last output")?;
- print_now!("\n");
- }
- ReplCmd::ReadFile(file) => {
- let mut contents = String::new();
- let mut file = fs::File::open(file).with_context(|| "Unable to open file")?;
- file.read_to_string(&mut contents)
- .with_context(|| "Unable to read file")?;
- self.handle(ReplCmd::Submit(contents))?;
- }
- }
- Ok(())
- }
-
- fn copy(&self, text: &str) -> Result<()> {
- match self.clipboard.as_ref() {
- Err(err) => bail!("{}", err),
- Ok(clip) => {
- clip.borrow_mut().set_text(text)?;
- Ok(())
- }
- }
- }
-}
-
-pub struct ReplyStreamHandler {
- sender: Sender<ReplyStreamEvent>,
- buffer: String,
- abort: SharedAbortSignal,
-}
-
-impl ReplyStreamHandler {
- pub fn new(sender: Sender<ReplyStreamEvent>, abort: SharedAbortSignal) -> Self {
- Self {
- sender,
- abort,
- buffer: String::new(),
- }
- }
-
- pub fn text(&mut self, text: &str) -> Result<()> {
- if self.buffer.is_empty() && text == "\n\n" {
- return Ok(());
- }
- self.buffer.push_str(text);
- let ret = self
- .sender
- .send(ReplyStreamEvent::Text(text.to_string()))
- .with_context(|| "Failed to send StreamEvent:Text");
- self.safe_ret(ret)?;
- Ok(())
- }
-
- pub fn done(&mut self) -> Result<()> {
- let ret = self
- .sender
- .send(ReplyStreamEvent::Done)
- .with_context(|| "Failed to send StreamEvent:Done");
- self.safe_ret(ret)?;
- 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,
-}
diff --git a/src/repl/init.rs b/src/repl/init.rs
deleted file mode 100644
index fbbcf22..0000000
--- a/src/repl/init.rs
+++ /dev/null
@@ -1,85 +0,0 @@
-use super::{
- highlighter::ReplHighlighter, prompt::ReplPrompt, validator::ReplValidator, REPL_COMMANDS,
-};
-
-use crate::config::SharedConfig;
-
-use anyhow::Result;
-use reedline::{
- default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
- ColumnarMenu, DefaultCompleter, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline,
- ReedlineEvent, ReedlineMenu, Vi,
-};
-
-const MENU_NAME: &str = "completion_menu";
-
-pub struct Repl {
- pub(crate) editor: Reedline,
- pub(crate) prompt: ReplPrompt,
-}
-
-impl Repl {
- pub fn init(config: SharedConfig) -> Result<Self> {
- let commands: Vec<String> = REPL_COMMANDS
- .into_iter()
- .map(|(v, _)| v.to_string())
- .collect();
-
- let completer = Self::create_completer(&config, &commands);
- let highlighter = ReplHighlighter::new(commands, config.clone());
- let menu = Self::create_menu();
- 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))
- };
- 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)
- .with_validator(Box::new(ReplValidator))
- .with_ansi_colors(true);
- editor.enable_bracketed_paste()?;
- let prompt = ReplPrompt::new(config);
- Ok(Self { editor, prompt })
- }
-
- fn create_completer(config: &SharedConfig, commands: &[String]) -> DefaultCompleter {
- let mut completion = commands.to_vec();
- completion.extend(config.read().repl_completions());
- let mut completer =
- DefaultCompleter::with_inclusions(&['.', '-', '_', ':']).set_min_word_len(2);
- completer.insert(completion.clone());
- completer
- }
-
- 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::CONTROL,
- KeyCode::Char('s'),
- ReedlineEvent::Submit,
- );
- }
-
- fn create_menu() -> ReedlineMenu {
- let completion_menu = ColumnarMenu::default().with_name(MENU_NAME);
- ReedlineMenu::EngineCompleter(Box::new(completion_menu))
- }
-}
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index eb99fda..c9c3c73 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -1,24 +1,35 @@
mod abort;
-mod handler;
mod highlighter;
-mod init;
mod prompt;
mod validator;
-pub use self::abort::*;
-pub use self::handler::*;
-pub use self::init::Repl;
+pub use self::abort::{create_abort_signal, AbortSignal};
+use self::highlighter::ReplHighlighter;
+use self::prompt::ReplPrompt;
+use self::validator::ReplValidator;
+
+use crate::client::init_client;
use crate::config::SharedConfig;
-use crate::print_now;
+use crate::render::{render_error, render_stream};
-use anyhow::Result;
+use anyhow::{bail, Context, Result};
+use arboard::Clipboard;
+use crossbeam::sync::WaitGroup;
use fancy_regex::Regex;
use lazy_static::lazy_static;
use reedline::Signal;
-use std::rc::Rc;
+use reedline::{
+ default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
+ ColumnarMenu, DefaultCompleter, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline,
+ ReedlineEvent, ReedlineMenu, Vi,
+};
+use std::cell::RefCell;
+use std::io::Read;
+
+const MENU_NAME: &str = "completion_menu";
-pub const REPL_COMMANDS: [(&str, &str); 14] = [
+const REPL_COMMANDS: [(&str, &str); 14] = [
(".help", "Print this help message"),
(".info", "Print system info"),
(".edit", "Multi-line editing (CTRL+S to finish)"),
@@ -40,128 +51,284 @@ lazy_static! {
static ref EDIT_RE: Regex = Regex::new(r"^\s*\.edit\s*").unwrap();
}
+pub struct Repl {
+ config: SharedConfig,
+ editor: Reedline,
+ prompt: ReplPrompt,
+ abort: AbortSignal,
+ clipboard: std::result::Result<RefCell<Clipboard>, arboard::Error>,
+}
+
impl Repl {
- pub fn run(&mut self, config: SharedConfig) -> Result<()> {
- let abort = AbortSignal::new();
- let handler = ReplCmdHandler::init(config, abort.clone())?;
- print_now!("Welcome to aichat {}\n", env!("CARGO_PKG_VERSION"));
- print_now!("Type \".help\" for more information.\n");
+ pub fn init(config: SharedConfig) -> Result<Self> {
+ let commands: Vec<String> = REPL_COMMANDS
+ .into_iter()
+ .map(|(v, _)| v.to_string())
+ .collect();
+
+ let completer = Self::create_completer(&config, &commands);
+ let highlighter = ReplHighlighter::new(commands, config.clone());
+ let menu = Self::create_menu();
+ 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))
+ };
+ 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)
+ .with_validator(Box::new(ReplValidator))
+ .with_ansi_colors(true);
+
+ editor.enable_bracketed_paste()?;
+
+ let prompt = ReplPrompt::new(config.clone());
+
+ let abort = create_abort_signal();
+
+ let clipboard = Clipboard::new().map(RefCell::new);
+
+ Ok(Self {
+ config,
+ editor,
+ prompt,
+ clipboard,
+ abort,
+ })
+ }
+
+ pub fn run(&mut self) -> Result<()> {
+ self.banner();
+
let mut already_ctrlc = false;
- let handler = Rc::new(handler);
+
loop {
- if abort.aborted_ctrld() {
+ if self.abort.aborted_ctrld() {
break;
}
- if abort.aborted_ctrlc() && !already_ctrlc {
+ if self.abort.aborted_ctrlc() && !already_ctrlc {
already_ctrlc = true;
}
let sig = self.editor.read_line(&self.prompt);
match sig {
Ok(Signal::Success(line)) => {
already_ctrlc = false;
- abort.reset();
- match self.handle_line(&handler, &line) {
+ self.abort.reset();
+ match self.handle(&line) {
Ok(quit) => {
if quit {
break;
}
}
Err(err) => {
- let err = format!("{err:?}");
- print_now!("Error: {}\n\n", err.trim());
+ render_error(err, self.config.read().highlight);
}
}
}
Ok(Signal::CtrlC) => {
- abort.set_ctrlc();
+ self.abort.set_ctrlc();
if already_ctrlc {
break;
}
already_ctrlc = true;
- print_now!("(To exit, press Ctrl+C again or Ctrl+D or type .exit)\n\n");
+ println!("(To exit, press Ctrl+C again or Ctrl+D or type .exit)\n");
}
Ok(Signal::CtrlD) => {
- abort.set_ctrld();
+ self.abort.set_ctrld();
break;
}
_ => {}
}
}
- handler.handle(ReplCmd::ExitSession)?;
+ self.handle(".exit session")?;
Ok(())
}
- fn handle_line(&mut self, handler: &Rc<ReplCmdHandler>, line: &str) -> Result<bool> {
+ fn handle(&self, line: &str) -> Result<bool> {
match parse_command(line) {
Some((cmd, args)) => match cmd {
".help" => {
dump_repl_help();
}
".info" => match args {
- Some("role") => handler.handle(ReplCmd::RoleInfo)?,
- Some("session") => handler.handle(ReplCmd::SessionInfo)?,
- Some(_) => unknown_command(),
+ 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 => {
- handler.handle(ReplCmd::Info)?;
+ let output = self.config.read().sys_info()?;
+ println!("{}", output);
}
},
".edit" => {
if let Some(text) = args {
- handler.handle(ReplCmd::Submit(text.to_string()))?;
+ self.ask(text)?;
}
}
".model" => match args {
- Some(name) => handler.handle(ReplCmd::SetModel(name.to_string()))?,
- None => print_now!("Usage: .model <name>\n\n"),
+ Some(name) => {
+ self.config.write().set_model(name)?;
+ }
+ None => println!("Usage: .model <name>"),
},
".role" => match args {
- Some(name) => handler.handle(ReplCmd::SetRole(name.to_string()))?,
- None => print_now!("Usage: .role <name>\n\n"),
+ Some(name) => {
+ self.config.write().set_role(name)?;
+ }
+ None => println!("Usage: .role <name>"),
},
".session" => {
- handler.handle(ReplCmd::StartSession(args.map(|v| v.to_string())))?;
+ self.config.write().start_session(args)?;
}
".set" => {
- handler.handle(ReplCmd::Set(args.unwrap_or_default().to_string()))?;
+ if let Some(args) = args {
+ self.config.write().update(args)?;
+ }
}
".copy" => {
- handler.handle(ReplCmd::Copy)?;
+ let config = self.config.read();
+ self.copy(config.last_reply())
+ .with_context(|| "Failed to copy the last output")?;
}
".read" => match args {
- Some(file) => handler.handle(ReplCmd::ReadFile(file.to_string()))?,
- None => print_now!("Usage: .read <file name>\n\n"),
+ Some(file) => {
+ let mut content = String::new();
+ let mut file =
+ std::fs::File::open(file).with_context(|| "Unable to open file")?;
+ file.read_to_string(&mut content)
+ .with_context(|| "Unable to read file")?;
+ self.ask(&content)?;
+ }
+ None => println!("Usage: .read <textfile>"),
},
".exit" => match args {
- Some("role") => handler.handle(ReplCmd::ExitRole)?,
- Some("session") => handler.handle(ReplCmd::ExitSession)?,
- Some(_) => unknown_command(),
+ Some("role") => {
+ self.config.write().clear_role()?;
+ }
+ Some("session") => {
+ self.config.write().end_session()?;
+ }
+ Some(_) => unknown_command()?,
None => {
return Ok(true);
}
},
- // deprecated
+ // deprecated this command
".clear" => match args {
Some("role") => {
- print_now!("Deprecated. Use '.exit role' instead.\n\n");
+ println!(r#"Deprecated. Use ".exit role" instead."#);
}
- Some("session") => {
- print_now!("Deprecated. Use '.exit session' instead.\n\n");
+ Some("conversation") => {
+ println!(r#"Deprecated. Use ".exit session" instead."#);
}
- _ => unknown_command(),
+ _ => unknown_command()?,
},
- _ => unknown_command(),
+ _ => unknown_command()?,
},
None => {
- handler.handle(ReplCmd::Submit(line.to_string()))?;
+ self.ask(line)?;
}
}
+ println!();
+
Ok(false)
}
+
+ fn ask(&self, input: &str) -> Result<()> {
+ if input.is_empty() {
+ return Ok(());
+ }
+ self.config.read().maybe_print_send_tokens(input);
+ let wg = WaitGroup::new();
+ let client = init_client(self.config.clone())?;
+ let ret = render_stream(
+ input,
+ client.as_ref(),
+ &self.config,
+ true,
+ self.abort.clone(),
+ wg.clone(),
+ );
+ wg.wait();
+ let buffer = ret?;
+ self.config.write().save_message(input, &buffer)?;
+ if self.config.read().auto_copy {
+ let _ = self.copy(&buffer);
+ }
+ Ok(())
+ }
+
+ fn banner(&self) {
+ let version = env!("CARGO_PKG_VERSION");
+ print!(
+ r#"Welcome to aichat {version}
+Type ".help" for more information.
+"#
+ )
+ }
+
+ fn create_completer(config: &SharedConfig, commands: &[String]) -> DefaultCompleter {
+ let mut completion = commands.to_vec();
+ completion.extend(config.read().repl_completions());
+ let mut completer =
+ DefaultCompleter::with_inclusions(&['.', '-', '_', ':']).set_min_word_len(2);
+ completer.insert(completion.clone());
+ completer
+ }
+
+ 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::CONTROL,
+ KeyCode::Char('s'),
+ ReedlineEvent::Submit,
+ );
+ }
+
+ 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!("No text")
+ }
+ match self.clipboard.as_ref() {
+ Err(err) => bail!("{}", err),
+ Ok(clip) => {
+ clip.borrow_mut().set_text(text)?;
+ Ok(())
+ }
+ }
+ }
}
-fn unknown_command() {
- print_now!("Unknown command. Try `.help`.\n\n");
+fn unknown_command() -> Result<()> {
+ bail!(r#"Unknown command. Type ".help" for more information."#);
}
fn dump_repl_help() {
@@ -170,12 +337,10 @@ fn dump_repl_help() {
.map(|(name, desc)| format!("{name:<24} {desc}"))
.collect::<Vec<String>>()
.join("\n");
- print_now!(
+ println!(
r###"{head}
-Press Ctrl+C to abort readline, Ctrl+D to exit the REPL
-
-"###,
+Press Ctrl+C to abort readline, Ctrl+D to exit the REPL"###,
);
}
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index ab922e3..938e07c 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -6,24 +6,9 @@ pub use self::prompt_input::*;
pub use self::split_line::*;
pub use self::tiktoken::cl100k_base_singleton;
-use chrono::prelude::*;
-use std::io::{stdout, Write};
-
-#[macro_export]
-macro_rules! print_now {
- ($($arg:tt)*) => {
- $crate::utils::print_now(&format!($($arg)*))
- };
-}
-
-pub fn print_now<T: ToString>(text: &T) {
- print!("{}", text.to_string());
- let _ = stdout().flush();
-}
-
pub fn now() -> String {
- let now = Local::now();
- now.to_rfc3339_opts(SecondsFormat::Secs, false)
+ let now = chrono::Local::now();
+ now.to_rfc3339_opts(chrono::SecondsFormat::Secs, false)
}
pub fn get_env_name(key: &str) -> String {
generated by cgit v1.2.3 (git 2.39.1) at 2026-10-03 02:10:45 +0000