refactor: improve code regarding tools/agents (#1021)

3 files changed, 36 insertions, 19 deletions

diff --git a/src/config/agent.rs b/src/config/agent.rs
index 49d5835..29121ee 100644
--- a/src/config/agent.rs
+++ b/src/config/agent.rs
@@ -10,15 +10,17 @@ use serde::{Deserialize, Serialize};
 
 const DEFAULT_AGENT_NAME: &str = "rag";
 
+pub type AgentVariables = IndexMap<String, String>;
+
 #[derive(Debug, Clone, Serialize)]
 pub struct Agent {
     name: String,
     config: AgentConfig,
     definition: AgentDefinition,
     #[serde(skip)]
-    shared_variables: IndexMap<String, String>,
+    shared_variables: AgentVariables,
     #[serde(skip)]
-    session_variables: Option<IndexMap<String, String>>,
+    session_variables: Option<AgentVariables>,
     #[serde(skip)]
     functions: Functions,
     #[serde(skip)]
@@ -108,9 +110,9 @@ impl Agent {
 
     pub fn init_agent_variables(
         agent_variables: &[AgentVariable],
-        variables: &IndexMap<String, String>,
+        variables: &AgentVariables,
         no_interaction: bool,
-    ) -> Result<IndexMap<String, String>> {
+    ) -> Result<AgentVariables> {
         let mut output = IndexMap::new();
         if agent_variables.is_empty() {
             return Ok(output);
@@ -224,26 +226,41 @@ impl Agent {
         self.config.agent_prelude = value;
     }
 
-    pub fn variables(&self) -> &IndexMap<String, String> {
+    pub fn variables(&self) -> &AgentVariables {
         match &self.session_variables {
             Some(variables) => variables,
             None => &self.shared_variables,
         }
     }
 
-    pub fn config_variables(&self) -> &IndexMap<String, String> {
+    pub fn variable_envs(&self) -> HashMap<String, String> {
+        self.variables()
+            .iter()
+            .map(|(k, v)| {
+                (
+                    format!("LLM_AGENT_VAR_{}", normalize_env_name(k)),
+                    v.clone(),
+                )
+            })
+            .collect()
+    }
+
+    pub fn config_variables(&self) -> &AgentVariables {
         &self.config.variables
     }
 
-    pub fn shared_variables(&self) -> &IndexMap<String, String> {
+    pub fn shared_variables(&self) -> &AgentVariables {
         &self.shared_variables
     }
 
-    pub fn set_shared_variables(&mut self, shared_variables: IndexMap<String, String>) {
+    pub fn set_shared_variables(&mut self, shared_variables: AgentVariables) {
         self.shared_variables = shared_variables;
     }
 
-    pub fn set_session_variables(&mut self, session_variables: Option<IndexMap<String, String>>) {
+    pub fn set_session_variables(&mut self, session_variables: Option<AgentVariables>) {
+        if self.shared_variables.is_empty() {
+            self.shared_variables = session_variables.clone().unwrap_or_default();
+        }
         self.session_variables = session_variables;
     }
 
@@ -319,7 +336,7 @@ pub struct AgentConfig {
     pub use_tools: Option<String>,
     pub agent_prelude: Option<String>,
     #[serde(default)]
-    pub variables: IndexMap<String, String>,
+    pub variables: AgentVariables,
 }
 
 impl AgentConfig {
@@ -419,7 +436,7 @@ impl AgentDefinition {
         )
     }
 
-    fn interpolated_instructions(&self, variables: &IndexMap<String, String>) -> String {
+    fn interpolated_instructions(&self, variables: &AgentVariables) -> String {
         let mut output = self.instructions.clone();
         for (k, v) in variables {
             output = output.replace(&format!("{{{{{k}}}}}"), v)
diff --git a/src/config/mod.rs b/src/config/mod.rs
index e578620..31337cd 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -3,7 +3,7 @@ mod input;
 mod role;
 mod session;
 
-pub use self::agent::{list_agents, Agent};
+pub use self::agent::{list_agents, Agent, AgentVariables};
 pub use self::input::Input;
 pub use self::role::{
     Role, RoleLike, CODE_ROLE, CREATE_TITLE_ROLE, EXPLAIN_SHELL_ROLE, SHELL_ROLE,
@@ -1487,10 +1487,13 @@ impl Config {
         if parts.len() != 2 {
             bail!("Usage: .variable <key> <value>");
         }
-        let key = parts[0];
-        let value = parts[1];
         match self.agent.as_mut() {
             Some(agent) => {
+                if let Some(session) = self.session.as_ref() {
+                    session.guard_empty()?;
+                }
+                let key = parts[0];
+                let value = parts[1];
                 agent.set_variable(key, value)?;
                 if let Some(session) = self.session.as_mut() {
                     session.sync_agent(agent, true);
@@ -2009,9 +2012,6 @@ impl Config {
             &all_variables,
             self.print_info_only,
         )?;
-        if shared_variables.is_empty() {
-            agent.set_shared_variables(new_variables.clone());
-        }
         agent.set_session_variables(Some(new_variables));
         session.sync_agent(agent, false);
         Ok(())
diff --git a/src/config/session.rs b/src/config/session.rs
index 23658fd..4b3da76 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -35,7 +35,7 @@ pub struct Session {
     #[serde(skip_serializing_if = "Option::is_none")]
     role_name: Option<String>,
     #[serde(default, skip_serializing_if = "IndexMap::is_empty")]
-    agent_variables: IndexMap<String, String>,
+    agent_variables: AgentVariables,
 
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
     compressed_messages: Vec<Message>,
@@ -284,7 +284,7 @@ impl Session {
         }
     }
 
-    pub fn agent_variables(&self) -> &IndexMap<String, String> {
+    pub fn agent_variables(&self) -> &AgentVariables {
         &self.agent_variables
     }