feat: added append tool, fixed patch tool

gptme/tools/__init__.py

@@ -44,6 +44,9 @@ def execute_codeblock(codeblock: str, ask: bool) -> Generator[Message, None, Non
     elif lang_or_fn.startswith("patch "):
         fn = lang_or_fn[len("patch ") :]
         yield from execute_patch(f"```{codeblock}```", fn, ask=ask)
+    elif lang_or_fn.startswith("append "):
+        fn = lang_or_fn[len("append ") :]
+        yield from execute_save(fn, codeblock_content, ask=ask, append=True)
     elif is_filename:
         yield from execute_save(lang_or_fn, codeblock_content, ask=ask)
     else:
@@ -78,6 +81,8 @@ def is_supported_codeblock(codeblock: str) -> bool:
         return True
     elif lang_or_fn.startswith("patch "):
         return True
+    elif lang_or_fn.startswith("append "):
+        return True
     elif is_filename:
         return True
     else:

gptme/tools/patch.py

@@ -44,18 +44,20 @@ def apply(codeblock: str, content: str) -> str:
     # TODO: support multiple patches in one file, or make it clear that this is not supported (one patch per codeblock)
     codeblock = codeblock.strip()
 
-    # get the original chunk
+    # get the original and modified chunks
     original = re.split("\n<<<<<<< ORIGINAL\n", codeblock)[1]
-    original = re.split("\n=======\n", original)[0]
+    original, modified = re.split("\n=======\n", original)
+    assert ">>>>>>> UPDATED\n" in modified, f"Invalid patch: {codeblock}"
+    modified = re.split(">>>>>>> UPDATED\n", modified)[0].rstrip("\n")
 
-    # get the modified chunk
-    modified = re.split("\n=======\n", codeblock)[1]
-    modified = re.split("\n>>>>>>> UPDATED\n", modified)[0]
+    # the modified chunk may contain "// ..." to refer to chunks in the original
+    # replace these with the original chunks
 
     # replace the original chunk with the modified chunk
-    content = content.replace(original, modified)
+    new = content.replace(original, modified)
+    assert new != content, "Patch did not change the file"
 
-    return content
+    return new
 
 
 def apply_file(codeblock, filename):

gptme/tools/save.py

@@ -5,25 +5,38 @@
 from ..util import ask_execute
 
 
-def execute_save(fn: str, code: str, ask: bool) -> Generator[Message, None, None]:
+def execute_save(
+    fn: str, code: str, ask: bool, append: bool = False
+) -> Generator[Message, None, None]:
     """Save the code to a file."""
+    action = "save" if not append else "append"
     # strip leading newlines
     code = code.lstrip("\n")
 
     if ask:
-        confirm = ask_execute(f"Save to {fn}?")
+        confirm = ask_execute(f"{action.capitalize()} to {fn}?")
         print()
     else:
         confirm = True
-        print("Skipping save confirmation.")
+        print(f"Skipping {action} confirmation.")
 
     if ask and not confirm:
         # early return
-        yield Message("system", "Save cancelled.")
+        yield Message("system", f"{action.capitalize()} cancelled.")
         return
 
     path = Path(fn).expanduser()
 
+    if append:
+        if not path.exists():
+            yield Message("system", f"File {fn} doesn't exist, can't append to it.")
+            return
+
+        with open(path, "a") as f:
+            f.write("\n" + code)
+        yield Message("system", f"Appended to {fn}")
+        return
+
     # if the file exists, ask to overwrite
     if path.exists():
         overwrite = ask_execute("File exists, overwrite?")

tests/test_tools_patch.py

@@ -46,3 +46,21 @@ def hello(name="world"):
     print(f"hello {name}")
 """
     )
+
+    # only remove code in patch
+    codeblock = """
+```patch test.py
+<<<<<<< ORIGINAL
+def hello():
+    print("hello")
+=======
+>>>>>>> UPDATED
+```
+"""
+    print(content)
+    result = apply(codeblock, content)
+    newline = "\n"
+    newline_escape = "\\n"
+    assert result.startswith(
+        "\n\n"
+    ), f"result: {result.replace(newline, newline_escape)}"