chat engine

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: