dev: code cleanup, maybe better exit behaviour

zhudotexe · Oct 3, 2024 · 7294719 · 7294719
1 parent 97eb175
commit 7294719
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/kani/prompts/impl/mistral.py b/kani/prompts/impl/mistral.py
@@ -261,7 +261,8 @@ async def stream(self, messages: list[ChatMessage], functions: list[AIFunction]
             content, tool_calls = self._parse_tool_calls(content)
             if inner_completion:
                 tool_calls = (inner_completion.message.tool_calls or []) + tool_calls
-            yield Completion(ChatMessage.assistant(content.removesuffix(self.eos_token).strip(), tool_calls=tool_calls))
+            clean_content = content.removesuffix(self.eos_token).strip()
+            yield Completion(ChatMessage.assistant(clean_content, tool_calls=tool_calls))
 
 
 MistralFunctionCallingAdapter = MixtralFunctionCallingAdapter
diff --git a/kani/utils/cli.py b/kani/utils/cli.py
@@ -73,8 +73,11 @@ async def chat_in_terminal_async(
                     # function
                     elif msg.role == ChatRole.FUNCTION and show_function_returns:
                         print_width(msg.text, width=width, prefix="FUNC: ")
-    except KeyboardInterrupt:
-        await kani.engine.close()
+    except (KeyboardInterrupt, asyncio.CancelledError):
+        # we won't close the engine here since it's common enough that people close the session in colab
+        # and if the process is closing then this will clean itself up anyway
+        # await kani.engine.close()
+        return
 
 
 @overload