From 5d763fc10c4e3d40268ead473ccb0000b4923bd1 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Sat, 20 Apr 2024 09:21:52 +0800
Subject: [PATCH] feat: add `.prompt` repl command (#420)

---
 src/config/mod.rs  |  5 +++++
 src/config/role.rs |  9 +++++++++
 src/repl/mod.rs    | 15 +++++++++++++--
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 1fe97af..2a3a866 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -301,6 +301,11 @@ impl Config {
         Ok(path)
     }
 
+    pub fn set_prompt(&mut self, prompt: &str) -> Result<()> {
+        let role = Role::temp(prompt);
+        self.set_role_obj(role)
+    }
+
     pub fn set_role(&mut self, name: &str) -> Result<()> {
         let role = self.retrieve_role(name)?;
         self.set_role_obj(role)
diff --git a/src/config/role.rs b/src/config/role.rs
index 5a6ff3a..50d5b5e 100644
--- a/src/config/role.rs
+++ b/src/config/role.rs
@@ -7,6 +7,7 @@ use crate::{
 use anyhow::{Context, Result};
 use serde::{Deserialize, Serialize};
 
+pub const TEMP_ROLE: &str = "%%";
 pub const SHELL_ROLE: &str = "%shell%";
 pub const EXPLAIN_ROLE: &str = "%explain%";
 pub const CODE_ROLE: &str = "%code%";
@@ -24,6 +25,14 @@ pub struct Role {
 }
 
 impl Role {
+    pub fn temp(prompt: &str) -> Self {
+        Self {
+            name: TEMP_ROLE.into(),
+            prompt: prompt.into(),
+            temperature: None,
+        }
+    }
+
     pub fn find_system_role(name: &str) -> Option<Self> {
         match name {
             SHELL_ROLE => Some(Self::shell()),
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index 3f30317..564bdf9 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -25,10 +25,15 @@ use std::{env, process};
 const MENU_NAME: &str = "completion_menu";
 
 lazy_static! {
-    static ref REPL_COMMANDS: [ReplCommand; 15] = [
+    static ref REPL_COMMANDS: [ReplCommand; 16] = [
         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(
+            ".prompt",
+            "Use a temp role with this prompt",
+            State::able_change_role()
+        ),
         ReplCommand::new(".role", "Use a role", State::able_change_role()),
         ReplCommand::new(".info role", "Show the role info", State::in_role(),),
         ReplCommand::new(".exit role", "Leave current role", State::in_role(),),
@@ -169,6 +174,12 @@ impl Repl {
                     }
                     None => println!("Usage: .model <name>"),
                 },
+                ".prompt" => match args {
+                    Some(text) => {
+                        self.config.write().set_prompt(text)?;
+                    }
+                    None => println!("Usage: .prompt <text>..."),
+                },
                 ".role" => match args {
                     Some(args) => match args.split_once(|c| c == '\n' || c == ' ') {
                         Some((name, text)) => {
@@ -181,7 +192,7 @@ impl Repl {
                             self.config.write().set_role(args)?;
                         }
                     },
-                    None => println!(r#"Usage: .role <name> [text...]"#),
+                    None => println!(r#"Usage: .role <name> [text]..."#),
                 },
                 ".session" => {
                     self.config.write().start_session(args)?;