From c5506fe39377d864f80175821712b33949241583 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Sun, 21 Apr 2024 07:57:02 +0800
Subject: [PATCH] refactor: tell user session compression (#426)

---
 src/main.rs          | 20 ++------------------
 src/repl/mod.rs      | 23 ++++++++++++++++++++---
 src/utils/mod.rs     |  2 +-
 src/utils/spinner.rs | 22 +++++++++++++++++++++-
 4 files changed, 44 insertions(+), 23 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 600381b..6d5fe9b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -25,9 +25,8 @@ use is_terminal::IsTerminal;
 use parking_lot::RwLock;
 use std::io::{stderr, stdin, stdout, Read};
 use std::sync::{mpsc, Arc};
-use std::time::Duration;
 use std::{process, thread};
-use utils::Spinner;
+use utils::run_spinner;
 
 fn main() -> Result<()> {
     let cli = Cli::parse();
@@ -155,7 +154,7 @@ fn execute(config: &GlobalConfig, mut input: Input) -> Result<()> {
     let client = init_client(config)?;
     config.read().maybe_print_send_tokens(&input);
     let (tx, rx) = mpsc::sync_channel::<()>(0);
-    thread::spawn(move || run_spinner(rx));
+    thread::spawn(move || run_spinner(" Generating", rx));
     let ret = client.send_message(input.clone());
     tx.send(())?;
     let mut eval_str = ret?;
@@ -246,18 +245,3 @@ fn create_input(
     };
     Ok(Some(input))
 }
-
-fn run_spinner(rx: mpsc::Receiver<()>) -> Result<()> {
-    let mut writer = stdout();
-    let mut spinner = Spinner::new(" Generating");
-    loop {
-        spinner.step(&mut writer)?;
-        if let Ok(()) = rx.try_recv() {
-            spinner.stop(&mut writer)?;
-            break;
-        }
-        thread::sleep(Duration::from_millis(50))
-    }
-    spinner.stop(&mut writer)?;
-    Ok(())
-}
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index 73196bc..412d0be 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -14,6 +14,7 @@ use crate::utils::{create_abort_signal, set_text, AbortSignal};
 use anyhow::{bail, Context, Result};
 use fancy_regex::Regex;
 use lazy_static::lazy_static;
+use nu_ansi_term::Color;
 use reedline::{
     default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
     ColumnarMenu, EditCommand, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline,
@@ -198,11 +199,14 @@ impl Repl {
                         }
                     }
                 }
-                ".set" => {
-                    if let Some(args) = args {
+                ".set" => match args {
+                    Some(args) => {
                         self.config.write().update(args)?;
                     }
-                }
+                    _ => {
+                        println!("Usage: .set <key> <value>...")
+                    }
+                },
                 ".copy" => {
                     let config = self.config.read();
                     self.copy(config.last_reply())
@@ -266,6 +270,19 @@ impl Repl {
         self.config.read().maybe_copy(&output);
         if self.config.write().should_compress_session() {
             let config = self.config.clone();
+            let color = if config.read().light_theme {
+                Color::LightGray
+            } else {
+                Color::DarkGray
+            };
+            print!(
+                "\n📢 {}{}{}\n",
+                color.normal().paint(
+                    "Session compression is being activated because the current tokens exceed `"
+                ),
+                color.italic().paint("compress_threshold"),
+                color.normal().paint("`."),
+            );
             std::thread::spawn(move || -> anyhow::Result<()> {
                 let _ = compress_session(&config);
                 config.write().end_compressing_session();
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 3df1a09..ba438a7 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -9,7 +9,7 @@ pub use self::abort_signal::{create_abort_signal, AbortSignal};
 pub use self::clipboard::set_text;
 pub use self::prompt_input::*;
 pub use self::render_prompt::render_prompt;
-pub use self::spinner::Spinner;
+pub use self::spinner::{run_spinner, Spinner};
 pub use self::tiktoken::cl100k_base_singleton;
 
 use fancy_regex::Regex;
diff --git a/src/utils/spinner.rs b/src/utils/spinner.rs
index c684901..3715eff 100644
--- a/src/utils/spinner.rs
+++ b/src/utils/spinner.rs
@@ -1,6 +1,11 @@
 use anyhow::Result;
 use crossterm::{cursor, queue, style, terminal};
-use std::io::{Stdout, Write};
+use std::{
+    io::{stdout, Stdout, Write},
+    sync::mpsc,
+    thread,
+    time::Duration,
+};
 
 pub struct Spinner {
     index: usize,
@@ -50,3 +55,18 @@ impl Spinner {
         Ok(())
     }
 }
+
+pub fn run_spinner(message: &str, rx: mpsc::Receiver<()>) -> Result<()> {
+    let mut writer = stdout();
+    let mut spinner = Spinner::new(message);
+    loop {
+        spinner.step(&mut writer)?;
+        if let Ok(()) = rx.try_recv() {
+            spinner.stop(&mut writer)?;
+            break;
+        }
+        thread::sleep(Duration::from_millis(50))
+    }
+    spinner.stop(&mut writer)?;
+    Ok(())
+}