Merge pull request #54 from duylam/openai-org-support

Support including OpenAI Org ID in the request to OpenAI API endpoints

2 files changed, 26 insertions, 5 deletions

diff --git a/README.md b/README.md
index 1ff7973..326d70c 100644
--- a/README.md
+++ b/README.md
@@ -32,8 +32,14 @@ You only share and pay for what you specifically select, for prompts and chat co
 # save api key to `~/.config/openai.token` file
 echo "YOUR_OPENAI_API_KEY" > ~/.config/openai.token
 
+# or save api key and org id to `~/.config/openai.token` file
+#echo "YOUR_OPENAI_API_KEY,YOUR_OPENAI_ORG_ID" > ~/.config/openai.token
+
 # alternatively set it as an environment variable
 export OPENAI_API_KEY="YOUR_OPENAI_API_KEY"
+
+# or set with org id
+#export OPENAI_API_KEY="YOUR_OPENAI_API_KEY,YOUR_OPENAI_ORG_ID"
 ```
 
 ### Using `vim-plug`
diff --git a/py/utils.py b/py/utils.py
index 76ae1e4..2e1f975 100644
--- a/py/utils.py
+++ b/py/utils.py
@@ -16,15 +16,26 @@ debug_log_file = vim.eval("g:vim_ai_debug_log_file")
 
 def load_api_key():
     config_file_path = os.path.join(os.path.expanduser("~"), ".config/openai.token")
-    api_key = os.getenv("OPENAI_API_KEY")
+    api_key_param_value = os.getenv("OPENAI_API_KEY")
     try:
         with open(config_file_path, 'r') as file:
-            api_key = file.read()
+            api_key_param_value = file.read()
     except Exception:
         pass
-    if not api_key:
+
+    if not api_key_param_value:
         raise Exception("Missing OpenAI API key")
-    return api_key.strip()
+
+    # The text is in format of "<api key>,<org id>" and the
+    # <org id> part is optional
+    elements = api_key_param_value.strip().split(",")
+    api_key = elements[0].strip()
+    org_id = None
+
+    if len(elements) > 1:
+        org_id = elements[1].strip()
+
+    return (api_key, org_id)
 
 def normalize_config(config):
     normalized = { **config }
@@ -119,13 +130,17 @@ def printDebug(text, *args):
 
 OPENAI_RESP_DATA_PREFIX = 'data: '
 OPENAI_RESP_DONE = '[DONE]'
-OPENAI_API_KEY = load_api_key()
+(OPENAI_API_KEY, OPENAI_ORG_ID) = load_api_key()
 
 def openai_request(url, data, options):
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {OPENAI_API_KEY}"
     }
+
+    if OPENAI_ORG_ID is not None:
+        headers["OpenAI-Organization"] =  f"{OPENAI_ORG_ID}"
+
     request_timeout=options['request_timeout']
     req = urllib.request.Request(
         url,