Add Azure Open AI support

scripts/bash_setup.sh

@@ -15,6 +15,7 @@ Usage: source bash_setup.sh [optional parameters]
     -e engineId  Set the OpenAI engine id.
     -d           Print some system information for debugging.
     -h           Print this help content.
+    -a           Use Azure Open AI
 
 To uninstall Codex CLI use bash_cleanup.sh.
 For more information visit https://github.com/microsoft/Codex-CLI
@@ -29,6 +30,7 @@ readParameters()
             -o ) shift; ORG_ID=$1 ;;
             -k ) shift; SECRET_KEY=$1 ;;
             -e ) shift; ENGINE_ID=$1 ;;
+            -a ) shift; USE_AZURE=$1 ;;
             -d ) systemInfo
                  exitScript
                 ;;
@@ -55,29 +57,53 @@ askSettings()
     fi
 }
 
-# Call OpenAI API with the given settings to verify everythin is in order
+# Call (Azure) OpenAI API with the given settings to verify everythin is in order
 validateSettings()
 {
-    echo -n "*** Testing Open AI access... "
-    local TEST=$(curl -s 'https://api.openai.com/v1/engines' -H "Authorization: Bearer $SECRET_KEY" -H "OpenAI-Organization: $ORG_ID" -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 OpenAI API key (https://beta.openai.com/account/api-keys)" 
-        echo "and Organization ID (https://beta.openai.com/account/org-settings)."
-        echo "*************"
-        exitScript
-        return
-    fi
-    local ENGINE_FOUND=$(echo "$TEST"|grep '"id"'|grep "\"$ENGINE_ID\"")
-    if [ -z "$ENGINE_FOUND" ]; then
-        echo "ERROR"
-        echo "Cannot find OpenAI engine: $ENGINE_ID" 
-        echo "Please check the OpenAI engine id (https://beta.openai.com/docs/engines/codex-series-private-beta)."
-        echo "*************"
-        exitScript
-        return
+       if [ -n "USE_AZURE" ]; then
+        echo -n "*** Testing Azure Open AI access... "
+        URL="${ORG_ID}openai/deployments?api-version=${USE_AZURE}"
+        local TEST=$(curl -s $URL -H "api-key: $SECRET_KEY" -w '%{http_code}')
+        local STATUS_CODE=$(echo "$TEST"|tail -n 1 | sed s'/}//g')
+        if [ $STATUS_CODE -ne 200 ]; then
+            echo "ERROR [$STATUS_CODE]"
+            echo "Failed to access Azure OpenAI API, result: $STATUS_CODE"
+            echo "Please check your Azure OpenAI Endpoint and API key (https://portal.azure.com)" 
+            echo "*************"
+            exitScript
+            return
+        fi
+        local DEPLOYMENT_FOUND=$(echo "$TEST"|grep '"id"'|grep "\"$ENGINE_ID\"")
+        if [ -z "$DEPLOYMENT_FOUND" ]; then
+            echo "ERROR"
+            echo "Cannot find Azure OpenAI deployment engine: $ENGINE_ID" 
+            echo "Please check the Azure OpenAI deployment engine id (hhttps://portal.azure.com)."
+            echo "*************"
+            exitScript
+            return
+        fi
+    else
+        echo -n "*** Testing Open AI access... "
+        local TEST=$(curl -s 'https://api.openai.com/v1/engines' -H "Authorization: Bearer $SECRET_KEY" -H "OpenAI-Organization: $ORG_ID" -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 OpenAI API key (https://beta.openai.com/account/api-keys)" 
+            echo "and Organization ID (https://beta.openai.com/account/org-settings)."
+            echo "*************"
+            exitScript
+            return
+        fi
+        local ENGINE_FOUND=$(echo "$TEST"|grep '"id"'|grep "\"$ENGINE_ID\"")
+        if [ -z "$ENGINE_FOUND" ]; then
+            echo "ERROR"
+            echo "Cannot find OpenAI engine: $ENGINE_ID" 
+            echo "Please check the OpenAI engine id (https://beta.openai.com/docs/engines/codex-series-private-beta)."
+            echo "*************"
+            exitScript
+            return
+        fi
     fi
     echo "OK ***"
 }
@@ -90,6 +116,9 @@ configureApp()
     echo "organization_id=$ORG_ID" >> $OPENAI_RC_FILE
     echo "secret_key=$SECRET_KEY" >> $OPENAI_RC_FILE
     echo "engine=$ENGINE_ID" >> $OPENAI_RC_FILE
+    if [ -n "$USE_AZURE" ]; then
+        echo "use_azure=$USE_AZURE" >> $OPENAI_RC_FILE
+    fi
     chmod +x "$CODEX_CLI_PATH/src/codex_query.py"
 }
 
@@ -149,7 +178,7 @@ systemInfo()
 # Remove variables and functions from the environment, in case the script was sourced
 cleanupEnv()
 {
-    unset ORG_ID SECRET_KEY ENGINE_ID SOURCED OPENAI_RC_FILE BASH_RC_FILE
+    unset ORG_ID SECRET_KEY ENGINE_ID USE_AZURE SOURCED OPENAI_RC_FILE BASH_RC_FILE
     unset -f askSettings validateSettings configureApp configureBash enableApp readParameters
 }
 

src/codex_query.py

@@ -57,10 +57,18 @@ def initialize():
     config = configparser.ConfigParser()
     config.read(API_KEYS_LOCATION)
 
+    """
+    Check if Azure Open AI is to be used
+    If so get the version and set base URL based on the organization
+    """
+    if 'use_azure' in config['openai']:
+        openai.api_type = "azure"
+        openai.api_base = config['openai']['organization_id'].strip('"').strip("'")
+        openai.api_version = config['openai']['use_azure'].strip('"').strip("'")
+    else:
+        openai.organization = config['openai']['organization_id'].strip('"').strip("'")
     openai.api_key = config['openai']['secret_key'].strip('"').strip("'")
-    openai.organization = config['openai']['organization_id'].strip('"').strip("'")
     ENGINE = config['openai']['engine'].strip('"').strip("'")
-
     prompt_config = {
         'engine': ENGINE,
         'temperature': TEMPERATURE,
@@ -143,6 +151,7 @@ def get_query(prompt_file):
         entry = input("prompt: ") + '\n'
     else:
         entry = sys.stdin.read()
+
     # first we check if the input is a command
     command_result, prompt_file = get_command_result(entry, prompt_file)
 
@@ -171,10 +180,8 @@ def detect_shell():
 if __name__ == '__main__':
     detect_shell()
     prompt_file = initialize()
-
     try:
         user_query, prompt_file = get_query(prompt_file)
-
         config = prompt_file.config if prompt_file else {
             'engine': ENGINE,
             'temperature': TEMPERATURE,
@@ -199,13 +206,13 @@ def detect_shell():
             prefix = '#' + config['shell'] + '\n\n'
 
         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="#")
-
+       
         completion_all = response['choices'][0]['text']
 
-        if is_sensitive_content(user_query + '\n' + completion_all):
+        if openai.api_type == "" and is_sensitive_content(user_query + '\n' + completion_all):
             print("\n#   Sensitive content detected, response has been redacted")
         else:
             print(completion_all)