feat: add config `dangerously_functions` (#582)

4 files changed, 40 insertions, 35 deletions

diff --git a/src/config/bot.rs b/src/config/bot.rs
index 19fd5cd..80d3a5d 100644
--- a/src/config/bot.rs
+++ b/src/config/bot.rs
@@ -105,6 +105,10 @@ impl Bot {
         &self.name
     }
 
+    pub fn config(&self) -> &BotConfig {
+        &self.config
+    }
+
     pub fn functions(&self) -> &Functions {
         &self.functions
     }
@@ -170,6 +174,8 @@ pub struct BotConfig {
     pub temperature: Option<f64>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub top_p: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub dangerously_functions: Option<FunctionsFilter>,
 }
 
 impl BotConfig {
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 7c99acb..3af4467 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -12,15 +12,13 @@ use crate::client::{
     create_client_config, list_chat_models, list_client_types, ClientConfig, Model,
     OPENAI_COMPATIBLE_PLATFORMS,
 };
-use crate::function::{FunctionDeclaration, Functions, ToolCallResult};
+use crate::function::{FunctionDeclaration, Functions, FunctionsFilter, ToolCallResult};
 use crate::rag::Rag;
 use crate::render::{MarkdownRender, RenderOptions};
-use crate::utils::{
-    format_option_value, fuzzy_match, get_env_name, light_theme_from_colorfgbg, now, render_prompt,
-    set_text, warning_text, AbortSignal, IS_STDOUT_TERMINAL,
-};
+use crate::utils::*;
 
 use anyhow::{anyhow, bail, Context, Result};
+use fancy_regex::Regex;
 use inquire::{Confirm, Select};
 use parking_lot::RwLock;
 use serde::Deserialize;
@@ -96,6 +94,7 @@ pub struct Config {
     pub rag_top_k: usize,
     pub rag_template: Option<String>,
     pub function_calling: bool,
+    pub dangerously_functions: Option<FunctionsFilter>,
     pub compress_threshold: usize,
     pub summarize_prompt: Option<String>,
     pub summary_prompt: Option<String>,
@@ -144,6 +143,7 @@ impl Default for Config {
             rag_top_k: 4,
             rag_template: None,
             function_calling: false,
+            dangerously_functions: None,
             compress_threshold: 4000,
             summarize_prompt: None,
             summary_prompt: None,
@@ -965,6 +965,26 @@ impl Config {
         functions
     }
 
+    pub fn is_dangerously_function(&self, name: &str) -> bool {
+        if get_env_bool("no_dangerously_functions") {
+            return false;
+        }
+        let dangerously_functions = match &self.bot {
+            Some(bot) => bot.config().dangerously_functions.as_ref(),
+            None => self.dangerously_functions.as_ref(),
+        };
+        match dangerously_functions {
+            None => false,
+            Some(regex) => {
+                let regex = match Regex::new(&format!("^({regex})$")) {
+                    Ok(v) => v,
+                    Err(_) => return false,
+                };
+                regex.is_match(name).unwrap_or(false)
+            }
+        }
+    }
+
     pub fn buffer_editor(&self) -> Option<String> {
         self.buffer_editor
             .clone()
diff --git a/src/function.rs b/src/function.rs
index 0896540..f16fe3f 100644
--- a/src/function.rs
+++ b/src/function.rs
@@ -1,9 +1,6 @@
 use crate::{
     config::{Config, GlobalConfig},
-    utils::{
-        dimmed_text, get_env_bool, indent_text, run_command, run_command_with_output, warning_text,
-        IS_STDOUT_TERMINAL,
-    },
+    utils::*,
 };
 
 use anyhow::{anyhow, bail, Context, Result};
@@ -20,6 +17,7 @@ use std::{
 
 pub const FUNCTION_ALL_MATCHER: &str = ".*";
 pub type ToolResults = (Vec<ToolCallResult>, String);
+pub type FunctionsFilter = String;
 
 pub fn eval_tool_calls(
     config: &GlobalConfig,
@@ -171,6 +169,7 @@ impl ToolCall {
 
     pub fn eval(&self, config: &GlobalConfig) -> Result<Value> {
         let function_name = self.name.clone();
+        let is_dangerously = config.read().is_dangerously_function(&function_name);
         let (call_name, cmd_name, mut cmd_args) = match &config.read().bot {
             Some(bot) => {
                 if !bot.functions().contains(&function_name) {
@@ -219,11 +218,11 @@ impl ToolCall {
         #[cfg(windows)]
         let cmd_name = polyfill_cmd_name(&cmd_name, &bin_dir);
 
-        let output = if self.is_execute() {
+        let output = if is_dangerously {
             if *IS_STDOUT_TERMINAL {
                 println!("{prompt}");
                 let answer = Text::new("[1] Run, [2] Run & Retrieve, [3] Skip:")
-                    .with_default("1")
+                    .with_default("2")
                     .with_validator(|input: &str| match matches!(input, "1" | "2" | "3") {
                         true => Ok(Validation::Valid),
                         false => Ok(Validation::Invalid(
@@ -239,7 +238,7 @@ impl ToolCall {
                         }
                         Value::Null
                     }
-                    "2" => run_and_retrieve(&cmd_name, &cmd_args, envs, &prompt)?,
+                    "2" => run_and_retrieve(&cmd_name, &cmd_args, envs)?,
                     _ => Value::Null,
                 }
             } else {
@@ -248,35 +247,23 @@ impl ToolCall {
             }
         } else {
             println!("{}", dimmed_text(&prompt));
-            run_and_retrieve(&cmd_name, &cmd_args, envs, &prompt)?
+            run_and_retrieve(&cmd_name, &cmd_args, envs)?
         };
 
         Ok(output)
     }
-
-    pub fn is_execute(&self) -> bool {
-        if get_env_bool("function_auto_execute") {
-            false
-        } else {
-            self.name.starts_with("may_") || self.name.contains("__may_")
-        }
-    }
 }
 
 fn run_and_retrieve(
     cmd_name: &str,
     cmd_args: &[String],
     envs: HashMap<String, String>,
-    prompt: &str,
 ) -> Result<Value> {
     let (success, stdout, stderr) = run_command_with_output(cmd_name, cmd_args, Some(envs))?;
 
     if success {
         if !stderr.is_empty() {
-            eprintln!(
-                "{}",
-                warning_text(&format!("{prompt}:\n{}", indent_text(&stderr, 4)))
-            );
+            eprintln!("{}", warning_text(&stderr));
         }
         let value = if !stdout.is_empty() {
             serde_json::from_str(&stdout)
@@ -296,7 +283,7 @@ fn run_and_retrieve(
         } else {
             &stderr
         };
-        bail!("{}", &format!("{prompt}:\n{}", indent_text(err, 4)));
+        bail!("{err}");
     }
 }
 
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index cd00cc5..5c34a54 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -151,14 +151,6 @@ pub fn dimmed_text(input: &str) -> String {
     nu_ansi_term::Style::new().dimmed().paint(input).to_string()
 }
 
-pub fn indent_text(text: &str, spaces: usize) -> String {
-    let indent_size = " ".repeat(spaces);
-    text.lines()
-        .map(|line| format!("{}{}", indent_size, line))
-        .collect::<Vec<String>>()
-        .join("\n")
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;