From e443905dc1f16202b29a8cdc783069ea72c2ac81 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Sun, 3 Mar 2024 15:25:21 +0800
Subject: feat: add `.clear messages` to clear session messages (#332)

---
 src/config/mod.rs     | 16 ++++++++++++++--
 src/config/session.rs |  9 +++++++++
 src/repl/mod.rs       | 13 ++++++++++---
 3 files changed, 33 insertions(+), 5 deletions(-)

(limited to 'src')

diff --git a/src/config/mod.rs b/src/config/mod.rs
index d6ea6db..804a2fd 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -632,6 +632,13 @@ impl Config {
         Ok(())
     }
 
+    pub fn clear_session_messages(&mut self) -> Result<()> {
+        if let Some(session) = self.session.as_mut() {
+            session.clear_messgaes();
+        }
+        Ok(())
+    }
+
     pub fn list_sessions(&self) -> Vec<String> {
         let sessions_dir = match Self::sessions_dir() {
             Ok(dir) => dir,
@@ -951,10 +958,15 @@ impl State {
             .collect()
     }
 
-    pub fn can_change_role() -> Vec<Self> {
+    pub fn unable_change_role() -> Vec<Self> {
+        vec![Self::Session]
+    }
+
+    pub fn able_change_role() -> Vec<Self> {
+        let excludes: HashSet<_> = Self::unable_change_role().into_iter().collect();
         Self::all()
             .into_iter()
-            .filter(|v| *v != Self::Session)
+            .filter(|v| !excludes.contains(v))
             .collect()
     }
 
diff --git a/src/config/session.rs b/src/config/session.rs
index cbfe1c2..076fa27 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -258,6 +258,15 @@ impl Session {
         Ok(())
     }
 
+    pub fn clear_messgaes(&mut self) {
+        if self.messages.is_empty() {
+            return;
+        }
+        self.messages.clear();
+        self.data_urls.clear();
+        self.dirty = true;
+    }
+
     pub fn echo_messages(&self, input: &Input) -> String {
         let messages = self.build_emssages(input);
         serde_yaml::to_string(&messages).unwrap_or_else(|_| "Unable to echo message".into())
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index d6fa7f2..e11b495 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -25,11 +25,11 @@ use std::{env, process};
 const MENU_NAME: &str = "completion_menu";
 
 lazy_static! {
-    static ref REPL_COMMANDS: [ReplCommand; 13] = [
+    static ref REPL_COMMANDS: [ReplCommand; 14] = [
         ReplCommand::new(".help", "Print this help message", State::all()),
         ReplCommand::new(".info", "Print system info", State::all()),
         ReplCommand::new(".model", "Switch LLM model", State::all()),
-        ReplCommand::new(".role", "Use a role", State::can_change_role()),
+        ReplCommand::new(".role", "Use a role", State::able_change_role()),
         ReplCommand::new(".info role", "Show role info", State::in_role(),),
         ReplCommand::new(".exit role", "Leave current role", State::in_role(),),
         ReplCommand::new(
@@ -38,6 +38,11 @@ lazy_static! {
             State::notin_session(),
         ),
         ReplCommand::new(".info session", "Show session info", State::in_session(),),
+        ReplCommand::new(
+            ".clear messages",
+            "Clear messages in the session",
+            State::unable_change_role()
+        ),
         ReplCommand::new(
             ".exit session",
             "End the current session",
@@ -225,8 +230,10 @@ impl Repl {
                         return Ok(true);
                     }
                 },
-                // deprecated this command
                 ".clear" => match args {
+                    Some("messages") => {
+                        self.config.write().clear_session_messages()?;
+                    }
                     Some("role") => {
                         println!(r#"Deprecated. Use ".exit role" instead."#);
                     }
-- 
cgit v1.2.3