Skip to content

Commit

Permalink
feat: .agent accept variable key-value pairs (#1090)
Browse files Browse the repository at this point in the history
  • Loading branch information
@sigoden
sigoden authored Jan 17, 2025
1 parent 5f0840a commit d3ae561
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 28 deletions.
13 changes: 6 additions & 7 deletions src/config/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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(())
}
Expand Down Expand Up @@ -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(
Expand All @@ -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(
Expand Down Expand Up @@ -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());
}
Expand Down
6 changes: 4 additions & 2 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,15 +101,17 @@ 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()))
.collect(),
);
}

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)?;
Expand Down
62 changes: 43 additions & 19 deletions src/repl/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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) => {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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(),
Expand All @@ -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()], "")
);
}
Expand Down

0 comments on commit d3ae561

Please sign in to comment.