From d3ae5612a51ba6a5d31dfabfe7f638afb8a35f24 Mon Sep 17 00:00:00 2001 From: sigoden Date: Fri, 17 Jan 2025 22:06:57 +0800 Subject: [PATCH] feat: `.agent` accept variable key-value pairs (#1090) --- src/config/mod.rs | 13 +++++----- src/main.rs | 6 +++-- src/repl/mod.rs | 62 ++++++++++++++++++++++++++++++++--------------- 3 files changed, 53 insertions(+), 28 deletions(-) diff --git a/src/config/mod.rs b/src/config/mod.rs index 7e2f275a..6370f3ac 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, + pub agent_variables: Option, #[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 48eb2b9a..a7b51982 100644 --- a/src/main.rs +++ b/src/main.rs @@ -101,7 +101,7 @@ async fn run(config: GlobalConfig, cli: Cli, text: Option) -> 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) -> 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 87befd36..652ecb90 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 [session-name|null] [key=value]..."#) } - None => println!(r#"Usage: .agent [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, &str) { +pub fn split_args_text(line: &str, is_win: bool) -> (Vec, &str) { let mut words = Vec::new(); let mut word = String::new(); let mut unbalance: Option = 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()], "") ); }