From 7160b974643bfe3946018228116fc33d444d3208 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Sat, 18 Nov 2023 11:35:40 +0100
Subject: [PATCH] fix: fixed regression in 78e5f52

---
 gptme/tools/__init__.py | 2 +-
 gptme/tools/save.py     | 2 +-
 tests/test_util.py      | 6 ++++--
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/gptme/tools/__init__.py b/gptme/tools/__init__.py
index 71634a730..95907152f 100644
--- a/gptme/tools/__init__.py
+++ b/gptme/tools/__init__.py
@@ -86,7 +86,7 @@ def is_supported_codeblock(codeblock: str) -> bool:
 def get_codeblocks(content: str) -> Generator[str, None, None]:
     """Returns all codeblocks in a message."""
     for codeblock in ("\n" + content).split("\n```")[1::2]:
-        yield codeblock
+        yield codeblock + "\n"
 
 
 def init_tools() -> None:
diff --git a/gptme/tools/save.py b/gptme/tools/save.py
index 59a03d26c..cb95701cb 100644
--- a/gptme/tools/save.py
+++ b/gptme/tools/save.py
@@ -1,5 +1,5 @@
-from pathlib import Path
 from collections.abc import Generator
+from pathlib import Path
 
 from ..message import Message
 from ..util import ask_execute
diff --git a/tests/test_util.py b/tests/test_util.py
index 9851510ea..554bfe8e0 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -51,7 +51,8 @@ def test_get_codeblocks():
     assert (
         codeblocks[0]
         == """python
-print("hello world")"""
+print("hello world")
+"""
     )
 
     # test a codeblock which contains triple backticks
@@ -65,5 +66,6 @@ def test_get_codeblocks():
     assert (
         codeblocks[0]
         == """python
-print("hello ``` world")"""
+print("hello ``` world")
+"""
     )