From b05b730cb5c1acdb47611a67bb98325bd86c2969 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 11 Jun 2024 15:44:35 +0800
Subject: feat: add config `repl_prelude` and `bot_prelude` (#584)

---
 config.example.yaml | 20 ++++++++++++--------
 src/config/mod.rs   | 36 +++++++++++++++++++++++++-----------
 2 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/config.example.yaml b/config.example.yaml
index 6708243..d040c66 100644
--- a/config.example.yaml
+++ b/config.example.yaml
@@ -9,7 +9,9 @@ wrap: no                         # Controls text wrapping (no, auto, <max-width>
 wrap_code: false                 # Enables or disables wrapping of code blocks
 auto_copy: false                 # Enables or disables automatic copying the last LLM response to the clipboard 
 keybindings: emacs               # Choose keybinding style (emacs, vi)
-prelude: null                    # Set a default role or session to start with (role:<name>, session:<name>)
+
+prelude: null                    # Set a default role or session to start with (e.g. role:<name>, session:<name>)
+repl_prelude: null               # This overrides 'prelude' when using REPL
 
 # Command that will be used to edit the current line buffer with ctrl+o
 # if unset fallback to $EDITOR and $VISUAL
@@ -23,6 +25,15 @@ function_calling: false
 # e.g. 'execute_command|execute_js_code' 'execute_.*'
 dangerously_functions: null
 
+bot_prelude: null               # Set a session to use when starting a bot.  (e.g. temp, default)
+
+bots:
+  - name: todo-sh
+    model: null
+    temperature: null
+    top_p: null
+    dangerously_functions: null
+
 # Specifies the embedding model to use
 embedding_model: null
 
@@ -253,10 +264,3 @@ clients:
     name: together
     api_base: https://api.together.xyz/v1
     api_key: xxx                                      # ENV: {client}_API_KEY
-
-bots:
-  - name: todo-sh
-    model: null
-    temperature: null
-    top_p: null
-    dangerously_functions: null
\ No newline at end of file
diff --git a/src/config/mod.rs b/src/config/mod.rs
index b088282..22eaa2e 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -89,19 +89,21 @@ pub struct Config {
     pub auto_copy: bool,
     pub keybindings: Keybindings,
     pub prelude: Option<String>,
+    pub repl_prelude: Option<String>,
     pub buffer_editor: Option<String>,
+    pub function_calling: bool,
+    pub dangerously_functions: Option<FunctionsFilter>,
+    pub bot_prelude: Option<String>,
+    pub bots: Vec<BotConfig>,
     pub embedding_model: Option<String>,
     pub rag_top_k: usize,
     pub rag_template: Option<String>,
-    pub function_calling: bool,
-    pub dangerously_functions: Option<FunctionsFilter>,
     pub compress_threshold: usize,
     pub summarize_prompt: Option<String>,
     pub summary_prompt: Option<String>,
     pub left_prompt: Option<String>,
     pub right_prompt: Option<String>,
     pub clients: Vec<ClientConfig>,
-    pub bots: Vec<BotConfig>,
     #[serde(skip)]
     pub roles: Vec<Role>,
     #[serde(skip)]
@@ -138,19 +140,21 @@ impl Default for Config {
             auto_copy: false,
             keybindings: Default::default(),
             prelude: None,
+            repl_prelude: None,
             buffer_editor: None,
+            function_calling: false,
+            dangerously_functions: None,
+            bot_prelude: None,
+            bots: vec![],
             embedding_model: None,
             rag_top_k: 4,
             rag_template: None,
-            function_calling: false,
-            dangerously_functions: None,
             compress_threshold: 4000,
             summarize_prompt: None,
             summary_prompt: None,
             left_prompt: None,
             right_prompt: None,
             clients: vec![],
-            bots: vec![],
             roles: vec![],
             role: None,
             session: None,
@@ -901,9 +905,13 @@ impl Config {
         if config.read().bot.is_some() {
             bail!("Already in a bot, please run '.exit bot' first to exit the current bot.");
         }
+        let prelude = config.read().bot_prelude.clone();
         let bot = Bot::init(config, name, abort_signal).await?;
         config.write().rag = bot.rag();
         config.write().bot = Some(bot);
+        if let Some(session) = prelude {
+            config.write().use_session(Some(&session))?;
+        }
         Ok(())
     }
 
@@ -922,14 +930,20 @@ impl Config {
     }
 
     pub fn apply_prelude(&mut self) -> Result<()> {
-        let prelude = self.prelude.clone().unwrap_or_default();
-        if prelude.is_empty() {
-            return Ok(());
-        }
+        let prelude = match self.working_mode {
+            WorkingMode::Command => self.prelude.as_ref(),
+            WorkingMode::Repl => self.repl_prelude.as_ref().or(self.prelude.as_ref()),
+            WorkingMode::Serve => return Ok(()),
+        };
+        let prelude = match prelude {
+            Some(v) => v.to_string(),
+            None => return Ok(()),
+        };
+
         let err_msg = || format!("Invalid prelude '{}", prelude);
         match prelude.split_once(':') {
             Some(("role", name)) => {
-                if self.role.is_none() && self.session.is_none() {
+                if self.state().is_empty() {
                     self.use_role(name).with_context(err_msg)?;
                 }
             }
-- 
cgit v1.2.3