From 42deaa082fe026b955f1bc11b5168a2309aad99b Mon Sep 17 00:00:00 2001 From: sigoden Date: Tue, 5 Nov 2024 17:16:01 +0800 Subject: feat: support session-scoped agent variables (#969) --- src/config/agent.rs | 147 +++++++++++++++++++++++++++++++++------------------- 1 file changed, 95 insertions(+), 52 deletions(-) (limited to 'src/config/agent.rs') diff --git a/src/config/agent.rs b/src/config/agent.rs index 39c913a..5d7278b 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -8,12 +8,18 @@ use std::{fs::read_to_string, path::Path}; use serde::{Deserialize, Serialize}; +const DEFAULT_AGENT_NAME: &str = "rag"; + #[derive(Debug, Clone, Serialize)] pub struct Agent { name: String, config: AgentConfig, definition: AgentDefinition, #[serde(skip)] + shared_variables: IndexMap, + #[serde(skip)] + session_variables: Option>, + #[serde(skip)] functions: Functions, #[serde(skip)] rag: Option>, @@ -33,7 +39,7 @@ impl Agent { bail!("Unknown agent `{name}`"); } let functions_file_path = functions_dir.join("functions.json"); - let rag_path = Config::agent_rag_file(name, "rag")?; + let rag_path = Config::agent_rag_file(name, DEFAULT_AGENT_NAME)?; let config_path = Config::agent_config_file(name)?; let agent_config = if config_path.exists() { AgentConfig::load(&config_path)? @@ -41,9 +47,6 @@ impl Agent { AgentConfig::new(&config.read()) }; let mut definition = AgentDefinition::load(&definition_file_path)?; - init_variables(&mut definition.variables, &agent_config.variables) - .context("Failed to init variables")?; - let functions = if functions_file_path.exists() { Functions::init(&functions_file_path)? } else { @@ -60,7 +63,7 @@ impl Agent { }; let rag = if rag_path.exists() { - Some(Arc::new(Rag::load(config, "rag", &rag_path)?)) + Some(Arc::new(Rag::load(config, DEFAULT_AGENT_NAME, &rag_path)?)) } else if !definition.documents.is_empty() { let mut ans = false; if *IS_STDOUT_TERMINAL { @@ -93,16 +96,66 @@ impl Agent { name: name.to_string(), config: agent_config, definition, + shared_variables: Default::default(), + session_variables: None, functions, rag, model, }) } + pub fn init_agent_variables( + agent_variables: &[AgentVariable], + variables: &IndexMap, + ) -> Result> { + let mut output = IndexMap::new(); + if agent_variables.is_empty() { + return Ok(output); + } + let mut printed = false; + for agent_variable in agent_variables { + let key = agent_variable.name.clone(); + match variables.get(&key) { + Some(value) => { + output.insert(key, value.clone()); + } + None => { + if let Some(value) = agent_variable.default.clone() { + output.insert(key, value); + continue; + } + if *IS_STDOUT_TERMINAL { + if !printed { + println!("🚀 Init agent variables..."); + printed = true; + } + let value = Text::new(&agent_variable.description) + .with_validator(|input: &str| { + if input.trim().is_empty() { + Ok(Validation::Invalid("This field is required".into())) + } else { + Ok(Validation::Valid) + } + }) + .prompt()?; + output.insert(key, value); + } else { + bail!("Failed to init agent variables in non-interactive mode"); + } + } + } + } + Ok(output) + } + pub fn export(&self) -> Result { let mut agent = self.clone(); agent.definition.instructions = self.interpolated_instructions(); let mut value = serde_json::json!(agent); + let variables = self.variables(); + if !variables.is_empty() { + value["variables"] = serde_json::to_value(variables)?; + } value["functions_dir"] = Config::agent_functions_dir(&self.name)? .display() .to_string() @@ -140,7 +193,7 @@ impl Agent { } pub fn interpolated_instructions(&self) -> String { - self.definition.interpolated_instructions() + self.definition.interpolated_instructions(self.variables()) } pub fn agent_prelude(&self) -> Option<&str> { @@ -151,18 +204,43 @@ impl Agent { self.config.agent_prelude = value; } - pub fn variables(&self) -> &[AgentVariable] { - &self.definition.variables + pub fn variables(&self) -> &IndexMap { + match &self.session_variables { + Some(variables) => variables, + None => &self.shared_variables, + } + } + + pub fn config_variables(&self) -> &IndexMap { + &self.config.variables + } + + pub fn shared_variables(&self) -> &IndexMap { + &self.shared_variables + } + + pub fn set_shared_variables(&mut self, shared_variables: IndexMap) { + self.shared_variables = shared_variables; + } + + pub fn set_session_variables(&mut self, session_variables: Option>) { + self.session_variables = session_variables; } pub fn set_variable(&mut self, key: &str, value: &str) -> Result<()> { - match self.definition.variables.iter_mut().find(|v| v.name == key) { - Some(variable) => { - variable.value = value.to_string(); - Ok(()) - } - None => bail!("Unknown variable '{key}'"), + let variables = match self.session_variables.as_mut() { + Some(v) => v, + None => &mut self.shared_variables, + }; + if !variables.contains_key(key) { + bail!("Unknown variable: '{key}'") } + variables.insert(key.to_string(), value.to_string()); + Ok(()) + } + + pub fn defined_variables(&self) -> &[AgentVariable] { + &self.definition.variables } } @@ -296,10 +374,10 @@ impl AgentDefinition { ) } - fn interpolated_instructions(&self) -> String { + fn interpolated_instructions(&self, variables: &IndexMap) -> String { let mut output = self.instructions.clone(); - for variable in &self.variables { - output = output.replace(&format!("{{{{{}}}}}", variable.name), &variable.value) + for (k, v) in variables { + output = output.replace(&format!("{{{{{k}}}}}"), v) } interpolate_variables(&mut output); output @@ -356,38 +434,3 @@ fn list_agents_impl() -> Result> { .collect(); Ok(agents) } - -fn init_variables( - variables: &mut [AgentVariable], - config_variable: &IndexMap, -) -> Result<()> { - if variables.is_empty() { - return Ok(()); - } - for variable in variables.iter_mut() { - match config_variable.get(&variable.name) { - Some(value) => variable.value = value.to_string(), - None => { - if let Some(value) = variable.default.clone() { - variable.value = value; - continue; - } - if *IS_STDOUT_TERMINAL { - let value = Text::new(&variable.description) - .with_validator(|input: &str| { - if input.trim().is_empty() { - Ok(Validation::Invalid("This field is required".into())) - } else { - Ok(Validation::Valid) - } - }) - .prompt()?; - variable.value = value; - } else { - bail!("Failed to init agent variables in non-interactive mode"); - } - } - } - } - Ok(()) -} -- cgit v1.2.3