diff options
| author | sigoden <sigoden@gmail.com> | 2025-01-06 08:18:33 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2025-01-06 08:18:33 +0800 |
| commit | ec469cff3ba23b6035aebc74fa7de7b82e678a56 (patch) | |
| tree | c7f730a64246a6fd327f4851f178b659cf1710d9 | |
| parent | bb648d68845e0755edb03570237039433a752463 (diff) | |
| download | aichat-ec469cff3ba23b6035aebc74fa7de7b82e678a56.tar.gz | |
feat: `.file` supports external command execution (#1075)
| -rw-r--r-- | scripts/completions/aichat.fish | 2 | ||||
| -rw-r--r-- | scripts/completions/aichat.nu | 2 | ||||
| -rw-r--r-- | scripts/completions/aichat.ps1 | 4 | ||||
| -rw-r--r-- | scripts/completions/aichat.zsh | 4 | ||||
| -rw-r--r-- | src/cli.rs | 2 | ||||
| -rw-r--r-- | src/config/input.rs | 53 | ||||
| -rw-r--r-- | src/repl/mod.rs | 14 |
7 files changed, 54 insertions, 27 deletions
diff --git a/scripts/completions/aichat.fish b/scripts/completions/aichat.fish index 148af82..19446f8 100644 --- a/scripts/completions/aichat.fish +++ b/scripts/completions/aichat.fish @@ -11,7 +11,7 @@ complete -c aichat -l rebuild-rag -d 'Rebuild the RAG to sync document changes' 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' -complete -c aichat -s f -l file -d 'Include files with the message' -r -F +complete -c aichat -s f -l file -d 'Include files, directories, or URLs' -r -F complete -c aichat -s S -l no-stream -d 'Turn off stream mode' complete -c aichat -l dry-run -d 'Display the message without sending it' complete -c aichat -l info -d 'Display information' diff --git a/scripts/completions/aichat.nu b/scripts/completions/aichat.nu index 748ffc1..ca41518 100644 --- a/scripts/completions/aichat.nu +++ b/scripts/completions/aichat.nu @@ -49,7 +49,7 @@ module completions { --serve # Serve the LLM API and WebAPP --execute(-e) # Execute commands in natural language --code(-c) # Output code only - --file(-f): string # Include files with the message + --file(-f): string # Include files, directories, or URLs --no-stream(-S) # Turn off stream mode --dry-run # Display the message without sending it --info # Display information diff --git a/scripts/completions/aichat.ps1 b/scripts/completions/aichat.ps1 index 30c8c79..8827351 100644 --- a/scripts/completions/aichat.ps1 +++ b/scripts/completions/aichat.ps1 @@ -39,8 +39,8 @@ Register-ArgumentCompleter -Native -CommandName 'aichat' -ScriptBlock { [CompletionResult]::new('--execute', '--execute', [CompletionResultType]::ParameterName, 'Execute commands in natural language') [CompletionResult]::new('-c', '-c', [CompletionResultType]::ParameterName, 'Output code only') [CompletionResult]::new('--code', '--code', [CompletionResultType]::ParameterName, 'Output code only') - [CompletionResult]::new('-f', '-f', [CompletionResultType]::ParameterName, 'Include files with the message') - [CompletionResult]::new('--file', '--file', [CompletionResultType]::ParameterName, 'Include files with the message') + [CompletionResult]::new('-f', '-f', [CompletionResultType]::ParameterName, 'Include files, directories, or URLs') + [CompletionResult]::new('--file', '--file', [CompletionResultType]::ParameterName, 'Include files, directories, or URLs') [CompletionResult]::new('-S', '-S', [CompletionResultType]::ParameterName, 'Turn off stream mode') [CompletionResult]::new('--no-stream', '--no-stream', [CompletionResultType]::ParameterName, 'Turn off stream mode') [CompletionResult]::new('--dry-run', '--dry-run', [CompletionResultType]::ParameterName, 'Display the message without sending it') diff --git a/scripts/completions/aichat.zsh b/scripts/completions/aichat.zsh index 11da4e7..bb6fcda 100644 --- a/scripts/completions/aichat.zsh +++ b/scripts/completions/aichat.zsh @@ -34,8 +34,8 @@ _aichat() { '--execute[Execute commands in natural language]' \ '-c[Output code only]' \ '--code[Output code only]' \ -'*-f[Include files with the message]:FILE:_files' \ -'*--file[Include files with the message]:FILE:_files' \ +'*-f[Include files, directories, or URLs]:FILE:_files' \ +'*--file[Include files, directories, or URLs]:FILE:_files' \ '-S[Turn off stream mode]' \ '--no-stream[Turn off stream mode]' \ '--dry-run[Display the message without sending it]' \ @@ -42,7 +42,7 @@ pub struct Cli { /// Output code only #[clap(short = 'c', long)] pub code: bool, - /// Include files with the message + /// Include files, directories, or URLs #[clap(short = 'f', long, value_name = "FILE")] pub file: Vec<String>, /// Turn off stream mode diff --git a/src/config/input.rs b/src/config/input.rs index 591eb67..2af4c2b 100644 --- a/src/config/input.rs +++ b/src/config/input.rs @@ -8,7 +8,6 @@ use crate::function::ToolResult; use crate::utils::{base64_encode, sha256, AbortSignal}; use anyhow::{bail, Context, Result}; -use fancy_regex::Regex; use path_absolutize::Absolutize; use std::{collections::HashMap, fs::File, io::Read, path::Path}; use unicode_width::{UnicodeWidthChar, UnicodeWidthStr}; @@ -16,10 +15,6 @@ use unicode_width::{UnicodeWidthChar, UnicodeWidthStr}; const IMAGE_EXTS: [&str; 5] = ["png", "jpeg", "jpg", "webp", "gif"]; const SUMMARY_MAX_WIDTH: usize = 80; -lazy_static::lazy_static! { - static ref URL_RE: Regex = Regex::new(r"^[A-Za-z0-9_-]{2,}:/").unwrap(); -} - #[derive(Debug, Clone)] pub struct Input { config: GlobalConfig, @@ -64,15 +59,21 @@ impl Input { role: Option<Role>, ) -> Result<Self> { let mut raw_paths = vec![]; + let mut external_cmds = vec![]; let mut local_paths = vec![]; let mut remote_urls = vec![]; for path in paths { match resolve_local_path(&path) { Some(v) => { - if let Ok(path) = Path::new(&v).absolutize() { - raw_paths.push(path.display().to_string()); + if v.len() > 2 && v.starts_with('`') && v.ends_with('`') { + external_cmds.push(v[1..v.len() - 1].to_string()); + raw_paths.push(v); + } else { + if let Ok(path) = Path::new(&v).absolutize() { + raw_paths.push(path.display().to_string()); + } + local_paths.push(v); } - local_paths.push(v); } None => { raw_paths.push(path.clone()); @@ -80,8 +81,10 @@ impl Input { } } } - let ret = load_documents(config, local_paths, remote_urls).await; - let (files, medias, data_urls) = ret.context("Failed to load files")?; + let (files, medias, data_urls) = + load_documents(config, external_cmds, local_paths, remote_urls) + .await + .context("Failed to load files")?; let mut texts = vec![]; if !raw_text.is_empty() { texts.push(raw_text.to_string()); @@ -89,9 +92,9 @@ impl Input { if !files.is_empty() { texts.push(String::new()); } - for (path, contents) in files { + for (kind, path, contents) in files { texts.push(format!( - "============ PATH: {path} ============\n{contents}\n" + "============ {kind}: {path} ============\n{contents}\n" )); } let (role, with_session, with_agent) = resolve_role(&config.read(), role); @@ -379,14 +382,29 @@ fn resolve_role(config: &Config, role: Option<Role>) -> (Role, bool, bool) { async fn load_documents( config: &GlobalConfig, + external_cmds: Vec<String>, local_paths: Vec<String>, remote_urls: Vec<String>, -) -> Result<(Vec<(String, String)>, Vec<String>, HashMap<String, String>)> { +) -> Result<( + Vec<(&'static str, String, String)>, + Vec<String>, + HashMap<String, String>, +)> { let mut files = vec![]; let mut medias = vec![]; let mut data_urls = HashMap::new(); - let loaders = config.read().document_loaders.clone(); + for cmd in external_cmds { + let (success, stdout, stderr) = + run_command_with_output(&SHELL.cmd, &[&SHELL.arg, &cmd], None)?; + if !success { + let err = if !stderr.is_empty() { stderr } else { stdout }; + bail!("Failed to run `{cmd}`\n{err}"); + } + files.push(("CMD", cmd, stdout)); + } + let local_files = expand_glob_paths(&local_paths, true).await?; + let loaders = config.read().document_loaders.clone(); for file_path in local_files { if is_image(&file_path) { let data_url = read_media_to_data_url(&file_path) @@ -397,9 +415,10 @@ async fn load_documents( let document = load_file(&loaders, &file_path) .await .with_context(|| format!("Unable to read file '{file_path}'"))?; - files.push((file_path, document.contents)); + files.push(("FILE", file_path, document.contents)); } } + for file_url in remote_urls { let (contents, extension) = fetch(&loaders, &file_url, true) .await @@ -408,7 +427,7 @@ async fn load_documents( data_urls.insert(sha256(&contents), file_url); medias.push(contents) } else { - files.push((file_url, contents)); + files.push(("URL", file_url, contents)); } } Ok((files, medias, data_urls)) @@ -427,7 +446,7 @@ pub fn resolve_data_url(data_urls: &HashMap<String, String>, data_url: String) - } fn resolve_local_path(path: &str) -> Option<String> { - if let Ok(true) = URL_RE.is_match(path) { + if is_url(path) { return None; } let new_path = if let (Some(file), Some(home)) = (path.strip_prefix("~/"), dirs::home_dir()) { diff --git a/src/repl/mod.rs b/src/repl/mod.rs index 868325e..0089f9a 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -148,7 +148,7 @@ lazy_static::lazy_static! { ), ReplCommand::new( ".file", - "Include files with the message", + "Include files, directories, URLs or commands", AssertState::pass() ), ReplCommand::new(".continue", "Continue the response", AssertState::pass()), @@ -412,7 +412,15 @@ impl Repl { .await?; ask(&self.config, self.abort_signal.clone(), input, true).await?; } - None => println!("Usage: .file <files>... [-- <text>...]"), + None => println!( + r#"Usage: .file <file|dir|url|cmd>... [-- <text>...] + +.file /tmp/file.txt +.file src/ Cargo.toml -- analyze +.file https://example.com/file.txt -- summarize +.file https://example.com/image.png -- recongize text +.file `git diff` -- Generate git commit message"# + ), }, ".continue" => { let (mut input, output) = match self.config.read().last_message.clone() { @@ -736,7 +744,7 @@ fn split_files_text(line: &str, is_win: bool) -> (Vec<String>, &str) { word.clear(); } } - '\'' | '"' => { + '\'' | '"' | '`' => { word.push(char); unbalance = Some(char); } |
