feat: change agent config/variables (#967)

- abandon `.save agent-config` - add variables to agent config.yaml - don't save agent variables after initializing the agent and using `.variable <key> <value>`
sigoden · Nov 4, 2024 · 79973a2 · 79973a2
1 parent 9fa3d8c
commit 79973a2
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 84 deletions.
diff --git a/config.agent.example.yaml b/config.agent.example.yaml
@@ -0,0 +1,10 @@
+# Agent-specific configuration
+
+model: openai:gpt-4o             # Specify the LLM to use
+temperature: null                # Set default temperature parameter
+top_p: null                      # Set default top-p parameter, range (0, 1)
+use_tools: null                  # Which additional tools to use by agent. (e.g. 'fs,web_search')
+agent_prelude: null              # Set a session to use when starting the agent. (e.g. temp, default)
+
+variables:                       # Custom default values for the agent variables
+  <key>: <value>
diff --git a/src/config/agent.rs b/src/config/agent.rs
@@ -4,10 +4,7 @@ use crate::{client::Model, function::Functions};
 
 use anyhow::{Context, Result};
 use inquire::{validator::Validation, Text};
-use std::{
-    fs::{self, read_to_string},
-    path::Path,
-};
+use std::{fs::read_to_string, path::Path};
 
 use serde::{Deserialize, Serialize};
 
@@ -36,7 +33,6 @@ impl Agent {
             bail!("Unknown agent `{name}`");
         }
         let functions_file_path = functions_dir.join("functions.json");
-        let variables_path = Config::agent_variables_file(name)?;
         let rag_path = Config::agent_rag_file(name, "rag")?;
         let config_path = Config::agent_config_file(name)?;
         let agent_config = if config_path.exists() {
@@ -45,7 +41,7 @@ impl Agent {
             AgentConfig::new(&config.read())
         };
         let mut definition = AgentDefinition::load(&definition_file_path)?;
-        init_variables(&variables_path, &mut definition.variables)
+        init_variables(&mut definition.variables, &agent_config.variables)
             .context("Failed to init variables")?;
 
         let functions = if functions_file_path.exists() {
@@ -94,18 +90,6 @@ impl Agent {
         })
     }
 
-    pub fn save_config(&self) -> Result<()> {
-        let config_path = Config::agent_config_file(&self.name)?;
-        ensure_parent_exists(&config_path)?;
-        let content = serde_yaml::to_string(&self.config)?;
-        fs::write(&config_path, content).with_context(|| {
-            format!("Failed to save agent config to '{}'", config_path.display())
-        })?;
-
-        println!("✨ Saved agent config to '{}'.", config_path.display());
-        Ok(())
-    }
-
     pub fn export(&self) -> Result<String> {
         let mut agent = self.clone();
         agent.definition.instructions = self.interpolated_instructions();
@@ -118,7 +102,7 @@ impl Agent {
             .display()
             .to_string()
             .into();
-        value["variables_file"] = Config::agent_variables_file(&self.name)?
+        value["config_file"] = Config::agent_config_file(&self.name)?
             .display()
             .to_string()
             .into();
@@ -166,8 +150,6 @@ impl Agent {
         match self.definition.variables.iter_mut().find(|v| v.name == key) {
             Some(variable) => {
                 variable.value = value.to_string();
-                let variables_path = Config::agent_variables_file(&self.name)?;
-                save_variables(&variables_path, self.variables())?;
                 Ok(())
             }
             None => bail!("Unknown variable '{key}'"),
@@ -229,6 +211,8 @@ pub struct AgentConfig {
     pub top_p: Option<f64>,
     pub use_tools: Option<String>,
     pub agent_prelude: Option<String>,
+    #[serde(default)]
+    pub variables: IndexMap<String, String>,
 }
 
 impl AgentConfig {
@@ -364,65 +348,37 @@ fn list_agents_impl() -> Result<Vec<String>> {
     Ok(agents)
 }
 
-fn init_variables(variables_path: &Path, variables: &mut [AgentVariable]) -> Result<()> {
+fn init_variables(
+    variables: &mut [AgentVariable],
+    config_variable: &IndexMap<String, String>,
+) -> Result<()> {
     if variables.is_empty() {
         return Ok(());
     }
-    let variable_values = if variables_path.exists() {
-        let content = read_to_string(variables_path).with_context(|| {
-            format!(
-                "Failed to read variables from '{}'",
-                variables_path.display()
-            )
-        })?;
-        let variable_values: IndexMap<String, String> = serde_yaml::from_str(&content)?;
-        variable_values
-    } else {
-        Default::default()
-    };
-    let mut initialized = false;
     for variable in variables.iter_mut() {
-        match variable_values.get(&variable.name) {
+        match config_variable.get(&variable.name) {
             Some(value) => variable.value = value.to_string(),
             None => {
-                if !initialized {
-                    println!("The agent has the variables and is initializing them...");
-                    initialized = true;
+                if let Some(value) = variable.default.clone() {
+                    variable.value = value;
+                    continue;
                 }
                 if *IS_STDOUT_TERMINAL {
-                    let mut text =
-                        Text::new(&variable.description).with_validator(|input: &str| {
+                    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)
                             }
-                        });
-                    if let Some(default) = &variable.default {
-                        text = text.with_default(default);
-                    }
-                    let value = text.prompt()?;
+                        })
+                        .prompt()?;
                     variable.value = value;
                 } else {
                     bail!("Failed to init agent variables in the script mode.");
                 }
             }
         }
     }
-    if initialized {
-        save_variables(variables_path, variables)?;
-    }
-    Ok(())
-}
-
-fn save_variables(variables_path: &Path, variables: &[AgentVariable]) -> Result<()> {
-    ensure_parent_exists(variables_path)?;
-    let variable_values: IndexMap<String, String> = variables
-        .iter()
-        .map(|v| (v.name.clone(), v.value.clone()))
-        .collect();
-    let content = serde_yaml::to_string(&variable_values)?;
-    fs::write(variables_path, content)
-        .with_context(|| format!("Failed to save variables to '{}'", variables_path.display()))?;
     Ok(())
 }
diff --git a/src/config/mod.rs b/src/config/mod.rs
@@ -51,7 +51,6 @@ const FUNCTIONS_DIR_NAME: &str = "functions";
 const FUNCTIONS_FILE_NAME: &str = "functions.json";
 const FUNCTIONS_BIN_DIR_NAME: &str = "bin";
 const AGENTS_DIR_NAME: &str = "agents";
-const AGENT_VARIABLES_FILE_NAME: &str = "variables.yaml";
 
 pub const TEMP_ROLE_NAME: &str = "%%";
 pub const TEMP_RAG_NAME: &str = "temp";
@@ -367,10 +366,6 @@ impl Config {
         Ok(Self::agent_data_dir(agent_name)?.join(format!("{rag_name}.yaml")))
     }
 
-    pub fn agent_variables_file(name: &str) -> Result<PathBuf> {
-        Ok(Self::agent_data_dir(name)?.join(AGENT_VARIABLES_FILE_NAME))
-    }
-
     pub fn agents_functions_dir() -> Result<PathBuf> {
         Ok(Self::functions_dir()?.join(AGENTS_DIR_NAME))
     }
@@ -1419,14 +1414,6 @@ impl Config {
         Ok(())
     }
 
-    pub fn save_agent_config(&mut self) -> Result<()> {
-        let agent = match &self.agent {
-            Some(v) => v,
-            None => bail!("No agent"),
-        };
-        agent.save_config()
-    }
-
     pub fn exit_agent(&mut self) -> Result<()> {
         self.exit_session()?;
         if self.agent.take().is_some() {

diff --git a/src/repl/mod.rs b/src/repl/mod.rs
@@ -31,7 +31,7 @@ lazy_static::lazy_static! {
 const MENU_NAME: &str = "completion_menu";
 
 lazy_static::lazy_static! {
-    static ref REPL_COMMANDS: [ReplCommand; 35] = [
+    static ref REPL_COMMANDS: [ReplCommand; 34] = [
         ReplCommand::new(".help", "Show this help message", AssertState::pass()),
         ReplCommand::new(".info", "View system info", AssertState::pass()),
         ReplCommand::new(".model", "Change the current LLM", AssertState::pass()),
@@ -141,11 +141,6 @@ lazy_static::lazy_static! {
             "Set agent variable",
             AssertState::True(StateFlags::AGENT)
         ),
-        ReplCommand::new(
-            ".save agent-config",
-            "Save the current agent config to file",
-            AssertState::True(StateFlags::AGENT)
-        ),
         ReplCommand::new(
             ".info agent",
             "View agent info",
@@ -346,11 +341,8 @@ impl Repl {
                         Some(("session", name)) => {
                             self.config.write().save_session(name)?;
                         }
-                        Some(("agent-config", _)) => {
-                            self.config.write().save_agent_config()?;
-                        }
                         _ => {
-                            println!(r#"Usage: .save <role|session|agent-config> [name]"#)
+                            println!(r#"Usage: .save <role|session> [name]"#)
                         }
                     }
                 }