From 3c98d5d3a403c1c01802d4ab1d1d802861f1737a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Mon, 9 Sep 2024 18:19:18 +0200
Subject: [PATCH] fixed tests

---
 gptme/codeblock.py       |  2 +-
 gptme/reduce.py          |  2 +-
 gptme/tools/base.py      |  6 ++----
 gptme/tools/patch.py     | 33 +--------------------------------
 tests/test_codeblock.py  | 10 +++++-----
 tests/test_logmanager.py |  3 ++-
 6 files changed, 12 insertions(+), 44 deletions(-)

diff --git a/gptme/codeblock.py b/gptme/codeblock.py
index bb82e916..6503143c 100644
--- a/gptme/codeblock.py
+++ b/gptme/codeblock.py
@@ -1,5 +1,5 @@
-from dataclasses import dataclass
 from collections.abc import Generator
+from dataclasses import dataclass
 from xml.etree import ElementTree
 
 
diff --git a/gptme/reduce.py b/gptme/reduce.py
index 14f71453..e1ae9692 100644
--- a/gptme/reduce.py
+++ b/gptme/reduce.py
@@ -8,7 +8,7 @@
 from collections.abc import Generator
 from copy import copy
 
-from ..codeblock import Codeblock
+from .codeblock import Codeblock
 from .message import Message, len_tokens
 from .models import get_model
 
diff --git a/gptme/tools/base.py b/gptme/tools/base.py
index 3fac1d92..cf85e063 100644
--- a/gptme/tools/base.py
+++ b/gptme/tools/base.py
@@ -38,9 +38,7 @@ def get_doc(self, _doc="") -> str:
             _doc += "\n\n"
         _doc += f"ToolSpec: {self.name}\n{self.desc}"
         if self.instructions:
-            _doc += f"\n# Instructions: {self.instructions}"
+            _doc += f"\n\n# Instructions\n\n{self.instructions}"
         if self.examples:
-            _doc += (
-                f"\n# Examples: {transform_examples_to_chat_directives(self.examples)}"
-            )
+            _doc += f"\n\n# Examples\n\n{transform_examples_to_chat_directives(self.examples)}"
         return _doc
diff --git a/gptme/tools/patch.py b/gptme/tools/patch.py
index c8995009..ced02b79 100644
--- a/gptme/tools/patch.py
+++ b/gptme/tools/patch.py
@@ -1,40 +1,8 @@
 """
 Gives the LLM agent the ability to patch text files, by using a adapted version git conflict markers.
-
-The format is suitable for small changes to files we have written in the past, without having to rewrite the whole file.
-
-
-Format:
-
-.. code-block:: patch
-
-    ```patch <filename>
-    <<<<<<< ORIGINAL
-    original code
-    =======
-    modified code
-    >>>>>>> UPDATED
-    ```
-
-Example:
-
-.. chat::
-
-    User: edit `hello.py` to ask for the name to greet
-    Assistant:
-    ```patch hello.py
-    <<<<<<< ORIGINAL
-    print("Hello world")
-    =======
-    name = input("What is your name? ")
-    print(f"Hello {name}")
-    >>>>>>> UPDATED
-    ```
-    System: Patch applied
 """
 
 # TODO: support multiple patches in one codeblock (or make it clear that only one patch per codeblock is supported/applied)
-
 import re
 from collections.abc import Generator
 from pathlib import Path
@@ -90,6 +58,7 @@ def patch_to_output(patch: str, filename: str, append: bool = False) -> str:
     print(f"Hello {name}")
 >>>>>>> UPDATED
 ''')}
+> System: Patch applied
 
 > User: change the codeblock to append to the file
 > Assistant: {patch_to_output("patch.py", '''
diff --git a/tests/test_codeblock.py b/tests/test_codeblock.py
index bcef387d..aba2f98e 100644
--- a/tests/test_codeblock.py
+++ b/tests/test_codeblock.py
@@ -11,7 +11,7 @@ def hello():
 More text
 """
     assert Codeblock.iter_from_markdown(markdown) == [
-        ("python", 'def hello():\n    print("Hello, World!")')
+        Codeblock("python", 'def hello():\n    print("Hello, World!")')
     ]
 
 
@@ -31,11 +31,11 @@ def greet(name):
 ```
 """
     assert Codeblock.iter_from_markdown(markdown) == [
-        (
+        Codeblock(
             "java",
             'public class Main {\n    public static void main(String[] args) {\n        System.out.println("Hello, Java!");\n    }\n}',
         ),
-        ("python", 'def greet(name):\n    return f"Hello, {name}!"'),
+        Codeblock("python", 'def greet(name):\n    return f"Hello, {name}!"'),
     ]
 
 
@@ -51,7 +51,7 @@ def print_readme():
 ```
 """
     assert Codeblock.iter_from_markdown(markdown) == [
-        (
+        Codeblock(
             "python",
             "def print_readme():\n    print('''Usage:\n```javascript\ncallme()\n```\n''')",
         )
@@ -77,5 +77,5 @@ def hello():
 ```
 """
     assert Codeblock.iter_from_markdown(markdown) == [
-        ("", 'def hello():\n    print("Hello, World!")')
+        Codeblock("", 'def hello():\n    print("Hello, World!")')
     ]
diff --git a/tests/test_logmanager.py b/tests/test_logmanager.py
index 9ad8789e..2b1f59a3 100644
--- a/tests/test_logmanager.py
+++ b/tests/test_logmanager.py
@@ -1,3 +1,4 @@
+from gptme.codeblock import Codeblock
 from gptme.logmanager import LogManager, Message
 
 
@@ -18,7 +19,7 @@ def test_get_last_codeblock():
 """,
         )
     )
-    assert ("python", "print('world')") == log.get_last_codeblock()
+    assert Codeblock("python", "print('world')") == log.get_last_codeblock()
 
 
 def test_branch():