feat: make built-in roles selectable and rename them (#412)

2 files changed, 26 insertions, 48 deletions

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 454dde3..88235ac 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -4,6 +4,7 @@ mod session;
 
 pub use self::input::{Input, InputContext};
 use self::role::Role;
+pub use self::role::{CODE_ROLE, EXPLAIN_ROLE, SHELL_ROLE};
 use self::session::{Session, TEMP_SESSION_NAME};
 
 use crate::client::{
@@ -201,6 +202,7 @@ impl Config {
                 role.complete_prompt_args(name);
                 role
             })
+            .or_else(|| Role::find_system_role(name))
             .ok_or_else(|| anyhow!("Unknown role `{name}`"))
     }
 
@@ -296,27 +298,6 @@ impl Config {
         self.set_role_obj(role)
     }
 
-    pub fn set_execute_role(&mut self) -> Result<()> {
-        let role = self
-            .retrieve_role(Role::EXECUTE)
-            .unwrap_or_else(|_| Role::for_execute());
-        self.set_role_obj(role)
-    }
-
-    pub fn set_describe_command_role(&mut self) -> Result<()> {
-        let role = self
-            .retrieve_role(Role::DESCRIBE_COMMAND)
-            .unwrap_or_else(|_| Role::for_describe_command());
-        self.set_role_obj(role)
-    }
-
-    pub fn set_code_role(&mut self) -> Result<()> {
-        let role = self
-            .retrieve_role(Role::CODE)
-            .unwrap_or_else(|_| Role::for_code());
-        self.set_role_obj(role)
-    }
-
     pub fn set_role_obj(&mut self, role: Role) -> Result<()> {
         if let Some(session) = self.session.as_mut() {
             session.guard_empty()?;
@@ -502,12 +483,7 @@ impl Config {
     pub fn repl_complete(&self, cmd: &str, args: &[&str]) -> Vec<String> {
         let (values, filter) = if args.len() == 1 {
             let values = match cmd {
-                ".role" => self
-                    .roles
-                    .iter()
-                    .filter(|v| !v.is_system())
-                    .map(|v| v.name.clone())
-                    .collect(),
+                ".role" => self.roles.iter().map(|v| v.name.clone()).collect(),
                 ".model" => list_models(self).into_iter().map(|v| v.id()).collect(),
                 ".session" => self.list_sessions(),
                 ".set" => vec![
diff --git a/src/config/role.rs b/src/config/role.rs
index 9042da3..5a6ff3a 100644
--- a/src/config/role.rs
+++ b/src/config/role.rs
@@ -1,3 +1,4 @@
+use super::Input;
 use crate::{
     client::{Message, MessageContent, MessageRole},
     utils::{detect_os, detect_shell},
@@ -6,9 +7,11 @@ use crate::{
 use anyhow::{Context, Result};
 use serde::{Deserialize, Serialize};
 
-use super::Input;
+pub const SHELL_ROLE: &str = "%shell%";
+pub const EXPLAIN_ROLE: &str = "%explain%";
+pub const CODE_ROLE: &str = "%code%";
 
-const INPUT_PLACEHOLDER: &str = "__INPUT__";
+pub const INPUT_PLACEHOLDER: &str = "__INPUT__";
 
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct Role {
@@ -21,19 +24,25 @@ pub struct Role {
 }
 
 impl Role {
-    pub const EXECUTE: &'static str = "__execute__";
-    pub const DESCRIBE_COMMAND: &'static str = "__describe_command__";
-    pub const CODE: &'static str = "__code__";
+    pub fn find_system_role(name: &str) -> Option<Self> {
+        match name {
+            SHELL_ROLE => Some(Self::shell()),
+            EXPLAIN_ROLE => Some(Self::explain()),
+            CODE_ROLE => Some(Self::code()),
+            _ => None,
+        }
+    }
 
-    pub fn for_execute() -> Self {
+    pub fn shell() -> Self {
         let os = detect_os();
-        let (shell, _, _) = detect_shell();
-        let (shell, use_semicolon) = match (shell.as_str(), os.as_str()) {
+        let (detected_shell, _, _) = detect_shell();
+        let (shell, use_semicolon) = match (detected_shell.as_str(), os.as_str()) {
+            // GPT doesn’t know much about nushell
             ("nushell", "windows") => ("cmd", true),
             ("nushell", _) => ("bash", true),
             ("powershell", _) => ("powershell", true),
             ("pwsh", _) => ("powershell", false),
-            _ => (shell.as_str(), false),
+            _ => (detected_shell.as_str(), false),
         };
         let combine = if use_semicolon {
             "\nIf multiple steps required try to combine them together using ';'.\nIf it already combined with '&&' try to replace it with ';'.".to_string()
@@ -41,7 +50,7 @@ impl Role {
             "\nIf multiple steps required try to combine them together using &&.".to_string()
         };
         Self {
-            name: Self::EXECUTE.into(),
+            name: SHELL_ROLE.into(),
             prompt: format!(
                 r#"Provide only {shell} commands for {os} without any description.
 Ensure the output is a valid {shell} command. {combine}
@@ -52,9 +61,9 @@ Output plain text only, without any markdown formatting."#
         }
     }
 
-    pub fn for_describe_command() -> Self {
+    pub fn explain() -> Self {
         Self {
-            name: Self::DESCRIBE_COMMAND.into(),
+            name: EXPLAIN_ROLE.into(),
             prompt: r#"Provide a terse, single sentence description of the given shell command.
 Describe each argument and option of the command.
 Provide short responses in about 80 words.
@@ -64,9 +73,9 @@ APPLY MARKDOWN formatting when possible."#
         }
     }
 
-    pub fn for_code() -> Self {
+    pub fn code() -> Self {
         Self {
-            name: Self::CODE.into(),
+            name: CODE_ROLE.into(),
             prompt: r#"Provide only code, without comments or explanations.
 If there is a lack of details, provide most logical solution, without requesting further clarification."#
                 .into(),
@@ -74,13 +83,6 @@ If there is a lack of details, provide most logical solution, without requesting
         }
     }
 
-    pub fn is_system(&self) -> bool {
-        matches!(
-            self.name.as_str(),
-            Self::EXECUTE | Self::DESCRIBE_COMMAND | Self::CODE
-        )
-    }
-
     pub fn export(&self) -> Result<String> {
         let output = serde_yaml::to_string(&self)
             .with_context(|| format!("Unable to show info about role {}", &self.name))?;