feat: enhenced flexibility for use tools (#688)

4 files changed, 135 insertions, 60 deletions

diff --git a/src/config/agent.rs b/src/config/agent.rs
index c7ff9ad..a6257ca 100644
--- a/src/config/agent.rs
+++ b/src/config/agent.rs
@@ -1,9 +1,6 @@
 use super::*;
 
-use crate::{
-    client::Model,
-    function::{Functions, FunctionsFilter, SELECTED_ALL_FUNCTIONS},
-};
+use crate::{client::Model, function::Functions};
 
 use anyhow::{Context, Result};
 use std::{fs::read_to_string, path::Path};
@@ -150,11 +147,12 @@ impl RoleLike for Agent {
         self.config.top_p
     }
 
-    fn functions_filter(&self) -> Option<FunctionsFilter> {
-        if self.functions.is_empty() {
+    fn use_tools(&self) -> Option<String> {
+        let common_tools = &self.definition.common_tools;
+        if common_tools.is_empty() {
             None
         } else {
-            Some(SELECTED_ALL_FUNCTIONS.into())
+            Some(common_tools.join(","))
         }
     }
 
@@ -171,7 +169,7 @@ impl RoleLike for Agent {
         self.config.top_p = value;
     }
 
-    fn set_functions_filter(&mut self, _value: Option<FunctionsFilter>) {}
+    fn set_use_tools(&mut self, _value: Option<String>) {}
 }
 
 #[derive(Debug, Clone, Default, Deserialize, Serialize)]
@@ -184,7 +182,7 @@ pub struct AgentConfig {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub top_p: Option<f64>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub dangerously_functions_filter: Option<FunctionsFilter>,
+    pub dangerously_functions_filter: Option<String>,
 }
 
 impl AgentConfig {
@@ -208,6 +206,8 @@ pub struct AgentDefinition {
     pub conversation_starters: Vec<String>,
     #[serde(default)]
     pub documents: Vec<String>,
+    #[serde(default)]
+    pub common_tools: Vec<String>,
 }
 
 impl AgentDefinition {
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 80abbb8..d2804ed 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -12,13 +12,14 @@ use crate::client::{
     create_client_config, list_chat_models, list_client_types, list_reranker_models, ClientConfig,
     Model, OPENAI_COMPATIBLE_PLATFORMS,
 };
-use crate::function::{FunctionDeclaration, Functions, FunctionsFilter, ToolResult};
+use crate::function::{FunctionDeclaration, Functions, ToolResult};
 use crate::rag::Rag;
 use crate::render::{MarkdownRender, RenderOptions};
 use crate::utils::*;
 
 use anyhow::{anyhow, bail, Context, Result};
 use fancy_regex::Regex;
+use indexmap::IndexMap;
 use inquire::{Confirm, Select};
 use parking_lot::RwLock;
 use serde::Deserialize;
@@ -103,7 +104,9 @@ pub struct Config {
     pub summary_prompt: Option<String>,
 
     pub function_calling: bool,
-    pub dangerously_functions_filter: Option<FunctionsFilter>,
+    pub mapping_tools: IndexMap<String, String>,
+    pub use_tools: Option<String>,
+    pub dangerously_functions_filter: Option<String>,
     pub agents: Vec<AgentConfig>,
 
     pub rag_embedding_model: Option<String>,
@@ -164,6 +167,8 @@ impl Default for Config {
             agent_prelude: None,
 
             function_calling: true,
+            mapping_tools: Default::default(),
+            use_tools: None,
             dangerously_functions_filter: None,
             agents: vec![],
 
@@ -421,6 +426,9 @@ impl Config {
         if role.top_p().is_none() && self.top_p.is_some() {
             role.set_top_p(self.top_p);
         }
+        if role.use_tools().is_none() && self.use_tools.is_some() && self.agent.is_none() {
+            role.set_use_tools(self.use_tools.clone())
+        }
         role
     }
 
@@ -475,6 +483,7 @@ impl Config {
             ("save_session", format_option_value(&self.save_session)),
             ("compress_threshold", self.compress_threshold.to_string()),
             ("function_calling", self.function_calling.to_string()),
+            ("use_tools", format_option_value(&role.use_tools())),
             (
                 "rag_reranker_model",
                 format_option_value(&self.rag_reranker_model),
@@ -545,6 +554,13 @@ impl Config {
                 }
                 self.function_calling = value;
             }
+            "use_tools" => {
+                if self.agent.is_some() {
+                    bail!("This action cannot be performed within an agent.")
+                }
+                let value = parse_value(value)?;
+                self.set_use_tools(value);
+            }
             "compress_threshold" => {
                 let value = parse_value(value)?;
                 self.set_compress_threshold(value);
@@ -584,6 +600,13 @@ impl Config {
         }
     }
 
+    pub fn set_use_tools(&mut self, value: Option<String>) {
+        match self.role_like_mut() {
+            Some(role_like) => role_like.set_use_tools(value),
+            None => self.use_tools = value,
+        }
+    }
+
     pub fn set_save_session(&mut self, value: Option<bool>) {
         if let Some(session) = self.session.as_mut() {
             session.set_save_session(value);
@@ -1034,23 +1057,62 @@ impl Config {
     }
 
     pub fn select_functions(&self, model: &Model, role: &Role) -> Option<Vec<FunctionDeclaration>> {
-        let mut functions = None;
+        let mut functions = vec![];
         if self.function_calling {
-            let filter = role.functions_filter();
-            if let Some(filter) = filter {
-                functions = match &self.agent {
-                    Some(agent) => agent.functions().select(&filter),
-                    None => self.functions.select(&filter),
-                };
+            let use_tools = role.use_tools();
+            let declaration_names: HashSet<String> = self
+                .functions
+                .declarations()
+                .iter()
+                .map(|v| v.name.to_string())
+                .collect();
+            if let Some(use_tools) = use_tools {
+                let mut tool_names: HashSet<String> = Default::default();
+                for item in use_tools.split(',') {
+                    let item = item.trim();
+                    if item == "all" {
+                        tool_names.extend(declaration_names);
+                        break;
+                    } else if let Some(values) = self.mapping_tools.get(item) {
+                        tool_names.extend(
+                            values
+                                .split(',')
+                                .map(|v| v.to_string())
+                                .filter(|v| declaration_names.contains(v)),
+                        )
+                    } else if declaration_names.contains(item) {
+                        tool_names.insert(item.to_string());
+                    }
+                }
+                functions = self
+                    .functions
+                    .declarations()
+                    .iter()
+                    .filter_map(|v| {
+                        if tool_names.contains(&v.name) {
+                            Some(v.clone())
+                        } else {
+                            None
+                        }
+                    })
+                    .collect();
+                if let Some(agent) = &self.agent {
+                    let agent_functions = agent.functions().declarations().to_vec();
+                    functions = [agent_functions, functions].concat();
+                }
                 if !model.supports_function_calling() {
-                    functions = None;
+                    functions.clear();
                     if *IS_STDOUT_TERMINAL {
                         eprintln!("{}", warning_text("WARNING: This LLM or client does not support function calling, despite the context requiring it."));
                     }
                 }
             }
         };
-        functions
+        if functions.is_empty() {
+            None
+        } else {
+            Some(functions)
+        }
     }
 
     pub fn is_dangerously_function(&self, name: &str) -> bool {
@@ -1110,14 +1172,15 @@ impl Config {
                     "max_output_tokens",
                     "temperature",
                     "top_p",
-                    "rag_reranker_model",
-                    "rag_top_k",
-                    "function_calling",
-                    "compress_threshold",
+                    "dry_run",
                     "save",
                     "save_session",
+                    "compress_threshold",
+                    "function_calling",
+                    "use_tools",
+                    "rag_reranker_model",
+                    "rag_top_k",
                     "highlight",
-                    "dry_run",
                 ]
                 .into_iter()
                 .map(|v| (format!("{v} "), None))
@@ -1131,8 +1194,7 @@ impl Config {
                     Some(v) => vec![v.to_string()],
                     None => vec![],
                 },
-                "rag_reranker_model" => list_reranker_models(self).iter().map(|v| v.id()).collect(),
-                "function_calling" => complete_bool(self.function_calling),
+                "dry_run" => complete_bool(self.dry_run),
                 "save" => complete_bool(self.save),
                 "save_session" => {
                     let save_session = if let Some(session) = &self.session {
@@ -1142,8 +1204,26 @@ impl Config {
                     };
                     complete_option_bool(save_session)
                 }
+                "function_calling" => complete_bool(self.function_calling),
+                "use_tools" => {
+                    let mut prefix = String::new();
+                    if let Some((v, _)) = args[1].rsplit_once(',') {
+                        prefix = format!("{v},");
+                    }
+                    let mut values = vec![];
+                    if prefix.is_empty() {
+                        values.push("all".to_string());
+                    }
+                    values.extend(self.mapping_tools.keys().map(|v| v.to_string()));
+                    values.extend(self.functions.declarations().iter().map(|v| v.name.clone()));
+                    values
+                        .into_iter()
+                        .filter(|v| !prefix.contains(&format!("{v},")))
+                        .map(|v| format!("{prefix}{v}"))
+                        .collect()
+                }
+                "rag_reranker_model" => list_reranker_models(self).iter().map(|v| v.id()).collect(),
                 "highlight" => complete_bool(self.highlight),
-                "dry_run" => complete_bool(self.dry_run),
                 _ => vec![],
             };
             (values.into_iter().map(|v| (v, None)).collect(), args[1])
diff --git a/src/config/role.rs b/src/config/role.rs
index 29d3f3b..3fc7696 100644
--- a/src/config/role.rs
+++ b/src/config/role.rs
@@ -2,7 +2,6 @@ use super::*;
 
 use crate::{
     client::{Message, MessageContent, MessageRole, Model},
-    function::{FunctionsFilter, SELECTED_ALL_FUNCTIONS},
     utils::{detect_os, detect_shell},
 };
 
@@ -21,11 +20,11 @@ pub trait RoleLike {
     fn model_mut(&mut self) -> &mut Model;
     fn temperature(&self) -> Option<f64>;
     fn top_p(&self) -> Option<f64>;
-    fn functions_filter(&self) -> Option<FunctionsFilter>;
+    fn use_tools(&self) -> Option<String>;
     fn set_model(&mut self, model: &Model);
     fn set_temperature(&mut self, value: Option<f64>);
     fn set_top_p(&mut self, value: Option<f64>);
-    fn set_functions_filter(&mut self, value: Option<FunctionsFilter>);
+    fn set_use_tools(&mut self, value: Option<String>);
 }
 
 #[derive(Debug, Clone, Default, Deserialize, Serialize)]
@@ -43,7 +42,7 @@ pub struct Role {
     #[serde(skip_serializing_if = "Option::is_none")]
     top_p: Option<f64>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    functions_filter: Option<FunctionsFilter>,
+    use_tools: Option<String>,
 
     #[serde(skip)]
     model: Model,
@@ -85,17 +84,13 @@ async function timeout(ms) {
                 .into(),
                 None,
             ),
-            (
-                "%functions%",
-                String::new(),
-                Some(SELECTED_ALL_FUNCTIONS.into()),
-            ),
+            ("%functions%", String::new(), Some("all".into())),
         ]
         .into_iter()
-        .map(|(name, prompt, functions_filter)| Self {
+        .map(|(name, prompt, use_tools)| Self {
             name: name.into(),
             prompt,
-            functions_filter,
+            use_tools,
             ..Default::default()
         })
         .collect()
@@ -111,8 +106,8 @@ async function timeout(ms) {
         let model = role_like.model();
         let temperature = role_like.temperature();
         let top_p = role_like.top_p();
-        let functions_filter = role_like.functions_filter();
-        self.batch_set(model, temperature, top_p, functions_filter);
+        let use_tools = role_like.use_tools();
+        self.batch_set(model, temperature, top_p, use_tools);
     }
 
     pub fn batch_set(
@@ -120,7 +115,7 @@ async function timeout(ms) {
         model: &Model,
         temperature: Option<f64>,
         top_p: Option<f64>,
-        functions_filter: Option<FunctionsFilter>,
+        use_tools: Option<String>,
     ) {
         self.set_model(model);
         if temperature.is_some() {
@@ -129,8 +124,8 @@ async function timeout(ms) {
         if top_p.is_some() {
             self.set_top_p(top_p);
         }
-        if functions_filter.is_some() {
-            self.set_functions_filter(functions_filter);
+        if use_tools.is_some() {
+            self.set_use_tools(use_tools);
         }
     }
 
@@ -242,8 +237,8 @@ impl RoleLike for Role {
         self.top_p
     }
 
-    fn functions_filter(&self) -> Option<FunctionsFilter> {
-        self.functions_filter.clone()
+    fn use_tools(&self) -> Option<String> {
+        self.use_tools.clone()
     }
 
     fn set_model(&mut self, model: &Model) {
@@ -258,8 +253,8 @@ impl RoleLike for Role {
         self.top_p = value;
     }
 
-    fn set_functions_filter(&mut self, value: Option<FunctionsFilter>) {
-        self.functions_filter = value;
+    fn set_use_tools(&mut self, value: Option<String>) {
+        self.use_tools = value;
     }
 }
 
diff --git a/src/config/session.rs b/src/config/session.rs
index 9aa4ca4..80629c1 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -21,7 +21,7 @@ pub struct Session {
     #[serde(skip_serializing_if = "Option::is_none")]
     top_p: Option<f64>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    functions_filter: Option<FunctionsFilter>,
+    use_tools: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     save_session: Option<bool>,
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -134,8 +134,8 @@ impl Session {
         if let Some(top_p) = self.top_p() {
             data["top_p"] = top_p.into();
         }
-        if let Some(functions_filter) = self.functions_filter() {
-            data["functions_filter"] = functions_filter.into();
+        if let Some(use_tools) = self.use_tools() {
+            data["use_tools"] = use_tools.into();
         }
         if let Some(save_session) = self.save_session() {
             data["save_session"] = save_session.into();
@@ -171,8 +171,8 @@ impl Session {
             items.push(("top_p", top_p.to_string()));
         }
 
-        if let Some(functions_filter) = self.functions_filter() {
-            items.push(("functions_filter", functions_filter));
+        if let Some(use_tools) = self.use_tools() {
+            items.push(("use_tools", use_tools));
         }
 
         if let Some(save_session) = self.save_session() {
@@ -242,7 +242,7 @@ impl Session {
         self.model_id = role.model().id();
         self.temperature = role.temperature();
         self.top_p = role.top_p();
-        self.functions_filter = role.functions_filter();
+        self.use_tools = role.use_tools();
         self.model = role.model().clone();
         self.role_name = role.name().to_string();
         self.role_prompt = role.prompt().to_string();
@@ -345,7 +345,7 @@ impl Session {
 
     pub fn guard_empty(&self) -> Result<()> {
         if !self.is_empty() {
-            bail!("Cannot perform this action in a session with messages")
+            bail!("This action cannot be performed in a session with messages.")
         }
         Ok(())
     }
@@ -442,8 +442,8 @@ impl RoleLike for Session {
         self.top_p
     }
 
-    fn functions_filter(&self) -> Option<FunctionsFilter> {
-        self.functions_filter.clone()
+    fn use_tools(&self) -> Option<String> {
+        self.use_tools.clone()
     }
 
     fn set_model(&mut self, model: &Model) {
@@ -468,9 +468,9 @@ impl RoleLike for Session {
         }
     }
 
-    fn set_functions_filter(&mut self, value: Option<FunctionsFilter>) {
-        if self.functions_filter != value {
-            self.functions_filter = value;
+    fn set_use_tools(&mut self, value: Option<String>) {
+        if self.use_tools != value {
+            self.use_tools = value;
             self.dirty = true;
         }
     }