feat: `.agent` accept variable key-value pairs (#1090)

3 files changed, 53 insertions, 28 deletions

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 7e2f275..6370f3a 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -17,7 +17,7 @@ use crate::client::{
 use crate::function::{FunctionDeclaration, Functions, ToolResult};
 use crate::rag::Rag;
 use crate::render::{MarkdownRender, RenderOptions};
-use crate::repl::{run_repl_command, split_params_text};
+use crate::repl::{run_repl_command, split_args_text};
 use crate::utils::*;
 
 use anyhow::{anyhow, bail, Context, Result};
@@ -148,7 +148,7 @@ pub struct Config {
     #[serde(skip)]
     pub info_flag: bool,
     #[serde(skip)]
-    pub cli_agent_variables: Option<AgentVariables>,
+    pub agent_variables: Option<AgentVariables>,
 
     #[serde(skip)]
     pub model: Model,
@@ -219,7 +219,7 @@ impl Default for Config {
 
             macro_flag: false,
             info_flag: false,
-            cli_agent_variables: None,
+            agent_variables: None,
 
             model: Default::default(),
             functions: Default::default(),
@@ -1530,7 +1530,6 @@ impl Config {
         if self.agent.take().is_some() {
             self.rag.take();
             self.discontinuous_last_message();
-            self.cli_agent_variables = None;
         }
         Ok(())
     }
@@ -2071,7 +2070,7 @@ impl Config {
         };
         if !agent.defined_variables().is_empty() && agent.shared_variables().is_empty() {
             let mut config_variables = agent.config_variables().clone();
-            if let Some(v) = &self.cli_agent_variables {
+            if let Some(v) = &self.agent_variables {
                 config_variables.extend(v.clone());
             }
             let new_variables = Agent::init_agent_variables(
@@ -2097,7 +2096,7 @@ impl Config {
             let session_variables =
                 if !agent.defined_variables().is_empty() && shared_variables.is_empty() {
                     let mut config_variables = agent.config_variables().clone();
-                    if let Some(v) = &self.cli_agent_variables {
+                    if let Some(v) = &self.agent_variables {
                         config_variables.extend(v.clone());
                     }
                     let new_variables = Agent::init_agent_variables(
@@ -2395,7 +2394,7 @@ pub async fn macro_execute(
     abort_signal: AbortSignal,
 ) -> Result<()> {
     let macro_value = Config::load_macro(name)?;
-    let (mut new_args, text) = split_params_text(args.unwrap_or_default(), cfg!(windows));
+    let (mut new_args, text) = split_args_text(args.unwrap_or_default(), cfg!(windows));
     if !text.is_empty() {
         new_args.push(text.to_string());
     }
diff --git a/src/main.rs b/src/main.rs
index 48eb2b9..a7b5198 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -101,7 +101,7 @@ async fn run(config: GlobalConfig, cli: Cli, text: Option<String>) -> Result<()>
             None => TEMP_SESSION_NAME,
         });
         if !cli.agent_variable.is_empty() {
-            config.write().cli_agent_variables = Some(
+            config.write().agent_variables = Some(
                 cli.agent_variable
                     .chunks(2)
                     .map(|v| (v[0].to_string(), v[1].to_string()))
@@ -109,7 +109,9 @@ async fn run(config: GlobalConfig, cli: Cli, text: Option<String>) -> Result<()>
             );
         }
 
-        Config::use_agent(&config, agent, session, abort_signal.clone()).await?
+        let ret = Config::use_agent(&config, agent, session, abort_signal.clone()).await;
+        config.write().agent_variables = None;
+        ret?;
     } else {
         if let Some(prompt) = &cli.prompt {
             config.write().use_prompt(prompt)?;
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index 87befd3..652ecb9 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -8,7 +8,8 @@ use self::prompt::ReplPrompt;
 
 use crate::client::{call_chat_completions, call_chat_completions_streaming};
 use crate::config::{
-    macro_execute, AssertState, Config, GlobalConfig, Input, LastMessage, StateFlags,
+    macro_execute, AgentVariables, AssertState, Config, GlobalConfig, Input, LastMessage,
+    StateFlags,
 };
 use crate::render::render_error;
 use crate::utils::{
@@ -424,11 +425,34 @@ pub async fn run_repl_command(
                 Config::use_rag(config, args, abort_signal.clone()).await?;
             }
             ".agent" => match split_first_arg(args) {
-                Some((agent_name, session_name)) => {
-                    Config::use_agent(config, agent_name, session_name, abort_signal.clone())
-                        .await?;
+                Some((agent_name, args)) => {
+                    let (new_args, _) = split_args_text(args.unwrap_or_default(), cfg!(windows));
+                    let mut session_name = None;
+                    if let Some((name, variable_pairs)) = new_args.split_first() {
+                        if name != "null" && name != "none" {
+                            session_name = Some(name.as_str());
+                        }
+                        let variables: AgentVariables = variable_pairs
+                            .iter()
+                            .filter_map(|v| v.split_once('='))
+                            .map(|(key, value)| (key.to_string(), value.to_string()))
+                            .collect();
+                        if variables.len() != variable_pairs.len() {
+                            bail!("Some variable key-value pairs are invalid");
+                        }
+                        if !variables.is_empty() {
+                            config.write().agent_variables = Some(variables);
+                        }
+                    }
+                    let ret =
+                        Config::use_agent(config, agent_name, session_name, abort_signal.clone())
+                            .await;
+                    config.write().agent_variables = None;
+                    ret?;
+                }
+                None => {
+                    println!(r#"Usage: .agent <agent-name> [session-name|null] [key=value]..."#)
                 }
-                None => println!(r#"Usage: .agent <agent-name> [session-name]"#),
             },
             ".starter" => match args {
                 Some(value) => {
@@ -529,7 +553,7 @@ pub async fn run_repl_command(
             },
             ".file" => match args {
                 Some(args) => {
-                    let (files, text) = split_params_text(args, cfg!(windows));
+                    let (files, text) = split_args_text(args, cfg!(windows));
                     let input = Input::from_files_with_spinner(
                         config,
                         text,
@@ -733,7 +757,7 @@ fn split_first_arg(args: Option<&str>) -> Option<(&str, Option<&str>)> {
     })
 }
 
-pub fn split_params_text(line: &str, is_win: bool) -> (Vec<String>, &str) {
+pub fn split_args_text(line: &str, is_win: bool) -> (Vec<String>, &str) {
     let mut words = Vec::new();
     let mut word = String::new();
     let mut unbalance: Option<char> = None;
@@ -829,45 +853,45 @@ mod tests {
     }
 
     #[test]
-    fn test_split_params_text() {
-        assert_eq!(split_params_text("", false), (vec![], ""));
+    fn test_split_args_text() {
+        assert_eq!(split_args_text("", false), (vec![], ""));
         assert_eq!(
-            split_params_text("file.txt", false),
+            split_args_text("file.txt", false),
             (vec!["file.txt".into()], "")
         );
         assert_eq!(
-            split_params_text("file.txt --", false),
+            split_args_text("file.txt --", false),
             (vec!["file.txt".into()], "")
         );
         assert_eq!(
-            split_params_text("file.txt -- hello", false),
+            split_args_text("file.txt -- hello", false),
             (vec!["file.txt".into()], "hello")
         );
         assert_eq!(
-            split_params_text("file.txt -- \thello", false),
+            split_args_text("file.txt -- \thello", false),
             (vec!["file.txt".into()], "\thello")
         );
         assert_eq!(
-            split_params_text("file.txt --\nhello", false),
+            split_args_text("file.txt --\nhello", false),
             (vec!["file.txt".into()], "hello")
         );
         assert_eq!(
-            split_params_text("file.txt --\r\nhello", false),
+            split_args_text("file.txt --\r\nhello", false),
             (vec!["file.txt".into()], "hello")
         );
         assert_eq!(
-            split_params_text("file.txt --\rhello", false),
+            split_args_text("file.txt --\rhello", false),
             (vec!["file.txt".into()], "hello")
         );
         assert_eq!(
-            split_params_text(r#"file1.txt 'file2.txt' "file3.txt""#, false),
+            split_args_text(r#"file1.txt 'file2.txt' "file3.txt""#, false),
             (
                 vec!["file1.txt".into(), "file2.txt".into(), "file3.txt".into()],
                 ""
             )
         );
         assert_eq!(
-            split_params_text(r#"./file1.txt 'file1 - Copy.txt' file\ 2.txt"#, false),
+            split_args_text(r#"./file1.txt 'file1 - Copy.txt' file\ 2.txt"#, false),
             (
                 vec![
                     "./file1.txt".into(),
@@ -878,7 +902,7 @@ mod tests {
             )
         );
         assert_eq!(
-            split_params_text(r#".\file.txt C:\dir\file.txt"#, true),
+            split_args_text(r#".\file.txt C:\dir\file.txt"#, true),
             (vec![".\\file.txt".into(), "C:\\dir\\file.txt".into()], "")
         );
     }