Merge pull request #8 from dahromy/main

Add option to select AI provider for zsh-copilot

README.md

@@ -22,18 +22,34 @@ echo "source ~/.zsh-copilot/zsh-copilot.plugin.zsh" >> ~/.zshrc
 
 ### Configuration
 
-You need to have an OPENAI API key with access to `gpt-4` to use this plugin. Expose this via the `OPENAI_API_KEY` environment variable:
+You need to have an API key for either OpenAI or Anthropic to use this plugin. Expose this via the appropriate environment variable:
 
+For OpenAI (default):
 ```sh
-export OPENAI_API_KEY=<your-api-key>
+export OPENAI_API_KEY=<your-openai-api-key>
 ```
 
-I tried out using `gpt-3` but the results were garbage.
+For Anthropic:
+```sh
+export ANTHROPIC_API_KEY=<your-anthropic-api-key>
+```
 
-To see available configurations, run:
+You can configure the AI provider using the `ZSH_COPILOT_AI_PROVIDER` variable:
 
 ```sh
-zsh-copilot --help
+export ZSH_COPILOT_AI_PROVIDER="openai"  # or "anthropic"
+```
+
+Other configuration options:
+
+- `ZSH_COPILOT_KEY`: Key to press to get suggestions (default: ^z)
+- `ZSH_COPILOT_SEND_CONTEXT`: If `true`, zsh-copilot will send context information to the AI model (default: true)
+- `ZSH_COPILOT_DEBUG`: Enable debug logging (default: false)
+
+To see all available configurations and their current values, run:
+
+```sh
+zsh-copilot
 ```
 
 ## Usage

zsh-copilot.plugin.zsh

@@ -1,15 +1,19 @@
+# Default key binding
 (( ! ${+ZSH_COPILOT_KEY} )) &&
     typeset -g ZSH_COPILOT_KEY='^z'
 
+# Configuration options
 (( ! ${+ZSH_COPILOT_SEND_CONTEXT} )) &&
     typeset -g ZSH_COPILOT_SEND_CONTEXT=true
 
-# (( ! ${+ZSH_COPILOT_SEND_GIT_DIFF} )) &&
-#     typeset -g ZSH_COPILOT_SEND_GIT_DIFF=true
-
 (( ! ${+ZSH_COPILOT_DEBUG} )) &&
     typeset -g ZSH_COPILOT_DEBUG=false
 
+# New option to select AI provider
+(( ! ${+ZSH_COPILOT_AI_PROVIDER} )) &&
+    typeset -g ZSH_COPILOT_AI_PROVIDER="openai"
+
+# System prompt
 read -r -d '' SYSTEM_PROMPT <<- EOM
   You will be given the raw input of a shell command. 
   Your task is to either complete the command or provide a new command that you think the user is trying to type. 
@@ -29,8 +33,8 @@ read -r -d '' SYSTEM_PROMPT <<- EOM
   DO NOT INTERACT WITH THE USER IN NATURAL LANGUAGE! If you do, you will be banned from the system. 
   Note that the double quote sign is escaped. Keep this in mind when you create quotes. 
   Here are two examples: 
-    * User input: 'list files in current directory'; Your response: '=ls # ls is the builtin command for listing files' 
+    * User input: 'list files in current directory'; Your response: '=ls' (ls is the builtin command for listing files)
-    * User input: 'cd /tm'; Your response: '+p # /tmp is the standard temp folder on linux and mac'.
+    * User input: 'cd /tm'; Your response: '+p' (/tmp is the standard temp folder on linux and mac).
 EOM
 
 if [[ "$OSTYPE" == "darwin"* ]]; then
