ctrl c to cancel completion

2 files changed, 48 insertions, 37 deletions

diff --git a/py/chat.py b/py/chat.py
index 066227b..551bfb8 100644
--- a/py/chat.py
+++ b/py/chat.py
@@ -29,24 +29,29 @@ if not messages:
     file_content = ">>> user\n\n" + file_content
     messages.append({"role": "user", "content": file_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)
+try:
+    if messages[-1]["content"].strip():
+        vim.command("normal! Go\n<<< assistant\n\n")
         vim.command("redraw")
 
-    vim.command("normal! a\n\n>>> user\n\n")
-    vim.command("redraw")
+        print('Answering...')
+        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\n")
+        vim.command("redraw")
+except KeyboardInterrupt:
+    vim.command("normal! a Ctrl-C...")
diff --git a/py/complete.py b/py/complete.py
index 705dd90..9e46d63 100644
--- a/py/complete.py
+++ b/py/complete.py
@@ -8,22 +8,28 @@ prompt = vim.eval("prompt")
 
 openai.api_key = load_api_key()
 
-if prompt.strip():
-
-    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)
+try:
+    if prompt.strip():
+
+        print('Completing...')
         vim.command("redraw")
+
+        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")
+except KeyboardInterrupt:
+    vim.command("normal! a Ctrl-C...")