From d9b8eabf23ae96fb395a8059c1e08840609499f4 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Sun, 28 Apr 2024 19:18:56 +0800
Subject: [PATCH] feat: add `--prompt` option (#454)

---
 scripts/completions/aichat.bash | 6 +++++-
 scripts/completions/aichat.fish | 1 +
 scripts/completions/aichat.nu   | 1 +
 scripts/completions/aichat.ps1  | 1 +
 scripts/completions/aichat.zsh  | 1 +
 src/cli.rs                      | 3 +++
 src/main.rs                     | 4 +++-
 src/repl/mod.rs                 | 2 +-
 8 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/scripts/completions/aichat.bash b/scripts/completions/aichat.bash
index f0b9bc3..cb9b708 100644
--- a/scripts/completions/aichat.bash
+++ b/scripts/completions/aichat.bash
@@ -19,7 +19,7 @@ _aichat() {
 
     case "${cmd}" in
         aichat)
-            opts="-m -r -s -e -c -f -H -S -w -h -V --model --role --session --save-session --serve --execute --code --file --no-highlight --no-stream --wrap --light-theme --dry-run --info --list-models --list-roles --list-sessions --help --version"
+            opts="-m -r -s -e -c -f -H -S -w -h -V --model --prompt --role --session --save-session --serve --execute --code --file --no-highlight --no-stream --wrap --light-theme --dry-run --info --list-models --list-roles --list-sessions --help --version"
             if [[ ${cur} == -* || ${COMP_CWORD} -eq 1 ]] ; then
                 COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") )
                 return 0
@@ -49,6 +49,10 @@ _aichat() {
                     COMPREPLY=($(compgen -W "$("$1" --list-models)" -- "${cur}"))
                     return 0
                     ;;
+                --prompt)
+                    COMPREPLY=($(compgen -f "${cur}"))
+                    return 0
+                    ;;
                 -r|--role)
                     COMPREPLY=($(compgen -W "$("$1" --list-roles)" -- "${cur}"))
                     return 0
diff --git a/scripts/completions/aichat.fish b/scripts/completions/aichat.fish
index 47f3c6b..9ff6538 100644
--- a/scripts/completions/aichat.fish
+++ b/scripts/completions/aichat.fish
@@ -1,4 +1,5 @@
 complete -c aichat -s m -l model -x -a "(aichat --list-models)" -d 'Select a LLM model' -r
+complete -c aichat -l prompt -d 'Use the system prompt'
 complete -c aichat -s r -l role -x -a "(aichat --list-roles)" -d 'Select a role' -r
 complete -c aichat -s s -l session -x  -a"(aichat --list-sessions)" -d 'Start or join a session' -r
 complete -c aichat -s f -l file -d 'Include files with the message' -r -F
diff --git a/scripts/completions/aichat.nu b/scripts/completions/aichat.nu
index a8ac4bf..3b6ed9a 100644
--- a/scripts/completions/aichat.nu
+++ b/scripts/completions/aichat.nu
@@ -25,6 +25,7 @@ module completions {
   # All-in-one chat and copilot CLI that integrates 10+ AI platforms
   export extern aichat [
     --model(-m): string@"nu-complete aichat model"    # Select a LLM model
+    --prompt                                          # Use the system prompt
     --role(-r): string@"nu-complete aichat role"      # Select a role
     --session(-s): string@"nu-complete aichat role"   # Start or join a session
     --save-session                                    # Forces the session to be saved
diff --git a/scripts/completions/aichat.ps1 b/scripts/completions/aichat.ps1
index c3e8e96..10b986d 100644
--- a/scripts/completions/aichat.ps1
+++ b/scripts/completions/aichat.ps1
@@ -22,6 +22,7 @@ Register-ArgumentCompleter -Native -CommandName 'aichat' -ScriptBlock {
         'aichat' {
             [CompletionResult]::new('-m', '-m', [CompletionResultType]::ParameterName, 'Select a LLM model')
             [CompletionResult]::new('--model', '--model', [CompletionResultType]::ParameterName, 'Select a LLM model')
+            [CompletionResult]::new('--prompt', '--prompt', [CompletionResultType]::ParameterName, 'Use the system prompt')
             [CompletionResult]::new('-r', '-r', [CompletionResultType]::ParameterName, 'Select a role')
             [CompletionResult]::new('--role', '--role', [CompletionResultType]::ParameterName, 'Select a role')
             [CompletionResult]::new('-s', '-s', [CompletionResultType]::ParameterName, 'Start or join a session')
diff --git a/scripts/completions/aichat.zsh b/scripts/completions/aichat.zsh
index 378caf6..bbcc02b 100644
--- a/scripts/completions/aichat.zsh
+++ b/scripts/completions/aichat.zsh
@@ -16,6 +16,7 @@ _aichat() {
     local context curcontext="$curcontext" state line
      local common=(
 '-m+[Select a LLM model]:MODEL:->models' \
+'--prompt=[Use the system prompt]:PROMPT: ' \
 '--model=[Select a LLM model]:MODEL:->models' \
 '-r+[Select a role]:ROLE:->roles' \
 '--role=[Select a role]:ROLE:->roles' \
diff --git a/src/cli.rs b/src/cli.rs
index f6e12ab..8225b7f 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -6,6 +6,9 @@ pub struct Cli {
     /// Select a LLM model
     #[clap(short, long)]
     pub model: Option<String>,
+    /// Use the system prompt
+    #[clap(long)]
+    pub prompt: Option<String>,
     /// Select a role
     #[clap(short, long)]
     pub role: Option<String>,
diff --git a/src/main.rs b/src/main.rs
index f3d3846..f691b6f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -80,7 +80,9 @@ async fn main() -> Result<()> {
     if cli.dry_run {
         config.write().dry_run = true;
     }
-    if let Some(name) = &cli.role {
+    if let Some(prompt) = &cli.prompt {
+        config.write().set_prompt(prompt)?;
+    } else if let Some(name) = &cli.role {
         config.write().set_role(name)?;
     } else if cli.execute {
         config.write().set_role(SHELL_ROLE)?;
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index dd1c735..6306f0b 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -32,7 +32,7 @@ lazy_static! {
         ReplCommand::new(".model", "Change the current LLM", State::all()),
         ReplCommand::new(
             ".prompt",
-            "Make a temporary role using a prompt",
+            "Create a temporary role using a custom prompt",
             State::able_change_role()
         ),
         ReplCommand::new(