chat streaming, more py3 integration

4 files changed, 55 insertions, 88 deletions

diff --git a/plugin/vim-ai.vim b/plugin/vim-ai.vim
index a5ded7a..cf33596 100644
--- a/plugin/vim-ai.vim
+++ b/plugin/vim-ai.vim
@@ -80,47 +80,24 @@ endfunction
 
 function! AIChatRun(...) range
   let lines = trim(join(getline(a:firstline, a:lastline), "\n"))
-  let selection = trim(@*)
-  let is_selection = lines != "" && lines == selection
-  let has_instruction = a:0
+  let is_selection = lines != "" && lines == trim(@*)
+  let instruction = a:0 ? trim(a:1) : ""
 
-  let prompt = ""
-  if search('^>>> user$', 'nw') != 0
-    " inside chat window
-    let prompt = trim(join(getline(1, '$'), "\n"))
-  else
-    " outside chat window
+  if search('^>>> user$', 'nw') == 0
+    " outside of chat window
     call ScratchWindow()
-    if has_instruction
-      if is_selection
-        let prompt = a:1 . ":\n" . lines
-      else
-        let prompt = a:1
-      endif
-    else
-      if is_selection
-        let prompt = lines
-      else
-        execute "normal i>>> user\<Enter>\<Enter>"
-        return
-      endif
+    let delimiter = instruction != "" && is_selection ? ":\n" : ""
+    let selection = is_selection ? lines : ""
+    let prompt = join([instruction, delimiter, selection], "")
+    execute "normal i>>> user\<Enter>\<Enter>" . prompt
+    if prompt == ""
+      " empty prompt, just opens chat window TODO: handle in python
+      return
     endif
   endif
 
-  if g:vim_ai_debug
-    echo "Prompt:\n" . prompt . "\n"
-  endif
-
   echo "Answering..."
-  " WORKAROUND: without sleep is echo on prev line not displayed (when combining with py3)
-  execute 'silent sleep 1m'
   execute "py3file " . s:chat_py
-  let output = py3eval('output')
-
-  set paste
-  execute "normal! ggdG"
-  execute "normal! i" . output . "\<Esc>"
-  set nopaste
 endfunction
 
 command! -range -nargs=? AI <line1>,<line2>call AIRun(<f-args>)
diff --git a/py/chat.py b/py/chat.py
index b6dcd35..b6db8db 100644
--- a/py/chat.py
+++ b/py/chat.py
@@ -1,20 +1,12 @@
-import requests
-import sys
-import os
+import openai
 
-file_content = vim.eval("prompt")
+# import utils
+plugin_root = vim.eval("s:plugin_root")
+vim.command(f"py3file {plugin_root}/py/utils.py")
 
-config_file_path = os.path.join(os.path.expanduser("~"), ".config/openai.token")
+openai.api_key = load_api_key()
 
-api_key = os.getenv("OPENAI_API_KEY")
-
-try:
-    with open(config_file_path, 'r') as file:
-        api_key = file.read()
-except Exception:
-    pass
-
-api_key = api_key.strip()
+file_content = vim.eval('trim(join(getline(1, "$"), "\n"))')
 
 lines = file_content.splitlines()
 messages = []
@@ -37,19 +29,20 @@ if not messages:
     file_content = ">>> user\n\n" + file_content
     messages.append({"role": "user", "content": file_content })
 
+vim.command("normal! Go\n<<< assistant\n\n")
+vim.command("redraw")
 
-url = "https://api.openai.com/v1/chat/completions"
-headers = {
-    'Content-Type': 'application/json',
-    'Authorization': F"Bearer {api_key}"
-}
-data = {
-    "model": "gpt-3.5-turbo",
-    "messages": messages
-}
-response = requests.post(url, headers=headers, json=data)
-response = response.json()
+response = openai.ChatCompletion.create(
+  model="gpt-3.5-turbo",
+  messages=messages,
+  stream=True,
+)
 
-answer = response['choices'][0]['message']['content']
+for resp in response:
+    if 'content' in resp['choices'][0]['delta']:
+        text = resp['choices'][0]['delta']['content']
+        vim.command("normal! a" + text)
+        vim.command("redraw")
 
-output = f"{file_content.strip()}\n\n<<< assistant\n\n{answer.strip()}\n\n>>> user\n"
+vim.command("normal! a\n\n>>> user\n")
+vim.command("redraw")
diff --git a/py/complete.py b/py/complete.py
index 409d47f..c37cbd0 100644
--- a/py/complete.py
+++ b/py/complete.py
@@ -1,33 +1,18 @@
-import requests
-import sys
-import os
+import openai
 
-prompt = vim.eval("prompt")
-
-config_file_path = os.path.join(os.path.expanduser("~"), ".config/openai.token")
+# import utils
+plugin_root = vim.eval("s:plugin_root")
+vim.command(f"py3file {plugin_root}/py/utils.py")
 
-api_key = os.getenv("OPENAI_API_KEY")
-
-try:
-    with open(config_file_path, 'r') as file:
-        api_key = file.read()
-except Exception:
-    pass
+prompt = vim.eval("prompt")
 
-api_key = api_key.strip()
+openai.api_key = load_api_key()
 
-url = "https://api.openai.com/v1/completions"
-headers = {
-    'Content-Type': 'application/json',
-    'Authorization': f"Bearer {api_key}"
-}
-data = {
-    "model": "text-davinci-003",
-    "prompt":prompt,
-    "max_tokens": 1000,
-    "temperature": 0.1
-}
-response = requests.post(url, headers=headers, json=data)
-response = response.json()
+response = openai.Completion.create(
+    model="text-davinci-003",
+    prompt=prompt,
+    max_tokens=1000,
+    temperature=0.1
+)
 
 output = response['choices'][0]['text']
diff --git a/py/utils.py b/py/utils.py
new file mode 100644
index 0000000..c9e88f2
--- /dev/null
+++ b/py/utils.py
@@ -0,0 +1,12 @@
+import sys
+import os
+
+def load_api_key():
+    config_file_path = os.path.join(os.path.expanduser("~"), ".config/openai.token")
+    api_key = os.getenv("OPENAI_API_KEY")
+    try:
+        with open(config_file_path, 'r') as file:
+            api_key = file.read()
+    except Exception:
+        pass
+    return api_key.strip()