From 2174dc055f660799ed18950fac6a9006f8471fbb Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 21 May 2024 07:56:02 +0800
Subject: refactor: change prompt for executing tool functions (#529)

* refactor: change prompt for executing tool functions

* add AICHAT_FUNCTION_AUTO_EXECUTE to auto confirm
---
 src/function.rs  | 127 +++++++++++++++++++++++++++++++++----------------------
 src/utils/mod.rs |   8 ++++
 2 files changed, 85 insertions(+), 50 deletions(-)

diff --git a/src/function.rs b/src/function.rs
index fe7f088..f2eb0c6 100644
--- a/src/function.rs
+++ b/src/function.rs
@@ -1,12 +1,14 @@
 use crate::{
     config::GlobalConfig,
-    utils::{dimmed_text, indent_text, run_command, run_command_with_output, warning_text},
+    utils::{
+        dimmed_text, get_env_bool, indent_text, run_command, run_command_with_output, warning_text,
+    },
 };
 
 use anyhow::{anyhow, bail, Context, Result};
 use fancy_regex::Regex;
 use indexmap::{IndexMap, IndexSet};
-use inquire::Confirm;
+use inquire::{validator::Validation, Text};
 use is_terminal::IsTerminal;
 use lazy_static::lazy_static;
 use serde::{Deserialize, Serialize};
@@ -38,7 +40,7 @@ pub fn eval_tool_calls(
         return Ok(output);
     }
     calls = ToolCall::dedup(calls);
-    let parallel = calls.len() > 1 && calls.iter().all(|v| !v.is_execute_type());
+    let parallel = calls.len() > 1 && calls.iter().all(|v| !v.is_execute());
     if parallel {
         let (tx, rx) = channel();
         let calls_len = calls.len();
@@ -226,62 +228,87 @@ impl ToolCall {
         if let Some(env_path) = config.read().function.env_path.clone() {
             envs.insert("PATH".into(), env_path);
         };
-        let output = if self.is_execute_type() {
-            let proceed = if stdout().is_terminal() {
-                Confirm::new(&prompt).with_default(true).prompt()?
-            } else {
-                println!("{}", dimmed_text(&prompt));
-                true
-            };
-            if proceed {
-                #[cfg(windows)]
-                let name = polyfill_cmd_name(&name, &config.read().function.bin_dir);
-                let exit_code = run_command(&name, &[&arguments], Some(envs))?;
-                if exit_code != 0 {
-                    bail!("Exit {exit_code}");
+        #[cfg(windows)]
+        let name = polyfill_cmd_name(&name, &config.read().function.bin_dir);
+
+        let output = if self.is_execute() {
+            if stdout().is_terminal() {
+                println!("{prompt}");
+                let anwser = Text::new("[1] Run, [2] Run & Retrieve, [3] Skip:")
+                    .with_default("1")
+                    .with_validator(|input: &str| match matches!(input, "1" | "2" | "3") {
+                        true => Ok(Validation::Valid),
+                        false => Ok(Validation::Invalid(
+                            "Invalid input, please select 1, 2 or 3".into(),
+                        )),
+                    })
+                    .prompt()?;
+                match anwser.as_str() {
+                    "1" => {
+                        let exit_code = run_command(&name, &[arguments], Some(envs))?;
+                        if exit_code != 0 {
+                            bail!("Exit {exit_code}");
+                        }
+                        Value::Null
+                    }
+                    "2" => run_and_retrieve(&name, &arguments, envs, &prompt)?,
+                    _ => Value::Null,
                 }
+            } else {
+                println!("Skipped {prompt}");
+                Value::Null
             }
-            Value::Null
         } else {
             println!("{}", dimmed_text(&prompt));
-            #[cfg(windows)]
-            let name = polyfill_cmd_name(&name, &config.read().function.bin_dir);
-            let (success, stdout, stderr) =
-                run_command_with_output(&name, &[arguments], Some(envs))?;
-
-            if success {
-                if !stderr.is_empty() {
-                    eprintln!(
-                        "{}",
-                        warning_text(&format!("{prompt}:\n{}", indent_text(&stderr, 4)))
-                    );
-                }
-                if !stdout.is_empty() {
-                    serde_json::from_str(&stdout)
-                        .ok()
-                        .unwrap_or_else(|| json!({"output": stdout}))
-                } else {
-                    Value::Null
-                }
-            } else {
-                let err = if stderr.is_empty() {
-                    if stdout.is_empty() {
-                        "Something wrong"
-                    } else {
-                        &stdout
-                    }
-                } else {
-                    &stderr
-                };
-                bail!("{}", &format!("{prompt}:\n{}", indent_text(err, 4)));
-            }
+            run_and_retrieve(&name, &arguments, envs, &prompt)?
         };
 
         Ok(output)
     }
 
-    pub fn is_execute_type(&self) -> bool {
-        self.name.starts_with("may_") || self.name.contains("__may_")
+    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(
+    name: &str,
+    arguments: &str,
+    envs: HashMap<String, String>,
+    prompt: &str,
+) -> Result<Value> {
+    let (success, stdout, stderr) = run_command_with_output(name, &[arguments], Some(envs))?;
+
+    if success {
+        if !stderr.is_empty() {
+            eprintln!(
+                "{}",
+                warning_text(&format!("{prompt}:\n{}", indent_text(&stderr, 4)))
+            );
+        }
+        let value = if !stdout.is_empty() {
+            serde_json::from_str(&stdout)
+                .ok()
+                .unwrap_or_else(|| json!({"output": stdout}))
+        } else {
+            Value::Null
+        };
+        Ok(value)
+    } else {
+        let err = if stderr.is_empty() {
+            if stdout.is_empty() {
+                "Something wrong"
+            } else {
+                &stdout
+            }
+        } else {
+            &stderr
+        };
+        bail!("{}", &format!("{prompt}:\n{}", indent_text(err, 4)));
     }
 }
 
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 3ce2a78..fa67c63 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -35,6 +35,14 @@ pub fn get_env_name(key: &str) -> String {
     )
 }
 
+pub fn get_env_bool(key: &str) -> bool {
+    if let Ok(value) = env::var(get_env_name(key)) {
+        value == "1" || value == "true"
+    } else {
+        false
+    }
+}
+
 pub fn tokenize(text: &str) -> Vec<&str> {
     if text.is_ascii() {
         text.split_inclusive(|c: char| c.is_ascii_whitespace())
-- 
cgit v1.2.3