diff options
| -rw-r--r-- | scripts/completions/aichat.bash | 10 | ||||
| -rw-r--r-- | scripts/completions/aichat.fish | 10 | ||||
| -rw-r--r-- | scripts/completions/aichat.nu | 20 | ||||
| -rw-r--r-- | scripts/completions/aichat.ps1 | 15 | ||||
| -rw-r--r-- | scripts/completions/aichat.zsh | 13 | ||||
| -rw-r--r-- | src/cli.rs | 16 | ||||
| -rw-r--r-- | src/config/bot.rs | 20 | ||||
| -rw-r--r-- | src/config/mod.rs | 24 | ||||
| -rw-r--r-- | src/config/session.rs | 3 | ||||
| -rw-r--r-- | src/main.rs | 80 | ||||
| -rw-r--r-- | src/rag/mod.rs | 3 | ||||
| -rw-r--r-- | src/repl/mod.rs | 12 | ||||
| -rw-r--r-- | src/utils/abort_signal.rs | 4 |
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 @@ -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; |
