Add support for the config file

python/plugin.py

@@ -1,5 +1,8 @@
 import urllib, urllib.request
 import json
+import os
+import configparser
+import sys
 
 try:
   import vim
@@ -9,14 +12,46 @@
 
 
 import openai
-from AUTH import *
 
-openai.organization = ORGANIZATION_ID
-openai.api_key = SECRET_KEY
 MAX_SUPPORTED_INPUT_LENGTH = 4096
 USE_STREAM_FEATURE = True
 MAX_TOKENS_DEFAULT = 64
 
+CONFIG_DIR = os.getenv('XDG_CONFIG_HOME', os.path.expanduser('~/.config'))
+API_KEYS_LOCATION = os.path.join(CONFIG_DIR, 'openaiapirc')
+
+def create_template_ini_file():
+    """
+    If the ini file does not exist create it and add the organization_id and
+    secret_key
+    """
+    if not os.path.isfile(API_KEYS_LOCATION):
+        with open(API_KEYS_LOCATION, 'w') as f:
+            f.write('[openai]\n')
+            f.write('organization_id=\n')
+            f.write('secret_key=\n')
+
+        print('OpenAI API config file created at {}'.format(API_KEYS_LOCATION))
+        print('Please edit it and add your organization ID and secret key')
+        print('If you do not yet have an organization ID and secret key, you\n'
+               'need to register for OpenAI Codex: \n'
+                'https://openai.com/blog/openai-codex/')
+        sys.exit(1)
+
+
+def initialize_openai_api():
+    """
+    Initialize the OpenAI API
+    """
+    # Check if file at API_KEYS_LOCATION exists
+    create_template_ini_file()
+    config = configparser.ConfigParser()
+    config.read(API_KEYS_LOCATION)
+
+    openai.organization_id = config['openai']['organization_id'].strip('"').strip("'")
+    openai.api_key = config['openai']['secret_key'].strip('"').strip("'")
+
+
 def complete_input_max_length(input_prompt, max_input_length=MAX_SUPPORTED_INPUT_LENGTH, stop=None, max_tokens=64):
     input_prompt = input_prompt[-max_input_length:]
     response = openai.Completion.create(engine='davinci-codex', prompt=input_prompt, best_of=1, temperature=0.5, max_tokens=max_tokens, stream=USE_STREAM_FEATURE, stop=stop)
@@ -85,6 +120,13 @@ def get_first_line_below_cursor_with_text():
 
 
 def create_completion(stop=None): 
+    try:
+        from AUTH import ORGANIZATION_ID, SECRET_KEY
+
+        openai.organization = ORGANIZATION_ID
+        openai.api_key = SECRET_KEY
+    except ModuleNotFoundError:
+        initialize_openai_api()
     max_tokens = get_max_tokens()
     vim_buf = vim.current.buffer
     input_prompt = '\n'.join(vim_buf[:])