From 6dc332930bf2e6d31c937d87e7ac7ae14c53b7c6 Mon Sep 17 00:00:00 2001 From: sigoden Date: Sun, 14 Apr 2024 14:43:42 +0800 Subject: feat: make built-in roles selectable and rename them (#412) --- src/config/mod.rs | 30 +++--------------------------- src/config/role.rs | 44 +++++++++++++++++++++++--------------------- src/main.rs | 20 ++++++++++---------- 3 files changed, 36 insertions(+), 58 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 { 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 { + 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 { let output = serde_yaml::to_string(&self) .with_context(|| format!("Unable to show info about role {}", &self.name))?; diff --git a/src/main.rs b/src/main.rs index d203710..dff191b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,23 +10,23 @@ extern crate log; mod utils; use crate::cli::Cli; -use crate::config::{Config, GlobalConfig}; -use crate::utils::{extract_block, run_command, CODE_BLOCK_RE}; +use crate::client::{ensure_model_capabilities, init_client, list_models}; +use crate::config::{Config, GlobalConfig, Input, CODE_ROLE, EXPLAIN_ROLE, SHELL_ROLE}; +use crate::render::{render_error, render_stream, MarkdownRender}; +use crate::repl::Repl; +use crate::utils::{ + cl100k_base_singleton, create_abort_signal, extract_block, run_command, CODE_BLOCK_RE, +}; use anyhow::{bail, Result}; use clap::Parser; -use client::{ensure_model_capabilities, init_client, list_models}; -use config::Input; use inquire::validator::Validation; use inquire::Text; use is_terminal::IsTerminal; use parking_lot::RwLock; -use render::{render_error, render_stream, MarkdownRender}; -use repl::Repl; use std::io::{stderr, stdin, stdout, Read}; use std::process; use std::sync::Arc; -use utils::{cl100k_base_singleton, create_abort_signal}; fn main() -> Result<()> { let cli = Cli::parse(); @@ -63,9 +63,9 @@ fn main() -> Result<()> { if let Some(name) = &cli.role { config.write().set_role(name)?; } else if cli.execute { - config.write().set_execute_role()?; + config.write().set_role(SHELL_ROLE)?; } else if cli.code { - config.write().set_code_role()?; + config.write().set_role(CODE_ROLE)?; } if let Some(session) = &cli.session { config @@ -190,7 +190,7 @@ fn execute(config: &GlobalConfig, input: Input) -> Result<()> { } "D" | "d" => { if !describe { - config.write().set_describe_command_role()?; + config.write().set_role(EXPLAIN_ROLE)?; } let input = Input::from_str(&eval_str, config.read().input_context()); let abort = create_abort_signal(); -- cgit v1.2.3