@@ -41,13 +45,15 @@ fi
 
 function _suggest_ai() {
     local OPENAI_API_URL=${OPENAI_API_URL:-"api.openai.com"}
+    local ANTHROPIC_API_URL=${ANTHROPIC_API_URL:-"api.anthropic.com"}
 
+    local context_info=""
     if [[ "$ZSH_COPILOT_SEND_CONTEXT" == 'true' ]]; then
-        local PROMPT="$SYSTEM_PROMPT 
+        context_info="Context: You are user $(whoami) with id $(id) in directory $(pwd). 
-            Context: You are user $(whoami) with id $(id) in directory $(pwd). 
             Your shell is $(echo $SHELL) and your terminal is $(echo $TERM) running on $(uname -a).
             $SYSTEM"
     fi
+
     # Get input
     local input=$(echo "${BUFFER:0:$CURSOR}" | tr '\n' ';')
     input=$(echo "$input" | sed 's/"/\\"/g')
@@ -55,26 +61,18 @@ function _suggest_ai() {
     _zsh_autosuggest_clear
     zle -R "Thinking..."
 
-    PROMPT=$(echo "$PROMPT" | tr -d '\n')
+    local full_prompt=$(echo "$SYSTEM_PROMPT $context_info" | tr -d '\n')
-    # Wasn't able to get this to work :(
-    # if [[ "$ZSH_COPILOT_SEND_GIT_DIFF" == 'true' ]]; then
-    #     if [[ $(git rev-parse --is-inside-work-tree) == 'true' ]]; then
-    #         local git_diff=$(git diff --staged --no-color)
-    #         local git_exit_code=$?
-    #         git_diff=$(echo "$git_diff" | tr '\\' ' ' | sed 's/[\$\"\`]/\\&/g' | tr '\\' '\\\\' | tr -d '\n')
-    #
-    #         if [[ git_exit_code -eq 0 ]]; then
-    #             PROMPT="$PROMPT; This is the git diff: <---->$git_diff<----> You may provide a git commit message if the user is trying to commit changes. You are an expert at committing changes, you don't give generic messages. You give the best commit messages"
-    #         fi
-    #     fi
-    # fi
 
-    local data="{
+    local data
-            \"model\": \"gpt-4\",
+    local response
+
+    if [[ "$ZSH_COPILOT_AI_PROVIDER" == "openai" ]]; then
+        data="{
+            \"model\": \"gpt-4o-mini\",
             \"messages\": [
                 {
                     \"role\": \"system\",
-                    \"content\": \"$PROMPT\"
+                    \"content\": \"$full_prompt\"
                 },
                 {
                     \"role\": \"user\",
@@ -82,14 +80,35 @@ function _suggest_ai() {
                 }
             ]
         }"
-    local response=$(curl "https://${OPENAI_API_URL}/v1/chat/completions" \
+        response=$(curl "https://${OPENAI_API_URL}/v1/chat/completions" \
             --silent \
             -H "Content-Type: application/json" \
             -H "Authorization: Bearer $OPENAI_API_KEY" \
-        -d $data)
+            -d "$data")
         local message=$(echo "$response" | jq -r '.choices[0].message.content')
-
+    elif [[ "$ZSH_COPILOT_AI_PROVIDER" == "anthropic" ]]; then
-    # zle -U "$suggestion"
+        data="{
+            \"model\": \"claude-3-5-sonnet-20240620\",
+            \"max_tokens\": 1000,
+            \"system\": \"$full_prompt\",
+            \"messages\": [
+                {
+                    \"role\": \"user\",
+                    \"content\": \"$input\"
+                }
+            ]
+        }"
+        response=$(curl "https://${ANTHROPIC_API_URL}/v1/messages" \
+            --silent \
+            -H "Content-Type: application/json" \
+            -H "x-api-key: $ANTHROPIC_API_KEY" \
+            -H "anthropic-version: 2023-06-01" \
+            -d "$data")
+        local message=$(echo "$response" | jq -r '.content[0].text')
+    else
+        echo "Invalid AI provider selected. Please choose 'openai' or 'anthropic'."
+        return 1
+    fi
 
     local first_char=${message:0:1}
     local suggestion=${message:1:${#message}}
@@ -107,7 +126,6 @@ function _suggest_ai() {
         zle -U "$suggestion"
     elif [[ "$first_char" == '+' ]]; then
         _zsh_autosuggest_suggest "$suggestion"
-         # POSTDISPLAY="$suggestion"
     fi
 }
 
@@ -117,9 +135,8 @@ function zsh-copilot() {
     echo "Configurations:"
     echo "    - ZSH_COPILOT_KEY: Key to press to get suggestions (default: ^z, value: $ZSH_COPILOT_KEY)."
     echo "    - ZSH_COPILOT_SEND_CONTEXT: If \`true\`, zsh-copilot will send context information (whoami, shell, pwd, etc.) to the AI model (default: true, value: $ZSH_COPILOT_SEND_CONTEXT)."
-    # echo "    - ZSH_COPILOT_SEND_GIT_DIFF: If \`true\`, zsh-copilot will send the git diff (if available) to the AI model (default: true, value: $ZSH_COPILOT_SEND_GIT_DIFF)."
+    echo "    - ZSH_COPILOT_AI_PROVIDER: AI provider to use ('openai' or 'anthropic', default: openai, value: $ZSH_COPILOT_AI_PROVIDER)."
 }
 
 zle -N _suggest_ai
 bindkey $ZSH_COPILOT_KEY _suggest_ai
-