diff options
Diffstat (limited to 'src/config/mod.rs')
| -rw-r--r-- | src/config/mod.rs | 118 |
1 files changed, 99 insertions, 19 deletions
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]) |