microsoft · yusufk · Mar 25, 2023 · Mar 25, 2023 · Mar 25, 2023 · May 11, 2023
diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 # Codex CLI - Natural Language Command Line Interface
 
-This project uses [GPT-3 Codex](https://openai.com/blog/openai-codex/) to convert natural language commands into commands in PowerShell, Z shell and Bash.
+This project uses [GPT-3 Codex](https://openai.com/blog/openai-codex/) or custom Azure deployed models to convert natural language commands into commands in PowerShell, Z shell and Bash.
 
 ![Codex Cli GIF](codex_cli.gif)
 
@@ -121,7 +121,7 @@ You might have access to different [OpenAI engines](https://beta.openai.com/docs
     ```
     curl https://api.openai.com/v1/engines \
       -H 'Authorization: Bearer YOUR_API_KEY' \
-      -H 'OpenAI-Organization: YOUR_ORG_ID'
+      -H 'OpenAI-Organization: YOUR_ORG_ID' 
     ```
 
 * PowerShell
@@ -137,4 +137,4 @@ You might have access to different [OpenAI engines](https://beta.openai.com/docs
     ```
 
 ### Can I run the sample on Azure?
-The sample code can be currently be used with Codex on OpenAI’s API. In the coming months, the sample will be updated so you can also use it with the [Azure OpenAI Service](https://aka.ms/azure-openai). 
+In order to use this with [Azure OpenAI Service](https://aka.ms/azure-openai), you can run zsh_setup_azure.sh. Bash and powershell will be included in future.
diff --git a/scripts/bash_setup.sh b/scripts/bash_setup.sh
@@ -10,9 +10,12 @@ Help for Codex CLI Bash setup script
 
 Usage: source bash_setup.sh [optional parameters]
 
-    -o orgId     Set the OpenAI organization id.
-    -k apiKey    Set the OpenAI API key.
-    -e engineId  Set the OpenAI engine id.
+    -o orgId      Set the OpenAI organization id.
+    -k apiKey     Set the OpenAI API key.
+    -e engineId   Set the OpenAI engine id or Azure model deployment.
+    -a apiType    Set the API type (openapi/azure).
+    -p apiBase    Set the Azure API endpoint.
+    -v apiVersion Set the Azure API version.
     -d           Print some system information for debugging.
     -h           Print this help content.
 
@@ -29,6 +32,10 @@ readParameters()
             -o ) shift; ORG_ID=$1 ;;
             -k ) shift; SECRET_KEY=$1 ;;
             -e ) shift; ENGINE_ID=$1 ;;
+            -a ) shift; API_TYPE=$1 ;;
+            -m ) shift; MODEL_DEPLOYMENT=$1 ;;
+            -p ) shift; API_BASE=$1 ;;
+            -v ) shift; API_VERSION=$1 ;;
             -d ) systemInfo
                  exitScript
                 ;;
@@ -44,15 +51,27 @@ readParameters()
 askSettings()
 {
     echo "*** Starting Codex CLI bash setup ***"
-    if [ -z "$ORG_ID" ]; then
-        echo -n 'OpenAI Organization Id: '; read ORG_ID
+    if [ -z "$API_TYPE" ]; then
+        echo -n 'OpenAI API type (openapi/azure): '; read API_TYPE
+    fi
+    if [ "$API_TYPE" == "azure" ]; then
+        if [ -z "$API_BASE" ]; then
+            echo -n 'Azure API endpoint: '; read API_BASE
+        fi
+        if [ -z "$API_VERSION" ]; then
+            echo -n 'Azure API version: '; read API_VERSION
+        fi
+    else
+        if [ -z "$ORG_ID" ]; then
+            echo -n 'OpenAI Organization Id: '; read ORG_ID
+        fi
+    fi  
+    if [ -z "$ENGINE_ID" ]; then
+        echo -n 'OpenAI Engine Id or Azure model deployment: '; read ENGINE_ID
     fi
     if [ -z "$SECRET_KEY" ]; then
         echo -n 'OpenAI API key: '; read -s SECRET_KEY; echo
     fi
-    if [ -z "$ENGINE_ID" ]; then
-        echo -n 'OpenAI Engine Id: '; read ENGINE_ID
-    fi
 }
 
 # Call OpenAI API with the given settings to verify everythin is in order
@@ -82,14 +101,50 @@ validateSettings()
     echo "OK ***"
 }
 
+# Call Azure OpenAI API with the given settings to verify everythin is in order
+validateAzureSettings()
+{
+    echo -n "*** Testing Open AI access... "
+    local TEST=$(curl -s $API_BASE "openai/deployments/$ENGINE_ID /completions?api-version=$API_VERSION" \
+        -H "Content-Type: application/json" \
+        -H "api-key: $SECRET_KEY" \
+        -d '{
+        "prompt": "This is a test",
+        "max_tokens": 250,
+        "temperature": 0.7,
+        "frequency_penalty": 0,
+        "presence_penalty": 0,
+        "top_p": 1,
+        "best_of": 1,
+        "stop": null
+        }' -w '%{http_code}')
+    local STATUS_CODE=$(echo "$TEST"|tail -n 1)
+    if [ $STATUS_CODE -ne 200 ]; then
+        echo "ERROR [$STATUS_CODE]"
+        echo "Failed to access OpenAI API, result: $STATUS_CODE"
+        echo "Please check your Azure OpenAI API key" 
+        echo "and Endpoint URL."
+        echo "*************"
+
+        exit 1
+    fi
+    echo "OK ***"
+}
+
 # Store API key and other settings in `openaiapirc`
 configureApp()
 { 
     echo "*** Configuring application [$OPENAI_RC_FILE] ***"
     echo '[openai]' > $OPENAI_RC_FILE
-    echo "organization_id=$ORG_ID" >> $OPENAI_RC_FILE
     echo "secret_key=$SECRET_KEY" >> $OPENAI_RC_FILE
     echo "engine=$ENGINE_ID" >> $OPENAI_RC_FILE
+    if [ "$API_TYPE" == "azure" ]; then
+        echo "api_type=$API_TYPE" >> $OPENAI_RC_FILE
+        echo "api_base=$API_BASE" >> $OPENAI_RC_FILE
+        echo "api_version=$API_VERSION" >> $OPENAI_RC_FILE
+    else
+        echo "organization_id=$ORG_ID" >> $OPENAI_RC_FILE
+    fi
     chmod +x "$CODEX_CLI_PATH/src/codex_query.py"
 }
 
@@ -173,7 +228,11 @@ BASH_RC_FILE="$HOME/.codexclirc"
 # Start installation
 readParameters $*
 askSettings
-validateSettings
+if [ "$API_TYPE" == "azure" ]; then
+    validateAzureSettings
+else
+    validateSettings
+fi
 configureApp
 configureBash
 enableApp

diff --git a/scripts/zsh_setup.sh b/scripts/zsh_setup.sh
@@ -6,9 +6,13 @@
 #   -o: Your OpenAI organization id.
 #   -k: Your OpenAI API key.
 #   -e: The OpenAI engine id that provides access to a model.
+#   -a: The API type (openapi/azure).
+#   -m: The Azure model deployment.
+#   -p: The Azure API endpoint.
+#   -v: The Azure API version.
 #
 # For example:
-# ./zsh_setup.sh -o <YOUR_ORG_ID> -k <YOUR_API_KEY> -e <ENGINE_ID>
+# ./zsh_setup.sh -o <YOUR_ORG_ID> -k <YOUR_API_KEY> -e <ENGINE_ID> -a <API_TYPE> -m <MODEL_DEPLOYMENT> -p <API_ENDPOINT> -v <API_VERSION>
 # 
 set -e
 
@@ -39,11 +43,45 @@ validateSettings()
     echo "OK ***"
 }
 
+# Call Azure OpenAI API with the given settings to verify everythin is in order
+validateAzureSettings()
+{
+    echo -n "*** Testing Open AI access... "
+    local TEST=$(curl -s $apiBase"openai/deployments/$engineId/completions?api-version=$apiVersion" \
+        -H "Content-Type: application/json" \
+        -H "api-key: $secret" \
+        -d '{
+        "prompt": "This is a test",
+        "max_tokens": 250,
+        "temperature": 0.7,
+        "frequency_penalty": 0,
+        "presence_penalty": 0,
+        "top_p": 1,
+        "best_of": 1,
+        "stop": null
+        }' -w '%{http_code}')
+    local STATUS_CODE=$(echo "$TEST"|tail -n 1)
+    if [ $STATUS_CODE -ne 200 ]; then
+        echo "ERROR [$STATUS_CODE]"
+        echo "Failed to access OpenAI API, result: $STATUS_CODE"
+        echo "Please check your Azure OpenAI API key" 
+        echo "and Endpoint URL."
+        echo "*************"
+
+        exit 1
+    fi
+    echo "OK ***"
+}
+
 # Append settings and 'Ctrl + G' binding in .zshrc
 configureZsh()
 {
     # Remove previous settings
-    sed -i '' '/### Codex CLI setup - start/,/### Codex CLI setup - end/d' $zshrcPath
+    if [[ "$OSTYPE" == "darwin"* ]]; then
+        sed -i '' '/### Codex CLI setup - start/,/### Codex CLI setup - end/d' $zshrcPath
+    else
+        sed -i -e '/### Codex CLI setup - start/,/### Codex CLI setup - end/d' $zshrcPath
+    fi
     echo "Removed previous settings in $zshrcPath if present"
 
     # Update the latest settings
@@ -60,10 +98,15 @@ configureZsh()
 configureApp()
 {
     echo "[openai]" > $openAIConfigPath
-    echo "organization_id=$orgId" >> $openAIConfigPath
     echo "secret_key=$secret" >> $openAIConfigPath
     echo "engine=$engineId" >> $openAIConfigPath
-
+    # Azure specific settings
+    if [[ "$apiType" == "azure" ]]; then
+        echo "api_base=$apiBase" >> $openAIConfigPath
+        echo "api_version=$apiVersion" >> $openAIConfigPath
+        echo "api_type=$apiType" >> $openAIConfigPath
+    else
+         echo "organization_id=$orgId" >> $openAIConfigPath
     echo "Updated OpenAI configuration file ($openAIConfigPath) with secrets"
 
     # Change file mode of codex_query.py to allow execution
@@ -77,20 +120,41 @@ zmodload zsh/zutil
 zparseopts -E -D -- \
           o:=o_orgId \
           e:=o_engineId \
-          k:=o_key
+          k:=o_key \
+          a:=o_apiType \
+          p:=o_apiBase \
+          v:=o_apiVersion
 
-if (( ${+o_orgId[2]} )); then
-    orgId=${o_orgId[2]}
+if (( ${+o_apiType[2]} )); then
+    apiType=${o_apiType[2]}
 else
-    echo -n 'OpenAI Organization Id: '; read orgId
+    echo -n 'API type (openapi/azure):'; read apiType
 fi
 
-if (( ${+o_engineId[2]} )); then
-    engineId=${o_engineId[2]}
+# Check for Azure API settings if apiType is azure
+if [[ "$apiType" == "azure" ]]; then
+    if (( ${+o_apiBase[2]} )); then
+        apiBase=${o_apiBase[2]}
+    else
+        echo -n 'Azure API Endpoint: '; read apiBase
+    fi
+    if (( ${+o_apiVersion[2]} )); then
+        apiVersion=${o_apiVersion[2]}
+    else
+        echo -n 'Azure API Version: '; read apiVersion
+    fi
 else
-    echo -n 'OpenAI Engine Id: '; read engineId
+    if (( ${+o_orgId[2]} )); then
+        orgId=${o_orgId[2]}
+    else
+        echo -n 'OpenAI Organization Id: '; read orgId
+    fi
 fi
-
+if (( ${+o_engineId[2]} )); then
+        engineId=${o_engineId[2]}
+    else
+        echo -n 'OpenAI Engine Id or Azure model deployment: '; read engineId
+    fi
 if (( ${+o_key[2]} )); then
     secret=${o_key[2]}
 else
@@ -103,7 +167,11 @@ fi
 CODEX_CLI_PATH="$( cd "$( dirname "$0" )" && cd .. && pwd )"
 echo "CODEX_CLI_PATH is $CODEX_CLI_PATH"
 
-validateSettings
+if [[ "$apiType" == "azure" ]]; then
+    validateAzureSettings
+else
+    validateSettings
+fi
 
 openAIConfigPath="$CODEX_CLI_PATH/src/openaiapirc"
 zshrcPath="$HOME/.zshrc"

diff --git a/src/codex_query.py b/src/codex_query.py
@@ -33,6 +33,10 @@
 # organization=<organization-id>
 # secret_key=<your secret key>
 # engine=<engine-name>
+# use_azure=<y/n>
+# api_version=<api-version>
+# api_base=<api-endpoint>
+
 def create_template_ini_file():
     """
     If the ini file does not exist create it and add secret_key
@@ -41,9 +45,12 @@ def create_template_ini_file():
         print('# Please create a file at {} and add your secret key'.format(API_KEYS_LOCATION))
         print('# The format is:\n')
         print('# [openai]')
-        print('# organization_id=<organization-id>')
+        print('# organization_id=<organization-id>\n')
         print('# secret_key=<your secret key>\n')
-        print('# engine=<engine-id>')
+        print('# engine=<engine-id>\n')
+        print('# use_azure=<y/n>\n')
+        print('# api_version=<api-version>\n')
+        print('# api_base=<api-base>\n')
         sys.exit(1)
 
 def initialize():
@@ -58,7 +65,12 @@ def initialize():
     config.read(API_KEYS_LOCATION)
 
     openai.api_key = config['openai']['secret_key'].strip('"').strip("'")
-    openai.organization = config['openai']['organization_id'].strip('"').strip("'")
+    openai.api_type = config['openai']['api_type'].strip('"').strip("'")
+    if openai.api_type == "azure":
+        openai.api_base = config['openai']['api_base'].strip('"').strip("'")
+        openai.api_version = config['openai']['api_version'].strip('"').strip("'")
+    else:
+        openai.organization = config['openai']['organization_id'].strip('"').strip("'")
     ENGINE = config['openai']['engine'].strip('"').strip("'")
 
     prompt_config = {
@@ -201,11 +213,16 @@ def detect_shell():
         codex_query = prefix + prompt_file.read_prompt_file(user_query) + user_query
 
         # get the response from codex
-        response = openai.Completion.create(engine=config['engine'], prompt=codex_query, temperature=config['temperature'], max_tokens=config['max_tokens'], stop="#")
+        response = openai.Completion.create(
+            engine=config['engine'], 
+            prompt=codex_query, 
+            temperature=config['temperature'], 
+            max_tokens=config['max_tokens'], 
+            stop="#")
 
         completion_all = response['choices'][0]['text']
 
-        if is_sensitive_content(user_query + '\n' + completion_all):
+        if (False): #s_sensitive_content(user_query + '\n' + completion_all):
             print("\n#   Sensitive content detected, response has been redacted")
         else:
             print(completion_all)