diff options
Diffstat (limited to 'py/utils.py')
| -rw-r--r-- | py/utils.py | 41 |
1 files changed, 32 insertions, 9 deletions
diff --git a/py/utils.py b/py/utils.py index a0003a0..6fd161a 100644 --- a/py/utils.py +++ b/py/utils.py @@ -1,14 +1,9 @@ import datetime import sys import os - -try: - import openai -except ImportError: - raise Exception("OpenAI module not found. Please install it with pip.") - -is_debugging = vim.eval("g:vim_ai_debug") == "1" -debug_log_file = vim.eval("g:vim_ai_debug_log_file") +import json +import urllib.error +import urllib.request def load_api_key(): config_file_path = os.path.join(os.path.expanduser("~"), ".config/openai.token") @@ -22,12 +17,14 @@ def load_api_key(): raise Exception("Missing OpenAI API key") return api_key.strip() +is_debugging = vim.eval("g:vim_ai_debug") == "1" +debug_log_file = vim.eval("g:vim_ai_debug_log_file") + def make_request_options(options): request_options = {} request_options['model'] = options['model'] request_options['max_tokens'] = int(options['max_tokens']) request_options['temperature'] = float(options['temperature']) - request_options['request_timeout'] = float(options['request_timeout']) return request_options def render_text_chunks(chunks): @@ -71,3 +68,29 @@ def printDebug(text, *args): return with open(debug_log_file, "a") as file: file.write(f"[{datetime.datetime.now()}] " + text.format(*args) + "\n") + +OPENAI_RESP_DATA_PREFIX = 'data: ' +OPENAI_RESP_DONE = '[DONE]' +OPENAI_API_KEY = load_api_key() + +def openai_request(url, data): + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {OPENAI_API_KEY}" + } + req = urllib.request.Request( + url, + data=json.dumps(data).encode("utf-8"), + headers=headers, + method="POST", + ) + with urllib.request.urlopen(req) as response: + for line_bytes in response: + line = line_bytes.decode("utf-8", errors="replace") + if line.startswith(OPENAI_RESP_DATA_PREFIX): + line_data = line[len(OPENAI_RESP_DATA_PREFIX):-1] + if line_data == OPENAI_RESP_DONE: + pass + else: + openai_obj = json.loads(line_data) + yield openai_obj |