Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: .file supports external command execution #1075

Merged
merged 1 commit into from
Jan 6, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion scripts/completions/aichat.fish
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
2 changes: 1 addition & 1 deletion scripts/completions/aichat.nu
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions scripts/completions/aichat.ps1
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down
4 changes: 2 additions & 2 deletions scripts/completions/aichat.zsh
Original file line number Diff line number Diff line change
Expand Up @@ -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]' \
Expand Down
2 changes: 1 addition & 1 deletion src/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
53 changes: 36 additions & 17 deletions src/config/input.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,18 +8,13 @@ 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};

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,
Expand Down Expand Up @@ -64,34 +59,42 @@ 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());
remote_urls.push(path);
}
}
}
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());
};
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);
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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))
Expand All @@ -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()) {
Expand Down
14 changes: 11 additions & 3 deletions src/repl/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()),
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -736,7 +744,7 @@ fn split_files_text(line: &str, is_win: bool) -> (Vec<String>, &str) {
word.clear();
}
}
'\'' | '"' => {
'\'' | '"' | '`' => {
word.push(char);
unbalance = Some(char);
}
Expand Down
Loading