stream complete/edit commands

3 files changed, 57 insertions, 93 deletions

diff --git a/plugin/vim-ai.vim b/plugin/vim-ai.vim
index c183aa5..0f644b8 100644
--- a/plugin/vim-ai.vim
+++ b/plugin/vim-ai.vim
@@ -2,10 +2,6 @@ let s:plugin_root = expand('<sfile>:p:h:h')
 let s:complete_py = s:plugin_root . "/py/complete.py"
 let s:chat_py = s:plugin_root . "/py/chat.py"
 
-if !exists('g:vim_ai_debug')
-  let g:vim_ai_debug = 0
-endif
-
 function! ScratchWindow()
   below new
   setlocal buftype=nofile
@@ -14,90 +10,43 @@ function! ScratchWindow()
   setlocal ft=aichat
 endfunction
 
-function! AIRun(...) 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 prompt = ""
-  if has_instruction
-    if is_selection
-      let prompt = a:1 . ":\n" . lines
-    else
-      let prompt = a:1
-    endif
-  else
-    let prompt = lines
-  endif
+function! MakePrompt(lines, args)
+  let lines = trim(join(a:lines, "\n"))
+  let is_selection = lines != "" && lines == trim(@*)
+  let instruction = trim(get(a:args, 0))
+  let delimiter = instruction != "" && is_selection ? ":\n" : ""
+  let selection = is_selection ? lines : ""
+  let prompt = join([instruction, delimiter, selection], "")
+  return prompt
+endfunction
 
-  if g:vim_ai_debug
-    echo "Prompt:\n" . prompt . "\n"
-  endif
-  echo "Completing..."
-  " WORKAROUND: without sleep is echo on prev line not displayed (when combining with py3)
-  execute 'silent sleep 1m'
+function! AIRun(...) range
+  let prompt = MakePrompt(getline(a:firstline, a:lastline), a:000)
+  set paste
+  execute "normal! " . a:lastline . "Go"
   execute "py3file " . s:complete_py
-  let output = py3eval('output')
-  let output = trim(output)
-
   execute "normal! " . a:lastline . "G"
-  set paste
-  execute "normal! o" . output . "\<Esc>"
   set nopaste
-  execute "normal! " . a:lastline . "G"
 endfunction
 
 function! AIEditRun(...) range
-  let has_instruction = a:0
-  let prompt = trim(join(getline(a:firstline, a:lastline), "\n"))
-  if has_instruction
-    let prompt = a:1 . ":\n" . prompt
-  endif
-
-  let buff_lastline = line('$')
-
-  if g:vim_ai_debug
-    echo "Prompt:\n" . prompt . "\n"
-  endif
-
-  echo "Editing..."
-  " WORKAROUND: without sleep is echo on prev line not displayed (when combining with py3)
-  execute 'silent sleep 1m'
-  execute "py3file " . s:complete_py
-  let output = py3eval('output')
-  let output = trim(output)
-
-  execute a:firstline . ',' . a:lastline . 'd'
-
+  let prompt = MakePrompt(getline(a:firstline, a:lastline), a:000)
   set paste
-  if a:lastline == buff_lastline
-    execute "normal! o" . output . "\<Esc>"
-  else
-    execute "normal! O" . output . "\<Esc>"
-  endif
+  execute "normal! " . a:firstline . "GV" . a:lastline . "Gc"
+  execute "py3file " . s:complete_py
   set nopaste
 endfunction
 
 function! AIChatRun(...) range
-  let lines = trim(join(getline(a:firstline, a:lastline), "\n"))
-  let is_selection = lines != "" && lines == trim(@*)
-  let instruction = a:0 ? trim(a:1) : ""
-
+  let lines = getline(a:firstline, a:lastline)
   set paste
-  if search('^>>> user$', 'nw') == 0
-    " outside of chat window
+  let is_outside_of_chat_window = search('^>>> user$', 'nw') == 0
+  if is_outside_of_chat_window
     call ScratchWindow()
-    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
+    let prompt = MakePrompt(lines, a:000)
+    execute "normal i>>> user\n\n" . prompt
   endif
 
-  echo "Answering..."
   execute "py3file " . s:chat_py
   set nopaste
 endfunction
diff --git a/py/chat.py b/py/chat.py
index b6db8db..066227b 100644
--- a/py/chat.py
+++ b/py/chat.py
@@ -29,20 +29,24 @@ 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")
-
-response = openai.ChatCompletion.create(
-  model="gpt-3.5-turbo",
-  messages=messages,
-  stream=True,
-)
-
-for resp in response:
-    if 'content' in resp['choices'][0]['delta']:
-        text = resp['choices'][0]['delta']['content']
+if messages[-1]["content"].strip():
+
+    vim.command("normal! Go\n<<< assistant\n\n")
+    vim.command("redraw")
+
+    response = openai.ChatCompletion.create(
+      model="gpt-3.5-turbo",
+      messages=messages,
+      stream=True,
+    )
+
+    generating_text = False
+    for resp in response:
+        text = resp['choices'][0]['delta'].get('content', '')
+        if not text.strip() and not generating_text:
+            continue # trim newlines from the beginning
         vim.command("normal! a" + text)
         vim.command("redraw")
 
-vim.command("normal! a\n\n>>> user\n")
-vim.command("redraw")
+    vim.command("normal! a\n\n>>> user\n\n")
+    vim.command("redraw")
diff --git a/py/complete.py b/py/complete.py
index c37cbd0..705dd90 100644
--- a/py/complete.py
+++ b/py/complete.py
@@ -8,11 +8,22 @@ prompt = vim.eval("prompt")
 
 openai.api_key = load_api_key()
 
-response = openai.Completion.create(
-    model="text-davinci-003",
-    prompt=prompt,
-    max_tokens=1000,
-    temperature=0.1
-)
+if prompt.strip():
 
-output = response['choices'][0]['text']
+    response = openai.Completion.create(
+        model="text-davinci-003",
+        prompt=prompt,
+        max_tokens=1000,
+        temperature=0.1,
+        stream=True,
+    )
+
+    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")