summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--scripts/completions/aichat.bash10
-rw-r--r--scripts/completions/aichat.fish10
-rw-r--r--scripts/completions/aichat.nu20
-rw-r--r--scripts/completions/aichat.ps115
-rw-r--r--scripts/completions/aichat.zsh13
-rw-r--r--src/cli.rs16
-rw-r--r--src/config/bot.rs20
-rw-r--r--src/config/mod.rs24
-rw-r--r--src/config/session.rs3
-rw-r--r--src/main.rs80
-rw-r--r--src/rag/mod.rs3
-rw-r--r--src/repl/mod.rs12
-rw-r--r--src/utils/abort_signal.rs4
13 files changed, 158 insertions, 72 deletions
diff --git a/scripts/completions/aichat.bash b/scripts/completions/aichat.bash
index cb9b708..adfe57d 100644
--- a/scripts/completions/aichat.bash
+++ b/scripts/completions/aichat.bash
@@ -19,7 +19,7 @@ _aichat() {
case "${cmd}" in
aichat)
- opts="-m -r -s -e -c -f -H -S -w -h -V --model --prompt --role --session --save-session --serve --execute --code --file --no-highlight --no-stream --wrap --light-theme --dry-run --info --list-models --list-roles --list-sessions --help --version"
+ opts="-m -r -s -b -e -c -f -H -S -w -h -V --model --prompt --role --session --save-session --bot --rag --serve --execute --code --file --no-highlight --no-stream --wrap --light-theme --dry-run --info --list-models --list-roles --list-sessions --list-bots --list-rags --help --version"
if [[ ${cur} == -* || ${COMP_CWORD} -eq 1 ]] ; then
COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") )
return 0
@@ -61,6 +61,14 @@ _aichat() {
COMPREPLY=($(compgen -W "$("$1" --list-sessions)" -- "${cur}"))
return 0
;;
+ -b|--bot)
+ COMPREPLY=($(compgen -W "$("$1" --list-bots)" -- "${cur}"))
+ return 0
+ ;;
+ --rag)
+ COMPREPLY=($(compgen -W "$("$1" --list-rags)" -- "${cur}"))
+ return 0
+ ;;
-f|--file)
local oldifs
if [[ -v IFS ]]; then
diff --git a/scripts/completions/aichat.fish b/scripts/completions/aichat.fish
index 9da7391..d05b4e6 100644
--- a/scripts/completions/aichat.fish
+++ b/scripts/completions/aichat.fish
@@ -2,9 +2,11 @@ complete -c aichat -s m -l model -x -a "(aichat --list-models)" -d 'Select a LLM
complete -c aichat -l prompt -d 'Use the system prompt'
complete -c aichat -s r -l role -x -a "(aichat --list-roles)" -d 'Select a role' -r
complete -c aichat -s s -l session -x -a"(aichat --list-sessions)" -d 'Start or join a session' -r
+complete -c aichat -l save-session -d 'Forces the session to be saved'
+complete -c aichat -s b -l bot -x -a"(aichat --list-bots)" -d 'Start a bot' -r
+complete -c aichat -l rag -x -a"(aichat --list-rags)" -d 'Start a RAG' -r
complete -c aichat -s f -l file -d 'Include files with the message' -r -F
complete -c aichat -s w -l wrap -d 'Control text wrapping (no, auto, <max-width>)'
-complete -c aichat -l save-session -d 'Forces the session to be saved'
complete -c aichat -l serve -d 'Serve the LLM API and WebAPP'
complete -c aichat -s e -l execute -d 'Execute commands in natural language'
complete -c aichat -s c -l code -d 'Output code only'
@@ -14,7 +16,9 @@ complete -c aichat -l light-theme -d 'Use light theme'
complete -c aichat -l dry-run -d 'Display the message without sending it'
complete -c aichat -l info -d 'Display information'
complete -c aichat -l list-models -d 'List all available models'
-complete -c aichat -l list-roles -d 'List all available roles'
-complete -c aichat -l list-sessions -d 'List all available sessions'
+complete -c aichat -l list-roles -d 'List all roles'
+complete -c aichat -l list-sessions -d 'List all sessions'
+complete -c aichat -l list-bots -d 'List all bots'
+complete -c aichat -l list-rags -d 'List all RAGs'
complete -c aichat -s h -l help -d 'Print help'
complete -c aichat -s V -l version -d 'Print version'
diff --git a/scripts/completions/aichat.nu b/scripts/completions/aichat.nu
index 89ebb35..ec4b083 100644
--- a/scripts/completions/aichat.nu
+++ b/scripts/completions/aichat.nu
@@ -22,6 +22,18 @@ module completions {
| parse "{value}"
}
+ def "nu-complete aichat bot" [] {
+ ^aichat --list-bots |
+ | lines
+ | parse "{value}"
+ }
+
+ def "nu-complete aichat rag" [] {
+ ^aichat --list-rags |
+ | lines
+ | parse "{value}"
+ }
+
# All-in-one chat and copilot CLI that integrates 10+ AI platforms
export extern aichat [
--model(-m): string@"nu-complete aichat model" # Select a LLM model
@@ -29,6 +41,8 @@ module completions {
--role(-r): string@"nu-complete aichat role" # Select a role
--session(-s): string@"nu-complete aichat role" # Start or join a session
--save-session # Forces the session to be saved
+ --bot(-b): string@"nu-complete aichat bot" # Start a bot
+ --rag: string@"nu-complete aichat rag" # Start a RAG
--serve # Serve the LLM API and WebAPP
--execute(-e) # Execute commands in natural language
--code(-c) # Output code only
@@ -40,8 +54,10 @@ module completions {
--dry-run # Display the message without sending it
--info # Display information
--list-models # List all available models
- --list-roles # List all available roles
- --list-sessions # List all available sessions
+ --list-roles # List all roles
+ --list-sessions # List all sessions
+ --list-bots # List all bots
+ --list-rags # List all RAGs
...text: string # Input text
--help(-h) # Print help
--version(-V) # Print version
diff --git a/scripts/completions/aichat.ps1 b/scripts/completions/aichat.ps1
index 78f2f02..093f53f 100644
--- a/scripts/completions/aichat.ps1
+++ b/scripts/completions/aichat.ps1
@@ -27,11 +27,14 @@ Register-ArgumentCompleter -Native -CommandName 'aichat' -ScriptBlock {
[CompletionResult]::new('--role', '--role', [CompletionResultType]::ParameterName, 'Select a role')
[CompletionResult]::new('-s', '-s', [CompletionResultType]::ParameterName, 'Start or join a session')
[CompletionResult]::new('--session', '--session', [CompletionResultType]::ParameterName, 'Start or join a session')
+ [CompletionResult]::new('--save-session', '--save-session', [CompletionResultType]::ParameterName, 'Forces the session to be saved')
+ [CompletionResult]::new('-b', '-b', [CompletionResultType]::ParameterName, 'Start a bot')
+ [CompletionResult]::new('--bot', '--bot', [CompletionResultType]::ParameterName, 'Start a bot')
+ [CompletionResult]::new('--rag', '--rag', [CompletionResultType]::ParameterName, 'Start a RAG')
[CompletionResult]::new('-f', '-f', [CompletionResultType]::ParameterName, 'Include files with the message')
[CompletionResult]::new('--file', '--file', [CompletionResultType]::ParameterName, 'Include files with the message')
[CompletionResult]::new('-w', '-w', [CompletionResultType]::ParameterName, 'Control text wrapping (no, auto, <max-width>)')
[CompletionResult]::new('--wrap', '--wrap', [CompletionResultType]::ParameterName, 'Control text wrapping (no, auto, <max-width>)')
- [CompletionResult]::new('--save-session', '--save-session', [CompletionResultType]::ParameterName, 'Forces the session to be saved')
[CompletionResult]::new('--serve', '--serve', [CompletionResultType]::ParameterName, 'Serve the LLM API and WebAPP')
[CompletionResult]::new('-e', '-e', [CompletionResultType]::ParameterName, 'Execute commands in natural language')
[CompletionResult]::new('--execute', '--execute', [CompletionResultType]::ParameterName, 'Execute commands in natural language')
@@ -45,8 +48,10 @@ Register-ArgumentCompleter -Native -CommandName 'aichat' -ScriptBlock {
[CompletionResult]::new('--dry-run', '--dry-run', [CompletionResultType]::ParameterName, 'Display the message without sending it')
[CompletionResult]::new('--info', '--info', [CompletionResultType]::ParameterName, 'Display information')
[CompletionResult]::new('--list-models', '--list-models', [CompletionResultType]::ParameterName, 'List all available models')
- [CompletionResult]::new('--list-roles', '--list-roles', [CompletionResultType]::ParameterName, 'List all available roles')
- [CompletionResult]::new('--list-sessions', '--list-sessions', [CompletionResultType]::ParameterName, 'List all available sessions')
+ [CompletionResult]::new('--list-roles', '--list-roles', [CompletionResultType]::ParameterName, 'List all roles')
+ [CompletionResult]::new('--list-sessions', '--list-sessions', [CompletionResultType]::ParameterName, 'List all sessions')
+ [CompletionResult]::new('--list-bots', '--list-bots', [CompletionResultType]::ParameterName, 'List all bots')
+ [CompletionResult]::new('--list-rags', '--list-rags', [CompletionResultType]::ParameterName, 'List all RAGs')
[CompletionResult]::new('-h', '-h', [CompletionResultType]::ParameterName, 'Print help')
[CompletionResult]::new('--help', '--help', [CompletionResultType]::ParameterName, 'Print help')
[CompletionResult]::new('-V', '-V', [CompletionResultType]::ParameterName, 'Print version')
@@ -71,6 +76,10 @@ Register-ArgumentCompleter -Native -CommandName 'aichat' -ScriptBlock {
$completions = Get-AichatValues "--list-roles"
} elseif ($flag -eq "-s" -or $flag -eq "--session") {
$completions = Get-AichatValues "--list-sessions"
+ } elseif ($flag -eq "-b" -or $flag -eq "--bot") {
+ $completions = Get-AichatValues "--list-bots"
+ } elseif ($flag -eq "--rag") {
+ $completions = Get-AichatValues "--list-rags"
} elseif ($flag -eq "-f" -or $flag -eq "--file") {
$completions = @()
}
diff --git a/scripts/completions/aichat.zsh b/scripts/completions/aichat.zsh
index 043d1db..ae39481 100644
--- a/scripts/completions/aichat.zsh
+++ b/scripts/completions/aichat.zsh
@@ -22,11 +22,14 @@ _aichat() {
'--role=[Select a role]:ROLE:->roles' \
'-s+[Start or join a session]:SESSION:->sessions' \
'--session=[Start or join a session]:SESSION:->sessions' \
+'--save-session[Forces the session to be saved]' \
+'-b+[Start a bot]:BOT:->bots' \
+'--bot=[Start a bot]:BOT:->bots' \
+'--rag=[Start a RAG]:RAG:->rags' \
'*-f+[Include files with the message]:FILE:_files' \
'*--file=[Include files with the message]:FILE:_files' \
'-w+[Control text wrapping (no, auto, <max-width>)]:WRAP: ' \
'--wrap=[Control text wrapping (no, auto, <max-width>)]:WRAP: ' \
-'--save-session[Forces the session to be saved]' \
'--serve[Serve the LLM API and WebAPP]' \
'-e[Execute commands in natural language]' \
'--execute[Execute commands in natural language]' \
@@ -40,8 +43,10 @@ _aichat() {
'--dry-run[Display the message without sending it]' \
'--info[Display information]' \
'--list-models[List all available models]' \
-'--list-roles[List all available roles]' \
-'--list-sessions[List all available sessions]' \
+'--list-roles[List all roles]' \
+'--list-sessions[List all sessions]' \
+'--list-bots[List all bots]' \
+'--list-rags[List all RAGs]' \
'-h[Print help]' \
'--help[Print help]' \
'-V[Print version]' \
@@ -53,7 +58,7 @@ _aichat() {
_arguments "${_arguments_options[@]}" $common \
&& ret=0
case $state in
- models|roles|sessions)
+ models|roles|sessions|bots|rags)
local -a values expl
values=( ${(f)"$(_call_program values aichat --list-$state)"} )
_wanted values expl $state compadd -a values && ret=0
diff --git a/src/cli.rs b/src/cli.rs
index c84db2a..651d23a 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -18,6 +18,12 @@ pub struct Cli {
/// Forces the session to be saved
#[clap(long)]
pub save_session: bool,
+ /// Start a bot
+ #[clap(short = 'b', long)]
+ pub bot: Option<String>,
+ /// Start a RAG
+ #[clap(long)]
+ pub rag: Option<String>,
/// Serve the LLM API and WebAPP
#[clap(long, value_name = "ADDRESS")]
pub serve: Option<Option<String>>,
@@ -51,12 +57,18 @@ pub struct Cli {
/// List all available models
#[clap(long)]
pub list_models: bool,
- /// List all available roles
+ /// List all roles
#[clap(long)]
pub list_roles: bool,
- /// List all available sessions
+ /// List all sessions
#[clap(long)]
pub list_sessions: bool,
+ /// List all bots
+ #[clap(long)]
+ pub list_bots: bool,
+ /// List all RAGs
+ #[clap(long)]
+ pub list_rags: bool,
/// Input text
#[clap(trailing_var_arg = true)]
text: Vec<String>,
diff --git a/src/config/bot.rs b/src/config/bot.rs
index e1e1df6..e02cce7 100644
--- a/src/config/bot.rs
+++ b/src/config/bot.rs
@@ -53,22 +53,14 @@ impl Bot {
None => config.current_model().clone(),
}
};
-
let rag = if rag_path.exists() {
Some(Arc::new(Rag::load(config, "rag", &rag_path)?))
} else if embeddings_dir.is_dir() {
- println!("The bot has an embeddings directory, RAG is initializing...");
- let ans = Confirm::new("The bot attached embeddings, init RAG?")
- .with_default(true)
- .prompt()?;
- if ans {
- let doc_path = embeddings_dir.display().to_string();
- Some(Arc::new(
- Rag::init(config, "rag", &rag_path, &[doc_path], abort_signal).await?,
- ))
- } else {
- None
- }
+ println!("The bot uses an embeddings directory, initializing RAG...");
+ let doc_path = embeddings_dir.display().to_string();
+ Some(Arc::new(
+ Rag::init(config, "rag", &rag_path, &[doc_path], abort_signal).await?,
+ ))
} else {
None
};
@@ -121,7 +113,7 @@ impl Bot {
self.rag.clone()
}
- pub fn converstaion_staters(&self) -> &[String] {
+ pub fn conversation_staters(&self) -> &[String] {
&self.definition.conversation_starters
}
}
diff --git a/src/config/mod.rs b/src/config/mod.rs
index b1eb41f..571bf2c 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -442,7 +442,20 @@ impl Config {
}
pub fn info(&self) -> Result<String> {
- if let Some(session) = &self.session {
+ if let Some(bot) = &self.bot {
+ let output = bot.export()?;
+ if let Some(session) = &self.session {
+ let session = session
+ .export()?
+ .split('\n')
+ .map(|v| format!(" {v}"))
+ .collect::<Vec<_>>()
+ .join("\n");
+ Ok(format!("{output}session:\n{session}"))
+ } else {
+ Ok(output)
+ }
+ } else if let Some(session) = &self.session {
session.export()
} else if let Some(role) = &self.role {
role.export()
@@ -896,6 +909,7 @@ impl Config {
pub async fn use_bot(
config: &GlobalConfig,
name: &str,
+ session: Option<&str>,
abort_signal: AbortSignal,
) -> Result<()> {
if !config.read().function_calling {
@@ -904,11 +918,13 @@ impl Config {
if config.read().bot.is_some() {
bail!("Already in a bot, please run '.exit bot' first to exit the current bot.");
}
- let prelude = config.read().bot_prelude.clone();
let bot = Bot::init(config, name, abort_signal).await?;
config.write().rag = bot.rag();
config.write().bot = Some(bot);
- if let Some(session) = prelude {
+ let session = session
+ .map(|v| v.to_string())
+ .or_else(|| config.read().bot_prelude.clone());
+ if let Some(session) = session {
config.write().use_session(Some(&session))?;
}
Ok(())
@@ -1033,7 +1049,7 @@ impl Config {
".bot" => list_bots().into_iter().map(|v| (v, None)).collect(),
".starter" => match &self.bot {
Some(bot) => bot
- .converstaion_staters()
+ .conversation_staters()
.iter()
.map(|v| (v.clone(), None))
.collect(),
diff --git a/src/config/session.rs b/src/config/session.rs
index 2aff467..38855cc 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -126,9 +126,6 @@ impl Session {
}
pub fn export(&self) -> Result<String> {
- if self.path.is_none() {
- bail!("Not found session '{}'", self.name)
- }
let mut data = json!({
"path": self.path,
"model": self.model().id(),
diff --git a/src/main.rs b/src/main.rs
index 72be91c..3f197eb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,15 +16,13 @@ extern crate log;
use crate::cli::Cli;
use crate::client::{list_chat_models, send_stream, ChatCompletionsOutput};
use crate::config::{
- Config, GlobalConfig, Input, WorkingMode, CODE_ROLE, EXPLAIN_SHELL_ROLE, SHELL_ROLE,
+ list_bots, Config, GlobalConfig, Input, WorkingMode, CODE_ROLE, EXPLAIN_SHELL_ROLE, SHELL_ROLE,
+ TEMP_SESSION_NAME,
};
use crate::function::{eval_tool_calls, need_send_call_results};
use crate::render::{render_error, MarkdownRender};
use crate::repl::Repl;
-use crate::utils::{
- create_abort_signal, detect_shell, extract_block, run_command, run_spinner, Shell,
- CODE_BLOCK_RE, IS_STDOUT_TERMINAL,
-};
+use crate::utils::*;
use anyhow::{bail, Result};
use async_recursion::async_recursion;
@@ -53,9 +51,17 @@ async fn main() -> Result<()> {
crate::logger::setup_logger(working_mode)?;
let config = Arc::new(RwLock::new(Config::init(working_mode)?));
+ let abort_signal = create_abort_signal();
+
if let Some(addr) = cli.serve {
return serve::run(config, addr).await;
}
+ if cli.list_models {
+ for model in list_chat_models(&config.read()) {
+ println!("{}", model.id());
+ }
+ return Ok(());
+ }
if cli.list_roles {
config
.read()
@@ -64,15 +70,14 @@ async fn main() -> Result<()> {
.for_each(|v| println!("{}", v.name()));
return Ok(());
}
- if cli.list_models {
- for model in list_chat_models(&config.read()) {
- println!("{}", model.id());
- }
+ if cli.list_bots {
+ let bots = list_bots().join("\n");
+ println!("{bots}");
return Ok(());
}
- if cli.list_sessions {
- let sessions = config.read().list_sessions().join("\n");
- println!("{sessions}");
+ if cli.list_rags {
+ let rags = config.read().list_rags().join("\n");
+ println!("{rags}");
return Ok(());
}
if let Some(wrap) = &cli.wrap {
@@ -84,19 +89,36 @@ async fn main() -> Result<()> {
if cli.dry_run {
config.write().dry_run = true;
}
- if let Some(prompt) = &cli.prompt {
- config.write().use_prompt(prompt)?;
- } else if let Some(name) = &cli.role {
- config.write().use_role(name)?;
- } else if cli.execute {
- config.write().use_role(SHELL_ROLE)?;
- } else if cli.code {
- config.write().use_role(CODE_ROLE)?;
+
+ if let Some(bot) = &cli.bot {
+ let session = cli.session.as_ref().map(|v| match v {
+ Some(v) => v.as_str(),
+ None => TEMP_SESSION_NAME,
+ });
+ Config::use_bot(&config, bot, session, abort_signal.clone()).await?
+ } else {
+ if let Some(prompt) = &cli.prompt {
+ config.write().use_prompt(prompt)?;
+ } else if let Some(name) = &cli.role {
+ config.write().use_role(name)?;
+ } else if cli.execute {
+ config.write().use_role(SHELL_ROLE)?;
+ } else if cli.code {
+ config.write().use_role(CODE_ROLE)?;
+ }
+ if let Some(session) = &cli.session {
+ config
+ .write()
+ .use_session(session.as_ref().map(|v| v.as_str()))?;
+ }
+ if let Some(rag) = &cli.rag {
+ Config::use_rag(&config, Some(rag), abort_signal.clone()).await?;
+ }
}
- if let Some(session) = &cli.session {
- config
- .write()
- .use_session(session.as_ref().map(|v| v.as_str()))?;
+ if cli.list_sessions {
+ let sessions = config.read().list_sessions().join("\n");
+ println!("{sessions}");
+ return Ok(());
}
if let Some(model_id) = &cli.model {
config.write().set_model(model_id)?;
@@ -124,8 +146,9 @@ async fn main() -> Result<()> {
config.write().apply_prelude()?;
if let Err(err) = match no_input {
false => {
- let input = create_input(&config, text, file)?;
- start_directive(&config, input, cli.no_stream, cli.code).await
+ let mut input = create_input(&config, text, file)?;
+ input.use_embeddings(abort_signal.clone()).await?;
+ start_directive(&config, input, cli.no_stream, cli.code, abort_signal).await
}
true => start_interactive(&config).await,
} {
@@ -142,6 +165,7 @@ async fn start_directive(
mut input: Input,
no_stream: bool,
code_mode: bool,
+ abort_signal: AbortSignal,
) -> Result<()> {
let client = input.create_client()?;
let extract_code = !*IS_STDOUT_TERMINAL && code_mode;
@@ -167,8 +191,7 @@ async fn start_directive(
(text, vec![])
}
} else {
- let abort = create_abort_signal();
- send_stream(&input, client.as_ref(), config, abort).await?
+ send_stream(&input, client.as_ref(), config, abort_signal.clone()).await?
};
config
.write()
@@ -180,6 +203,7 @@ async fn start_directive(
input.merge_tool_call(output, tool_call_results),
no_stream,
code_mode,
+ abort_signal,
)
.await
} else {
diff --git a/src/rag/mod.rs b/src/rag/mod.rs
index f88e80e..c98bab7 100644
--- a/src/rag/mod.rs
+++ b/src/rag/mod.rs
@@ -50,6 +50,9 @@ impl Rag {
doc_paths: &[String],
abort_signal: AbortSignal,
) -> Result<Self> {
+ if !*IS_STDOUT_TERMINAL {
+ bail!("An interactive shell is required to initialize rag.")
+ }
debug!("init rag: {name}");
let model = select_embedding_model(config)?;
let chunk_size = set_chunk_size(&model)?;
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index 0b74e73..d56b790 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -79,7 +79,7 @@ lazy_static! {
),
ReplCommand::new(
".exit session",
- "End the current session",
+ "End the session",
AssertState::True(StateFlags::SESSION_EMPTY | StateFlags::SESSION)
),
ReplCommand::new(
@@ -105,7 +105,7 @@ lazy_static! {
),
ReplCommand::new(
".starter",
- "Use converstaion starters",
+ "Use the conversation starter",
AssertState::True(StateFlags::BOT)
),
ReplCommand::new(
@@ -138,14 +138,13 @@ impl Repl {
let editor = Self::create_editor(config)?;
let prompt = ReplPrompt::new(config);
-
- let abort = create_abort_signal();
+ let abort_signal = create_abort_signal();
Ok(Self {
config: config.clone(),
editor,
prompt,
- abort_signal: abort,
+ abort_signal,
})
}
@@ -254,7 +253,8 @@ impl Repl {
}
".bot" => match args {
Some(name) => {
- Config::use_bot(&self.config, name, self.abort_signal.clone()).await?;
+ Config::use_bot(&self.config, name, None, self.abort_signal.clone())
+ .await?;
}
None => println!(r#"Usage: .bot <name>"#),
},
diff --git a/src/utils/abort_signal.rs b/src/utils/abort_signal.rs
index ac93653..3a8007d 100644
--- a/src/utils/abort_signal.rs
+++ b/src/utils/abort_signal.rs
@@ -54,9 +54,9 @@ impl AbortSignalInner {
}
}
-pub async fn watch_abort_signal(abort: AbortSignal) {
+pub async fn watch_abort_signal(abort_signal: AbortSignal) {
loop {
- if abort.aborted() {
+ if abort_signal.aborted() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(100)).await;