adding edit and chat commands

3 files changed, 118 insertions, 4 deletions

diff --git a/doc/tags b/doc/tags
new file mode 100644
index 0000000..426fc7a
--- /dev/null
+++ b/doc/tags
@@ -0,0 +1,5 @@
+:AI	vim-ai.txt	/*:AI*
+vim-ai	vim-ai.txt	/*vim-ai*
+vim-ai-about	vim-ai.txt	/*vim-ai-about*
+vim-ai-commands	vim-ai.txt	/*vim-ai-commands*
+vim-ai.txt	vim-ai.txt	/*vim-ai.txt*
diff --git a/plugin/vim-ai.vim b/plugin/vim-ai.vim
index cd0cd6f..6fb0832 100644
--- a/plugin/vim-ai.vim
+++ b/plugin/vim-ai.vim
@@ -1,19 +1,49 @@
 let s:plugin_root = expand('<sfile>:p:h:h')
-let s:complete_script_file = s:plugin_root . "/py/complete.py"
+let s:complete_py = s:plugin_root . "/py/complete.py"
+let s:chat_py = s:plugin_root . "/py/chat.py"
+
+function! ScratchWindow()
+  below new
+  setlocal buftype=nofile
+  setlocal bufhidden=hide
+  setlocal noswapfile
+endfunction
 
 function! AIRun(...) range
   let prompt = getline(a:firstline, a:lastline)
   if a:0
-    " join arguments and prepend to the prompt
     let instruction = join(a:000, ", ") . ":"
     call insert(prompt, instruction, 0)
   endif
 
-  let buff_lastline = line('$')
   let prompt = join(prompt, "\n")
 
   echo "Completing..."
-  let output = system("echo " . shellescape(prompt) . " | python3 " . s:complete_script_file . " ")
+  let output = system("echo " . shellescape(prompt) . " | python3 " . s:complete_py . " ")
+  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
+  if !a:0
+    echo "Missing edit prompt instruction"
+    return
+  endif
+
+  let prompt = getline(a:firstline, a:lastline)
+  let instruction = join(a:000, ", ") . ":"
+  call insert(prompt, instruction, 0)
+
+  let buff_lastline = line('$')
+  let prompt = join(prompt, "\n")
+
+  echo "Editing..."
+  let output = system("echo " . shellescape(prompt) . " | python3 " . s:complete_py . " ")
   let output = trim(output)
 
   execute a:firstline . ',' . a:lastline . 'd'
@@ -25,7 +55,36 @@ function! AIRun(...) range
     execute "normal! O" . output . "\<Esc>"
   endif
   set nopaste
+endfunction
 
+function! AIChatRun(...)
+  let prompt = []
+  if search('^>>> user$', 'nw') != 0
+    " inside chat window
+    let prompt = getline(1, '$')
+  else
+    " outside chat window
+    call ScratchWindow()
+    if !a:0
+      execute "normal i>>> user\<Enter>\<Enter>"
+      return
+    endif
+    let instruction = join(a:000, ", ")
+    call insert(prompt, instruction, 0)
+  endif
+
+  let prompt = join(prompt, "\n")
+
+  echo "Answering..."
+  let output = system("echo " . shellescape(prompt) . " | python3 " . s:chat_py . " ")
+
+  set paste
+  execute "normal! ggdG"
+  execute "normal! i" . output . "\<Esc>"
+  set nopaste
 endfunction
 
 command! -range -nargs=? AI <line1>,<line2>call AIRun(<f-args>)
+command! -range -nargs=? AIComplete <line1>,<line2>call AIRun(<f-args>)
+command! -range -nargs=? AIEdit <line1>,<line2>call AIEditRun(<f-args>)
+command! -range -nargs=? AIChat <line1>,<line2>call AIChatRun(<f-args>)
diff --git a/py/chat.py b/py/chat.py
new file mode 100644
index 0000000..4c3d532
--- /dev/null
+++ b/py/chat.py
@@ -0,0 +1,50 @@
+import openai
+
+import sys
+import os
+import openai
+
+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
+
+openai.api_key = api_key.strip()
+
+lines = sys.stdin.readlines()
+
+file_content = "".join(lines)
+
+messages = []
+
+for line in lines:
+    if line.startswith(">>> system"):
+        messages.append({"role": "system", "content": ""})
+        continue
+    if line.startswith(">>> user"):
+        messages.append({"role": "user", "content": ""})
+        continue
+    if line.startswith("<<< assistant"):
+        messages.append({"role": "assistant", "content": ""})
+        continue
+    if not messages:
+        continue
+    messages[-1]["content"] += "\n" + line
+
+if not messages:
+    file_content = ">>> user\n\n" + file_content
+    messages.append({"role": "user", "content": file_content })
+
+response = openai.ChatCompletion.create(
+  model="gpt-3.5-turbo",
+  messages=messages
+)
+
+answer = response['choices'][0]['message']['content']
+
+print(f"{file_content.strip()}\n\n<<< assistant\n\n{answer.strip()}\n\n>>> user\n")