diff options
Diffstat (limited to '')
| -rw-r--r-- | py/complete.py | 35 |
1 files changed, 21 insertions, 14 deletions
diff --git a/py/complete.py b/py/complete.py index bddd8e6..96ae032 100644 --- a/py/complete.py +++ b/py/complete.py @@ -4,28 +4,35 @@ import openai plugin_root = vim.eval("s:plugin_root") vim.command(f"py3file {plugin_root}/py/utils.py") -prompt = vim.eval("prompt").strip() +engine = vim.eval("engine") +options = make_options() request_options = make_request_options() +prompt = vim.eval("prompt").strip() + openai.api_key = load_api_key() +def complete_engine(): + response = openai.Completion.create(stream=True, prompt=prompt, **request_options) + text_chunks = map(lambda resp: resp['choices'][0].get('text', ''), response) + return text_chunks + +def chat_engine(): + initial_prompt = options.get('initial_prompt', '') + chat_content = f"{initial_prompt}\n\n>>> user\n\n{prompt}".strip() + messages = parse_chat_messages(chat_content) + response = openai.ChatCompletion.create(messages=messages, stream=True, **request_options) + text_chunks = map(lambda resp: resp['choices'][0]['delta'].get('content', ''), response) + return text_chunks + +engines = {"chat": chat_engine, "complete": complete_engine} + try: if prompt: - print('Completing...') vim.command("redraw") - - response = openai.Completion.create(stream=True, prompt=prompt, **request_options) - - generating_text = False - for resp in response: - text = resp['choices'][0].get('text', '') - if not text.strip() and not generating_text: - continue # trim newlines from the beginning - - generating_text = True - vim.command("normal! a" + text) - vim.command("redraw") + text_chunks = engines[engine]() + render_text_chunks(text_chunks) except KeyboardInterrupt: vim.command("normal! a Ctrl-C...") except openai.error.Timeout: |