From b965c63be05c726c996ece298d6f5f2514a19316 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Mon, 17 Jun 2024 19:54:24 +0800
Subject: refactor: minor improvement (#611)

---
 src/config/mod.rs | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

(limited to 'src/config/mod.rs')

diff --git a/src/config/mod.rs b/src/config/mod.rs
index cb510ef..41e449a 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -583,7 +583,8 @@ impl Config {
     }
 
     pub fn use_prompt(&mut self, prompt: &str) -> Result<()> {
-        let role = Role::new(TEMP_ROLE_NAME, prompt);
+        let mut role = Role::new(TEMP_ROLE_NAME, prompt);
+        role.set_model(&self.model);
         self.use_role_obj(role)
     }
 
@@ -705,9 +706,8 @@ impl Config {
 
     pub fn exit_session(&mut self) -> Result<()> {
         if let Some(mut session) = self.session.take() {
-            let is_repl = self.working_mode == WorkingMode::Repl;
             let sessions_dir = self.sessions_dir()?;
-            session.exit(&sessions_dir, is_repl)?;
+            session.exit(&sessions_dir, self.working_mode.is_repl())?;
             self.last_message = None;
         }
         Ok(())
@@ -723,7 +723,7 @@ impl Config {
         };
         let session_path = self.session_file(&name)?;
         if let Some(session) = self.session.as_mut() {
-            session.save(&session_path)?;
+            session.save(&session_path, self.working_mode.is_repl())?;
         }
         Ok(())
     }
@@ -933,8 +933,8 @@ impl Config {
     }
 
     pub fn exit_bot(&mut self) -> Result<()> {
+        self.exit_session()?;
         if self.bot.take().is_some() {
-            self.exit_session()?;
             self.rag.take();
             self.last_message = None;
         }
@@ -1420,6 +1420,12 @@ pub enum WorkingMode {
     Serve,
 }
 
+impl WorkingMode {
+    pub fn is_repl(&self) -> bool {
+        *self == WorkingMode::Repl
+    }
+}
+
 bitflags::bitflags! {
     #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
     pub struct StateFlags: u32 {
-- 
cgit v1.2.3