diff options
Diffstat (limited to '')
| -rw-r--r-- | py/utils.py | 25 |
1 files changed, 20 insertions, 5 deletions
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, |