From 861529374726609a56f6b8af1145a0df5c924dcc Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 11 Jun 2024 14:01:45 +0800
Subject: feat: add config `dangerously_functions` (#582)

---
 src/config/bot.rs |  6 ++++++
 src/config/mod.rs | 30 +++++++++++++++++++++++++-----
 2 files changed, 31 insertions(+), 5 deletions(-)

(limited to 'src/config')

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()
-- 
cgit v1.2.3