feat: support multiple models (#71)

3 files changed, 14 insertions, 4 deletions

diff --git a/src/repl/handler.rs b/src/repl/handler.rs
index 159127c..014a851 100644
--- a/src/repl/handler.rs
+++ b/src/repl/handler.rs
@@ -12,6 +12,7 @@ use std::cell::RefCell;
 
 pub enum ReplCmd {
     Submit(String),
+    SetModel(String),
     SetRole(String),
     UpdateConfig(String),
     Prompt(String),
@@ -65,6 +66,10 @@ impl ReplCmdHandler {
                 self.config.write().save_conversation(&input, &buffer)?;
                 *self.reply.borrow_mut() = buffer;
             }
+            ReplCmd::SetModel(name) => {
+                self.config.write().set_model(&name)?;
+                print_now!("\n");
+            }
             ReplCmd::SetRole(name) => {
                 let output = self.config.write().change_role(&name)?;
                 print_now!("{}\n\n", output.trim_end());
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index 337afb3..663decb 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -19,9 +19,10 @@ use reedline::Signal;
 use std::borrow::Cow;
 use std::sync::Arc;
 
-pub const REPL_COMMANDS: [(&str, &str); 11] = [
+pub const REPL_COMMANDS: [(&str, &str); 12] = [
     (".info", "Print the information"),
     (".set", "Modify the configuration temporarily"),
+    (".model", "Choose a model"),
     (".prompt", "Add a GPT prompt"),
     (".role", "Select a role"),
     (".clear role", "Clear the currently selected role"),
@@ -109,6 +110,10 @@ impl Repl {
                     self.editor.print_history()?;
                     print_now!("\n");
                 }
+                ".model" => match args {
+                    Some(name) => handler.handle(ReplCmd::SetModel(name.to_string()))?,
+                    None => print_now!("Usage: .model <name>\n\n"),
+                },
                 ".role" => match args {
                     Some(name) => handler.handle(ReplCmd::SetRole(name.to_string()))?,
                     None => print_now!("Usage: .role <name>\n\n"),
diff --git a/src/repl/prompt.rs b/src/repl/prompt.rs
index 1a59ba3..8c670f4 100644
--- a/src/repl/prompt.rs
+++ b/src/repl/prompt.rs
@@ -76,10 +76,10 @@ impl Prompt for ReplPrompt {
     }
 
     fn render_prompt_right(&self) -> Cow<str> {
-        if let Some(conversation) = self.config.read().conversation.as_ref() {
-            conversation.reamind_tokens().to_string().into()
-        } else {
+        if self.config.read().conversation.is_none() {
             Cow::Borrowed("")
+        } else {
+            self.config.read().get_reamind_tokens().to_string().into()
         }
     }