feat: add cli option --agent-variable <name> <value> (#1027)

sigoden · Dec 1, 2024 · 6b6413b · 6b6413b
1 parent 7e8e2a6
commit 6b6413b
Show file tree

Hide file tree

Showing 9 changed files with 56 additions and 28 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -24,8 +24,11 @@ pub struct Cli {
     /// Start a agent
     #[clap(short = 'a', long)]
     pub agent: Option<String>,
+    /// Set agent variables
+    #[clap(long, value_names = ["NAME", "VALUE"], num_args = 2)]
+    pub agent_variable: Vec<String>,
     /// Start a RAG
-    #[clap(short = 'R', long)]
+    #[clap(long)]
     pub rag: Option<String>,
     /// Serve the LLM API and WebAPP
     #[clap(long, value_name = "ADDRESS")]

diff --git a/src/config/agent.rs b/src/config/agent.rs
@@ -15,24 +15,17 @@ const DEFAULT_AGENT_NAME: &str = "rag";
 
 pub type AgentVariables = IndexMap<String, String>;
 
-#[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<AgentVariables>,
-    #[serde(skip)]
     shared_dynamic_instructions: Option<String>,
-    #[serde(skip)]
     session_dynamic_instructions: Option<String>,
-    #[serde(skip)]
     functions: Functions,
-    #[serde(skip)]
     rag: Option<Arc<Rag>>,
-    #[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<String> {
         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
@@ -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<AgentVariables>,
 }
 
 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,18 +2029,24 @@ 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
                 } else {
                     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
@@ -65,7 +65,7 @@ async fn run(config: GlobalConfig, cli: Cli, text: Option<String>) -> 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<String>) -> 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 {