handle roles in python

2 files changed, 26 insertions, 27 deletions

diff --git a/plugin/vim-ai.vim b/plugin/vim-ai.vim
index 6192f1d..dba5998 100644
--- a/plugin/vim-ai.vim
+++ b/plugin/vim-ai.vim
@@ -116,14 +116,7 @@ function! AIChatRun(is_selection, ...) range
       let instruction = a:0 ? a:1 : ""
       let prompt = MakePrompt(a:is_selection, lines, instruction)
     endif
-    let is_outside_of_chat_window = search('^>>> user$', 'nw') == 0
-    if is_outside_of_chat_window
-      " write an user prompt
-      execute "normal! i>>> user\n\n" . prompt
-    else
-      " appending prompt into restored conversation
-      execute "normal! Gi" . prompt
-    endif
+    execute "normal! Gi" . prompt
   endif
 
   let s:last_command = "chat"
diff --git a/py/chat.py b/py/chat.py
index f8345b4..6b5ed67 100644
--- a/py/chat.py
+++ b/py/chat.py
@@ -5,30 +5,36 @@ plugin_root = vim.eval("s:plugin_root")
 vim.command(f"py3file {plugin_root}/py/utils.py")
 
 options = make_options()
-file_content = vim.eval('trim(join(getline(1, "$"), "\n"))')
 
 openai.api_key = load_api_key()
 
-lines = file_content.splitlines()
-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
+def parse_messages():
+    file_content = vim.eval('trim(join(getline(1, "$"), "\n"))')
+    lines = file_content.splitlines()
+    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
+    return messages
+
+messages = parse_messages()
 
 if not messages:
-    file_content = ">>> user\n\n" + file_content
-    messages.append({"role": "user", "content": file_content })
+    # roles not found, put whole file content as an user prompt
+    vim.command("normal! ggO>>> user\n")
+    vim.command("normal! G")
+    vim.command("redraw")
+    messages = parse_messages()
 
 for message in messages:
     # strip newlines from the content as it causes empty responses