fix: fixed regression in 78e5f52

ErikBjare · Nov 18, 2023 · 7160b97 · 7160b97
1 parent be40f10
commit 7160b97
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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")
+"""
     )