From 6b6413b02df588faf6557053751b7e9c591539f7 Mon Sep 17 00:00:00 2001 From: sigoden Date: Sun, 1 Dec 2024 11:35:12 +0800 Subject: [PATCH] feat: add cli option `--agent-variable ` (#1027) --- scripts/completions/aichat.bash | 2 +- scripts/completions/aichat.fish | 7 ++++--- scripts/completions/aichat.nu | 3 ++- scripts/completions/aichat.ps1 | 2 +- scripts/completions/aichat.zsh | 2 +- src/cli.rs | 5 ++++- src/config/agent.rs | 16 ++++++--------- src/config/mod.rs | 36 ++++++++++++++++++++++++--------- src/main.rs | 11 +++++++++- 9 files changed, 56 insertions(+), 28 deletions(-) diff --git a/scripts/completions/aichat.bash b/scripts/completions/aichat.bash index 260a9d42..cfa7ec81 100644 --- a/scripts/completions/aichat.bash +++ b/scripts/completions/aichat.bash @@ -17,7 +17,7 @@ _aichat() { case "${cmd}" in aichat) - opts="-m -r -s -a -R -e -c -f -S -h -V --model --prompt --role --session --empty-session --save-session --agent --rag --serve --execute --code --file --no-stream --dry-run --info --list-models --list-roles --list-sessions --list-agents --list-rags --help --version" + opts="-m -r -s -a -e -c -f -S -h -V --model --prompt --role --session --empty-session --save-session --agent --agent-variable --rag --serve --execute --code --file --no-stream --dry-run --info --list-models --list-roles --list-sessions --list-agents --list-rags --help --version" if [[ ${cur} == -* || ${cword} -eq 1 ]] ; then COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") ) return 0 diff --git a/scripts/completions/aichat.fish b/scripts/completions/aichat.fish index df902b3a..19bf96e3 100644 --- a/scripts/completions/aichat.fish +++ b/scripts/completions/aichat.fish @@ -1,11 +1,12 @@ complete -c aichat -s m -l model -x -a "(aichat --list-models)" -d 'Select a LLM model' -r complete -c aichat -l prompt -d 'Use the system prompt' complete -c aichat -s r -l role -x -a "(aichat --list-roles)" -d 'Select a role' -r -complete -c aichat -s s -l session -x -a"(aichat --list-sessions)" -d 'Start or join a session' -r +complete -c aichat -s s -l session -x -a "(aichat --list-sessions)" -d 'Start or join a session' -r complete -c aichat -l empty-session -d 'Ensure the session is empty' complete -c aichat -l save-session -d 'Ensure the new conversation is saved to the session' -complete -c aichat -s a -l agent -x -a"(aichat --list-agents)" -d 'Start a agent' -r -complete -c aichat -s R -l rag -x -a"(aichat --list-rags)" -d 'Start a RAG' -r +complete -c aichat -s a -l agent -x -a "(aichat --list-agents)" -d 'Start a agent' -r +complete -c aichat -l agent-variable -d 'Set agent variables' +complete -c aichat -l rag -x -a"(aichat --list-rags)" -d 'Start a RAG' -r complete -c aichat -l serve -d 'Serve the LLM API and WebAPP' complete -c aichat -s e -l execute -d 'Execute commands in natural language' complete -c aichat -s c -l code -d 'Output code only' diff --git a/scripts/completions/aichat.nu b/scripts/completions/aichat.nu index 382ee2db..447da497 100644 --- a/scripts/completions/aichat.nu +++ b/scripts/completions/aichat.nu @@ -43,7 +43,8 @@ module completions { --empty-session # Ensure the session is empty --save-session # Ensure the new conversation is saved to the session --agent(-a): string@"nu-complete aichat agent" # Start a agent - --rag(-R): string@"nu-complete aichat rag" # Start a RAG + --agent-variable # Set agent variables + --rag: string@"nu-complete aichat rag" # Start a RAG --serve # Serve the LLM API and WebAPP --execute(-e) # Execute commands in natural language --code(-c) # Output code only diff --git a/scripts/completions/aichat.ps1 b/scripts/completions/aichat.ps1 index e9a6259c..7b3cc82b 100644 --- a/scripts/completions/aichat.ps1 +++ b/scripts/completions/aichat.ps1 @@ -31,7 +31,7 @@ Register-ArgumentCompleter -Native -CommandName 'aichat' -ScriptBlock { [CompletionResult]::new('--save-session', '--save-session', [CompletionResultType]::ParameterName, 'Ensure the new conversation is saved to the session') [CompletionResult]::new('-a', '-a', [CompletionResultType]::ParameterName, 'Start a agent') [CompletionResult]::new('--agent', '--agent', [CompletionResultType]::ParameterName, 'Start a agent') - [CompletionResult]::new('-R', '-R', [CompletionResultType]::ParameterName, 'Start a RAG') + [CompletionResult]::new('--agent-variable', '--agent-variable', [CompletionResultType]::ParameterName, 'Set agent variables') [CompletionResult]::new('--rag', '--rag', [CompletionResultType]::ParameterName, 'Start a RAG') [CompletionResult]::new('--serve', '--serve', [CompletionResultType]::ParameterName, 'Serve the LLM API and WebAPP') [CompletionResult]::new('-e', '-e', [CompletionResultType]::ParameterName, 'Execute commands in natural language') diff --git a/scripts/completions/aichat.zsh b/scripts/completions/aichat.zsh index a8b1d0c6..ed7627e1 100644 --- a/scripts/completions/aichat.zsh +++ b/scripts/completions/aichat.zsh @@ -26,7 +26,7 @@ _aichat() { '--save-session[Ensure the new conversation is saved to the session]' \ '-a[Start a agent]:AGENT:->agents' \ '--agent[Start a agent]:AGENT:->agents' \ -'-R[Start a RAG]:RAG:->rags' \ +'--agent-variable[Set agent variables]' \ '--rag[Start a RAG]:RAG:->rags' \ '--serve[Serve the LLM API and WebAPP]' \ '-e[Execute commands in natural language]' \ diff --git a/src/cli.rs b/src/cli.rs index 21734b8a..c3af5cc2 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -24,8 +24,11 @@ pub struct Cli { /// Start a agent #[clap(short = 'a', long)] pub agent: Option, + /// Set agent variables + #[clap(long, value_names = ["NAME", "VALUE"], num_args = 2)] + pub agent_variable: Vec, /// Start a RAG - #[clap(short = 'R', long)] + #[clap(long)] pub rag: Option, /// Serve the LLM API and WebAPP #[clap(long, value_name = "ADDRESS")] diff --git a/src/config/agent.rs b/src/config/agent.rs index 87ac719e..5e54b269 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -15,24 +15,17 @@ const DEFAULT_AGENT_NAME: &str = "rag"; pub type AgentVariables = IndexMap; -#[derive(Debug, Clone, Serialize)] +#[derive(Debug, Clone)] pub struct Agent { name: String, config: AgentConfig, definition: AgentDefinition, - #[serde(skip)] shared_variables: AgentVariables, - #[serde(skip)] session_variables: Option, - #[serde(skip)] shared_dynamic_instructions: Option, - #[serde(skip)] session_dynamic_instructions: Option, - #[serde(skip)] functions: Functions, - #[serde(skip)] rag: Option>, - #[serde(skip)] model: Model, } @@ -75,7 +68,7 @@ impl Agent { let rag = if rag_path.exists() { Some(Arc::new(Rag::load(config, DEFAULT_AGENT_NAME, &rag_path)?)) - } else if !definition.documents.is_empty() && !config.read().print_info_only { + } else if !definition.documents.is_empty() && !config.read().cli_info_flag { let mut ans = false; if *IS_STDOUT_TERMINAL { ans = Confirm::new("The agent has the documents, init RAG?") @@ -182,11 +175,14 @@ impl Agent { pub fn export(&self) -> Result { let mut agent = self.clone(); agent.definition.instructions = self.interpolated_instructions(); - let mut value = serde_json::json!(agent); + let mut value = json!({}); + value["name"] = json!(self.name()); let variables = self.variables(); if !variables.is_empty() { value["variables"] = serde_json::to_value(variables)?; } + value["config"] = json!(self.config); + value["definition"] = json!(self.definition); value["functions_dir"] = Config::agent_functions_dir(&self.name) .display() .to_string() diff --git a/src/config/mod.rs b/src/config/mod.rs index e0e1efa8..c98c81d3 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -156,9 +156,12 @@ pub struct Config { #[serde(skip)] pub working_mode: WorkingMode, #[serde(skip)] - pub print_info_only: bool, - #[serde(skip)] pub last_message: Option<(Input, String)>, + + #[serde(skip)] + pub cli_info_flag: bool, + #[serde(skip)] + pub cli_agent_variables: Option, } impl Default for Config { @@ -218,8 +221,10 @@ impl Default for Config { model: Default::default(), functions: Default::default(), working_mode: WorkingMode::Cmd, - print_info_only: false, last_message: None, + + cli_info_flag: false, + cli_agent_variables: None, } } } @@ -1508,6 +1513,7 @@ impl Config { if self.agent.take().is_some() { self.rag.take(); self.last_message = None; + self.cli_agent_variables = None; } Ok(()) } @@ -1997,14 +2003,20 @@ impl Config { None => return Ok(()), }; 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 { + config_variables.extend(v.clone()); + } let new_variables = Agent::init_agent_variables( agent.defined_variables(), - agent.config_variables(), - self.print_info_only, + &config_variables, + self.cli_info_flag, )?; agent.set_shared_variables(new_variables); } - agent.update_shared_dynamic_instructions(false)?; + if !self.cli_info_flag { + agent.update_shared_dynamic_instructions(false)?; + } Ok(()) } @@ -2017,10 +2029,14 @@ impl Config { let shared_variables = agent.shared_variables().clone(); 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 { + config_variables.extend(v.clone()); + } let new_variables = Agent::init_agent_variables( agent.defined_variables(), - agent.config_variables(), - self.print_info_only, + &config_variables, + self.cli_info_flag, )?; agent.set_shared_variables(new_variables.clone()); new_variables @@ -2028,7 +2044,9 @@ impl Config { shared_variables }; agent.set_session_variables(session_variables); - agent.update_session_dynamic_instructions(None)?; + if !self.cli_info_flag { + agent.update_session_dynamic_instructions(None)?; + } session.sync_agent(agent); } else { let variables = session.agent_variables(); diff --git a/src/main.rs b/src/main.rs index 2cefaf31..77b72c22 100644 --- a/src/main.rs +++ b/src/main.rs @@ -65,7 +65,7 @@ async fn run(config: GlobalConfig, cli: Cli, text: Option) -> Result<()> return serve::run(config, addr).await; } if cli.info { - config.write().print_info_only = true; + config.write().cli_info_flag = true; } if cli.list_models { @@ -98,6 +98,15 @@ async fn run(config: GlobalConfig, cli: Cli, text: Option) -> Result<()> Some(v) => v.as_str(), None => TEMP_SESSION_NAME, }); + if !cli.agent_variable.is_empty() { + config.write().cli_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? } else { if let Some(prompt) = &cli.prompt {