diff options
| -rw-r--r-- | src/config/agent.rs | 39 | ||||
| -rw-r--r-- | src/config/mod.rs | 12 | ||||
| -rw-r--r-- | src/config/session.rs | 4 | ||||
| -rw-r--r-- | src/function.rs | 112 |
4 files changed, 96 insertions, 71 deletions
diff --git a/src/config/agent.rs b/src/config/agent.rs index 49d5835..29121ee 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -10,15 +10,17 @@ use serde::{Deserialize, Serialize}; const DEFAULT_AGENT_NAME: &str = "rag"; +pub type AgentVariables = IndexMap<String, String>; + #[derive(Debug, Clone, Serialize)] pub struct Agent { name: String, config: AgentConfig, definition: AgentDefinition, #[serde(skip)] - shared_variables: IndexMap<String, String>, + shared_variables: AgentVariables, #[serde(skip)] - session_variables: Option<IndexMap<String, String>>, + session_variables: Option<AgentVariables>, #[serde(skip)] functions: Functions, #[serde(skip)] @@ -108,9 +110,9 @@ impl Agent { pub fn init_agent_variables( agent_variables: &[AgentVariable], - variables: &IndexMap<String, String>, + variables: &AgentVariables, no_interaction: bool, - ) -> Result<IndexMap<String, String>> { + ) -> Result<AgentVariables> { let mut output = IndexMap::new(); if agent_variables.is_empty() { return Ok(output); @@ -224,26 +226,41 @@ impl Agent { self.config.agent_prelude = value; } - pub fn variables(&self) -> &IndexMap<String, String> { + pub fn variables(&self) -> &AgentVariables { match &self.session_variables { Some(variables) => variables, None => &self.shared_variables, } } - pub fn config_variables(&self) -> &IndexMap<String, String> { + pub fn variable_envs(&self) -> HashMap<String, String> { + self.variables() + .iter() + .map(|(k, v)| { + ( + format!("LLM_AGENT_VAR_{}", normalize_env_name(k)), + v.clone(), + ) + }) + .collect() + } + + pub fn config_variables(&self) -> &AgentVariables { &self.config.variables } - pub fn shared_variables(&self) -> &IndexMap<String, String> { + pub fn shared_variables(&self) -> &AgentVariables { &self.shared_variables } - pub fn set_shared_variables(&mut self, shared_variables: IndexMap<String, String>) { + pub fn set_shared_variables(&mut self, shared_variables: AgentVariables) { self.shared_variables = shared_variables; } - pub fn set_session_variables(&mut self, session_variables: Option<IndexMap<String, String>>) { + pub fn set_session_variables(&mut self, session_variables: Option<AgentVariables>) { + if self.shared_variables.is_empty() { + self.shared_variables = session_variables.clone().unwrap_or_default(); + } self.session_variables = session_variables; } @@ -319,7 +336,7 @@ pub struct AgentConfig { pub use_tools: Option<String>, pub agent_prelude: Option<String>, #[serde(default)] - pub variables: IndexMap<String, String>, + pub variables: AgentVariables, } impl AgentConfig { @@ -419,7 +436,7 @@ impl AgentDefinition { ) } - fn interpolated_instructions(&self, variables: &IndexMap<String, String>) -> String { + fn interpolated_instructions(&self, variables: &AgentVariables) -> String { let mut output = self.instructions.clone(); for (k, v) in variables { output = output.replace(&format!("{{{{{k}}}}}"), v) diff --git a/src/config/mod.rs b/src/config/mod.rs index e578620..31337cd 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -3,7 +3,7 @@ mod input; mod role; mod session; -pub use self::agent::{list_agents, Agent}; +pub use self::agent::{list_agents, Agent, AgentVariables}; pub use self::input::Input; pub use self::role::{ Role, RoleLike, CODE_ROLE, CREATE_TITLE_ROLE, EXPLAIN_SHELL_ROLE, SHELL_ROLE, @@ -1487,10 +1487,13 @@ impl Config { if parts.len() != 2 { bail!("Usage: .variable <key> <value>"); } - let key = parts[0]; - let value = parts[1]; match self.agent.as_mut() { Some(agent) => { + if let Some(session) = self.session.as_ref() { + session.guard_empty()?; + } + let key = parts[0]; + let value = parts[1]; agent.set_variable(key, value)?; if let Some(session) = self.session.as_mut() { session.sync_agent(agent, true); @@ -2009,9 +2012,6 @@ impl Config { &all_variables, self.print_info_only, )?; - if shared_variables.is_empty() { - agent.set_shared_variables(new_variables.clone()); - } agent.set_session_variables(Some(new_variables)); session.sync_agent(agent, false); Ok(()) diff --git a/src/config/session.rs b/src/config/session.rs index 23658fd..4b3da76 100644 --- a/src/config/session.rs +++ b/src/config/session.rs @@ -35,7 +35,7 @@ pub struct Session { #[serde(skip_serializing_if = "Option::is_none")] role_name: Option<String>, #[serde(default, skip_serializing_if = "IndexMap::is_empty")] - agent_variables: IndexMap<String, String>, + agent_variables: AgentVariables, #[serde(default, skip_serializing_if = "Vec::is_empty")] compressed_messages: Vec<Message>, @@ -284,7 +284,7 @@ impl Session { } } - pub fn agent_variables(&self) -> &IndexMap<String, String> { + pub fn agent_variables(&self) -> &AgentVariables { &self.agent_variables } diff --git a/src/function.rs b/src/function.rs index 8d4a2a1..698c137 100644 --- a/src/function.rs +++ b/src/function.rs @@ -159,25 +159,17 @@ impl ToolCall { pub fn eval(&self, config: &GlobalConfig) -> Result<Value> { let function_name = self.name.clone(); - let (call_name, cmd_name, mut cmd_args, mut envs) = match &config.read().agent { + let (call_name, cmd_name, mut cmd_args, envs, agent_name) = match &config.read().agent { Some(agent) => match agent.functions().find(&function_name) { Some(function) => { + let agent_name = agent.name().to_string(); if function.agent { - let envs: HashMap<String, String> = agent - .variables() - .iter() - .map(|(k, v)| { - ( - format!("LLM_AGENT_VAR_{}", normalize_env_name(k)), - v.clone(), - ) - }) - .collect(); ( - format!("{}:{}", agent.name(), function_name), - agent.name().to_string(), + format!("{agent_name}-{function_name}"), + agent_name.clone(), vec![function_name], - envs, + agent.variable_envs(), + Some(agent_name), ) } else { ( @@ -185,10 +177,11 @@ impl ToolCall { function_name, vec![], Default::default(), + Some(agent_name), ) } } - None => bail!("Unexpected call {function_name} {}", self.arguments), + None => bail!("Unexpected call: {function_name} {}", self.arguments), }, None => match config.read().functions.contains(&function_name) { true => ( @@ -196,6 +189,7 @@ impl ToolCall { function_name, vec![], Default::default(), + None, ), false => bail!("Unexpected call: {function_name} {}", self.arguments), }, @@ -215,50 +209,64 @@ impl ToolCall { }; cmd_args.push(json_data.to_string()); - let prompt = format!("Call {cmd_name} {}", cmd_args.join(" ")); - let mut bin_dirs: Vec<PathBuf> = vec![]; - if let Some(agent) = config.read().agent.as_ref() { - let dir = Config::agent_functions_dir(agent.name()).join("bin"); - if dir.exists() { - bin_dirs.push(dir); - } - } - bin_dirs.push(Config::functions_bin_dir()); - let current_path = std::env::var("PATH").context("No PATH environment variable")?; - let prepend_path = bin_dirs - .iter() - .map(|v| format!("{}{PATH_SEP}", v.display())) - .collect::<Vec<_>>() - .join(""); - envs.insert("PATH".into(), format!("{prepend_path}{current_path}")); + let output = match run_llm_function(cmd_name, cmd_args, envs, agent_name)? { + Some(contents) => serde_json::from_str(&contents) + .ok() + .unwrap_or_else(|| json!({"result": contents})), + None => Value::Null, + }; - let temp_file = temp_file("-eval-", ""); - envs.insert("LLM_OUTPUT".into(), temp_file.display().to_string()); + Ok(output) + } +} - #[cfg(windows)] - let cmd_name = polyfill_cmd_name(&cmd_name, &bin_dirs); - if *IS_STDOUT_TERMINAL { - println!("{}", dimmed_text(&prompt)); - } - let exit_code = run_command(&cmd_name, &cmd_args, Some(envs)) - .map_err(|err| anyhow!("Unable to run {cmd_name}, {err}"))?; - if exit_code != 0 { - bail!("Tool call exit with {exit_code}"); +fn run_llm_function( + cmd_name: String, + cmd_args: Vec<String>, + mut envs: HashMap<String, String>, + agent_name: Option<String>, +) -> Result<Option<String>> { + let prompt = format!("Call {cmd_name} {}", cmd_args.join(" ")); + + let mut bin_dirs: Vec<PathBuf> = vec![]; + if let Some(agent_name) = agent_name { + let dir = Config::agent_functions_dir(&agent_name).join("bin"); + if dir.exists() { + bin_dirs.push(dir); } - let output = if temp_file.exists() { - let contents = - fs::read_to_string(temp_file).context("Failed to retrieve tool call output")?; + } + bin_dirs.push(Config::functions_bin_dir()); + let current_path = std::env::var("PATH").context("No PATH environment variable")?; + let prepend_path = bin_dirs + .iter() + .map(|v| format!("{}{PATH_SEP}", v.display())) + .collect::<Vec<_>>() + .join(""); + envs.insert("PATH".into(), format!("{prepend_path}{current_path}")); - serde_json::from_str(&contents) - .ok() - .unwrap_or_else(|| json!({"result": contents})) - } else { - Value::Null - }; + let temp_file = temp_file("-eval-", ""); + envs.insert("LLM_OUTPUT".into(), temp_file.display().to_string()); - Ok(output) + #[cfg(windows)] + let cmd_name = polyfill_cmd_name(&cmd_name, &bin_dirs); + if *IS_STDOUT_TERMINAL { + println!("{}", dimmed_text(&prompt)); + } + let exit_code = run_command(&cmd_name, &cmd_args, Some(envs)) + .map_err(|err| anyhow!("Unable to run {cmd_name}, {err}"))?; + if exit_code != 0 { + bail!("Tool call exit with {exit_code}"); } + let mut output = None; + if temp_file.exists() { + let contents = + fs::read_to_string(temp_file).context("Failed to retrieve tool call output")?; + if !contents.is_empty() { + output = Some(contents); + } + }; + Ok(output) } #[cfg(windows)] |