From a9268b600fb378400795fbaf98bf923372b4c19a Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 9 Jul 2024 21:22:51 +0800
Subject: feat: support agent variables (#692)

---
 src/config/mod.rs | 33 ++++++++++++++++++++++++++++++---
 1 file changed, 30 insertions(+), 3 deletions(-)

(limited to 'src/config/mod.rs')

diff --git a/src/config/mod.rs b/src/config/mod.rs
index a243a83..154dcd9 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -50,6 +50,7 @@ const FUNCTIONS_FILE_NAME: &str = "functions.json";
 const FUNCTIONS_BIN_DIR_NAME: &str = "bin";
 const AGENTS_DIR_NAME: &str = "agents";
 const AGENT_RAG_FILE_NAME: &str = "rag.bin";
+const AGENT_VARIABLES_FILE_NAME: &str = "variables.yaml";
 
 pub const TEMP_ROLE_NAME: &str = "%%";
 pub const TEMP_RAG_NAME: &str = "temp";
@@ -352,6 +353,10 @@ impl Config {
         Ok(Self::agent_config_dir(name)?.join(AGENT_RAG_FILE_NAME))
     }
 
+    pub fn agent_variables_file(name: &str) -> Result<PathBuf> {
+        Ok(Self::agent_config_dir(name)?.join(AGENT_VARIABLES_FILE_NAME))
+    }
+
     pub fn agents_functions_dir() -> Result<PathBuf> {
         match env::var(get_env_name("agents_functions_dir")) {
             Ok(value) => Ok(PathBuf::from(value)),
@@ -1014,6 +1019,20 @@ impl Config {
         }
     }
 
+    pub fn set_agent_variable(&mut self, data: &str) -> Result<()> {
+        let parts: Vec<&str> = data.split_whitespace().collect();
+        if parts.len() != 2 {
+            bail!("Usage: .variable <key> <value>");
+        }
+        let key = parts[0];
+        let value = parts[1];
+        match self.agent.as_mut() {
+            Some(agent) => agent.set_variable(key, value)?,
+            None => bail!("No agent"),
+        };
+        Ok(())
+    }
+
     pub fn exit_agent(&mut self) -> Result<()> {
         self.exit_session()?;
         if self.agent.take().is_some() {
@@ -1180,6 +1199,14 @@ impl Config {
                         .collect(),
                     None => vec![],
                 },
+                ".variable" => match &self.agent {
+                    Some(agent) => agent
+                        .variables()
+                        .iter()
+                        .map(|v| (v.name.clone(), Some(v.description.clone())))
+                        .collect(),
+                    None => vec![],
+                },
                 ".set" => vec![
                     "max_output_tokens",
                     "temperature",
@@ -1200,7 +1227,7 @@ impl Config {
                 _ => vec![],
             };
             (values, args[0])
-        } else if args.len() == 2 {
+        } else if args.len() == 2 && cmd == ".set" {
             let values = match args[0] {
                 "max_output_tokens" => match self.model.max_output_tokens() {
                     Some(v) => vec![v.to_string()],
@@ -1693,11 +1720,11 @@ pub(crate) fn ensure_parent_exists(path: &Path) -> Result<()> {
     }
     let parent = path
         .parent()
-        .ok_or_else(|| anyhow!("Failed to write to {}, No parent path", path.display()))?;
+        .ok_or_else(|| anyhow!("Failed to write to '{}', No parent path", path.display()))?;
     if !parent.exists() {
         create_dir_all(parent).with_context(|| {
             format!(
-                "Failed to write {}, Cannot create parent directory",
+                "Failed to write to '{}', Cannot create parent directory",
                 path.display()
             )
         })?;
-- 
cgit v1.2.3