recover for unfinished chat

2 files changed, 8 insertions, 0 deletions

diff --git a/py/chat.py b/py/chat.py
index 93a8664..6054484 100644
--- a/py/chat.py
+++ b/py/chat.py
@@ -27,6 +27,13 @@ def initialize_chat_window():
     vim_break_undo_sequence()
     vim.command("redraw")
 
+    file_content = vim.eval('trim(join(getline(1, "$"), "\n"))')
+    role_lines = re.findall(r'(^>>> user|^>>> system|^<<< assistant).*', file_content, flags=re.MULTILINE)
+    if not role_lines[-1].startswith(">>> user"):
+        # last role is not user, most likely completion was cancelled before
+        vim.command("normal! o")
+        vim.command("normal! i\n>>> user\n\n")
+
     if prompt:
         vim.command("normal! a" + prompt)
         vim_break_undo_sequence()
diff --git a/py/utils.py b/py/utils.py
index 9539d84..22d4079 100644
--- a/py/utils.py
+++ b/py/utils.py
@@ -5,6 +5,7 @@ import json
 import urllib.error
 import urllib.request
 import socket
+import re
 from urllib.error import URLError
 
 is_debugging = vim.eval("g:vim_ai_debug") == "1"