From bc532d99be0c407fe5e6ef82018350837ac2fb97 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Sun, 19 Jan 2025 17:32:22 +0800
Subject: feat: support `.edit config` and `.edit agent-config` repl commands
 (#1103)

---
 src/config/mod.rs     | 52 +++++++++++++++++++++++++++++++++++++++++++++++++--
 src/config/session.rs |  2 +-
 src/repl/mod.rs       | 16 ++++++++++++++--
 3 files changed, 65 insertions(+), 5 deletions(-)

(limited to 'src')

diff --git a/src/config/mod.rs b/src/config/mod.rs
index e2f9048..97e095c 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -474,6 +474,18 @@ impl Config {
         Ok((log_level, log_path))
     }
 
+    pub fn edit_config(&self) -> Result<()> {
+        let config_path = Self::config_file();
+        let editor = self.editor()?;
+        edit_file(&editor, &config_path)?;
+        println!(
+            "NOTE: Remember to restart {} if there are changes made to '{}",
+            env!("CARGO_CRATE_NAME"),
+            config_path.display(),
+        );
+        Ok(())
+    }
+
     pub fn current_model(&self) -> &Model {
         if let Some(session) = self.session.as_ref() {
             session.model()
@@ -953,6 +965,9 @@ impl Config {
         ensure_parent_exists(&role_path)?;
         let editor = self.editor()?;
         edit_file(&editor, &role_path)?;
+        if self.working_mode.is_repl() {
+            println!("✓ Saved the role to '{}'.", role_path.display());
+        }
         Ok(())
     }
 
@@ -1508,6 +1523,34 @@ impl Config {
         }
     }
 
+    pub fn edit_agent_config(&self) -> Result<()> {
+        let agent_name = match &self.agent {
+            Some(agent) => agent.name(),
+            None => bail!("No agent"),
+        };
+        let agent_config_path = Config::agent_config_file(agent_name);
+        ensure_parent_exists(&agent_config_path)?;
+        if !agent_config_path.exists() {
+            std::fs::write(
+                &agent_config_path,
+                "# see https://github.com/sigoden/aichat/blob/main/config.agent.example.yaml\n",
+            )
+            .with_context(|| {
+                format!(
+                    "Failed to write to agent config file at '{}'",
+                    agent_config_path.display()
+                )
+            })?;
+        }
+        let editor = self.editor()?;
+        edit_file(&editor, &agent_config_path)?;
+        println!(
+            "NOTE: Remember to reload the agent if there are changes made to '{}'",
+            agent_config_path.display()
+        );
+        Ok(())
+    }
+
     pub fn exit_agent(&mut self) -> Result<()> {
         self.exit_session()?;
         if self.agent.take().is_some() {
@@ -2550,7 +2593,12 @@ fn create_config_file(config_path: &Path) -> Result<()> {
     );
 
     ensure_parent_exists(config_path)?;
-    std::fs::write(config_path, config_data).with_context(|| "Failed to write to config file")?;
+    std::fs::write(config_path, config_data).with_context(|| {
+        format!(
+            "Failed to write to config file at '{}'",
+            config_path.display()
+        )
+    })?;
     #[cfg(unix)]
     {
         use std::os::unix::prelude::PermissionsExt;
@@ -2558,7 +2606,7 @@ fn create_config_file(config_path: &Path) -> Result<()> {
         std::fs::set_permissions(config_path, perms)?;
     }
 
-    println!("✓ Saved config file to '{}'.\n", config_path.display());
+    println!("✓ Saved the config file to '{}'.\n", config_path.display());
 
     Ok(())
 }
diff --git a/src/config/session.rs b/src/config/session.rs
index 13af751..2bc1b6c 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -438,7 +438,7 @@ impl Session {
         })?;
 
         if is_repl {
-            println!("✓ Saved session to '{}'.", session_path.display());
+            println!("✓ Saved the session to '{}'.", session_path.display());
         }
 
         if self.name() != session_name {
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index 908fa67..28f31bc 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -29,9 +29,10 @@ use std::{env, process};
 const MENU_NAME: &str = "completion_menu";
 
 lazy_static::lazy_static! {
-    static ref REPL_COMMANDS: [ReplCommand; 34] = [
+    static ref REPL_COMMANDS: [ReplCommand; 36] = [
         ReplCommand::new(".help", "Show this help message", AssertState::pass()),
         ReplCommand::new(".info", "View system info", AssertState::pass()),
+        ReplCommand::new(".edit config", "Edit the configuration file", AssertState::False(StateFlags::AGENT)),
         ReplCommand::new(".model", "Change the current LLM", AssertState::pass()),
         ReplCommand::new(
             ".prompt",
@@ -104,6 +105,11 @@ lazy_static::lazy_static! {
             "Use the conversation starter",
             AssertState::True(StateFlags::AGENT)
         ),
+        ReplCommand::new(
+            ".edit agent-config",
+            "Edit the agent-specific config file",
+            AssertState::True(StateFlags::AGENT),
+        ),
         ReplCommand::new(
             ".info agent",
             "View agent info",
@@ -490,6 +496,9 @@ pub async fn run_repl_command(
                     bail!("Cannot perform this operation because you are in a macro")
                 }
                 match args {
+                    Some("config") => {
+                        config.read().edit_config()?;
+                    }
                     Some("role") => {
                         config.write().edit_role()?;
                     }
@@ -499,8 +508,11 @@ pub async fn run_repl_command(
                     Some("rag-docs") => {
                         Config::edit_rag_docs(config, abort_signal.clone()).await?;
                     }
+                    Some("agent-config") => {
+                        config.write().edit_agent_config()?;
+                    }
                     _ => {
-                        println!(r#"Usage: .edit <role|session|rag-docs>"#)
+                        println!(r#"Usage: .edit <config|role|session|rag-docs|agent-config>"#)
                     }
                 }
             }
-- 
cgit v1.2.3