feat: added patch tool, added more integration test examples

Author: ErikBjare

gptme/commands.py

@@ -164,6 +164,7 @@ def handle_cmd(
             sys.exit(0)
         case "replay":
             log.undo(1, quiet=True)
+            log.write()
             print("Replaying conversation...")
             for msg in log.log:
                 if msg.role == "assistant":

gptme/prompts.py

@@ -7,6 +7,7 @@
 
 from .config import get_config
 from .message import Message
+from .tools import patch
 
 USER = os.environ.get("USER", None)
 
@@ -160,29 +161,9 @@ def initial_prompt(short: bool = False) -> Generator[Message, None, None]:
 python hello.py
 ```
 > stdout: `Hello world!`
-
-## Producing and applying patches
-Use `diff` and `patch` to produce and apply patches.
-
-> User: add a name input to hello.py
-First we create the patch:
-```hello.patch
-@@ -1,1 +1,2 @@
--print("Hello world!")
-+name = input("What is your name? ")
-+print(f"Hello {name}!")
-```
-
-Now, we apply it:
-```bash
-patch hello.py < hello.patch
-```
-
-Now, we can run it:
-```bash
-echo "John" | python hello.py
-```
-""".strip(),
+""".strip()
+            + "\n\n"
+            + patch.instructions,
             hide=True,
             pinned=True,
         )

gptme/tools/__init__.py

@@ -2,6 +2,7 @@
 from typing import Generator
 
 from ..message import Message
+from .patch import execute_patch
 from .python import execute_python, init_python
 from .save import execute_save
 from .shell import execute_shell
@@ -32,16 +33,19 @@ def execute_msg(msg: Message, ask: bool) -> Generator[Message, None, None]:
 def execute_codeblock(codeblock: str, ask: bool) -> Generator[Message, None, None]:
     """Executes a codeblock and returns the output."""
     lang_or_fn = codeblock.splitlines()[0].strip()
-    codeblock = codeblock[len(lang_or_fn) :]
+    codeblock_content = codeblock[len(lang_or_fn) :]
 
     is_filename = lang_or_fn.count(".") >= 1
 
     if lang_or_fn in ["python", "py"]:
-        yield from execute_python(codeblock, ask=ask)
+        yield from execute_python(codeblock_content, ask=ask)
     elif lang_or_fn in ["terminal", "bash", "sh"]:
-        yield from execute_shell(codeblock, ask=ask)
+        yield from execute_shell(codeblock_content, ask=ask)
+    elif lang_or_fn.startswith("patch "):
+        fn = lang_or_fn[len("patch ") :]
+        yield from execute_patch(f"```{codeblock}```", fn, ask=ask)
     elif is_filename:
-        yield from execute_save(lang_or_fn, codeblock, ask=ask)
+        yield from execute_save(lang_or_fn, codeblock_content, ask=ask)
     else:
         logger.warning(
             f"Unknown codeblock type '{lang_or_fn}', neither supported language or filename."

gptme/tools/patch.py

@@ -0,0 +1,135 @@
+"""
+Gives the LLM agent the ability to patch files, by using a adapted version git conflict markers.
+
+Inspired by aider.
+"""
+
+import re
+from pathlib import Path
+from typing import Generator
+
+from ..message import Message
+from ..util import ask_execute
+
+example_patch = """
+```patch filename.py
+<<<<<<< ORIGINAL
+original lines
+=======
+modified lines
+>>>>>>> UPDATED
+```
+"""
+
+instructions = """
+# Patching files
+
+The LLM agent can patch files, by using a adapted version git conflict markers.
+This can be used to make changes to files we have written in the past, without having to rewrite the whole file.
+
+## Example
+
+> User: Patch the file `hello.py` to ask for the name of the user.
+```hello.py
+def hello():
+    print("hello world")
+```
+
+> Assistant:
+```patch hello.py
+<<<<<<< ORIGINAL
+    print("hello world")
+=======
+    name = input("What is your name? ")
+    print(f"hello {name}")
+>>>>>>> UPDATED
+```
+"""
+
+
+def apply(codeblock: str, content: str) -> str:
+    """
+    Applies the patch to the file.
+    """
+    codeblock = codeblock.strip()
+
+    # get the original chunk
+    original = re.split("\n<<<<<<< ORIGINAL\n", codeblock)[1]
+    original = re.split("\n=======\n", original)[0]
+
+    # get the modified chunk
+    modified = re.split("\n=======\n", codeblock)[1]
+    modified = re.split("\n>>>>>>> UPDATED\n", modified)[0]
+
+    # replace the original chunk with the modified chunk
+    content = content.replace(original, modified)
+
+    return content
+
+
+def apply_file(codeblock, filename):
+    codeblock = codeblock.strip()
+    _patch, filename = codeblock.splitlines()[0].split()
+    assert _patch == "```patch"
+    assert Path(filename).exists()
+
+    with open(filename, "r") as f:
+        content = f.read()
+
+    result = apply(codeblock, content)
+
+    with open(filename, "w") as f:
+        f.write(result)
+
+    print(f"Applied patch to {filename}")
+
+
+def execute_patch(codeblock: str, fn: str, ask: bool) -> Generator[Message, None, None]:
+    """
+    Executes the patch.
+    """
+    if ask:
+        confirm = ask_execute()
+        if not confirm:
+            print("Patch not applied")
+            return
+
+    apply_file(codeblock, fn)
+    yield Message("system", "Patch applied")
+
+
+def test_apply_simple():
+    codeblock = example_patch
+    content = """original lines"""
+    result = apply(codeblock, content)
+    assert result == """modified lines"""
+
+
+def test_apply_function():
+    content = """
+def hello():
+    print("hello")
+
+if __name__ == "__main__":
+    hello()
+"""
+
+    codeblock = """
+```patch test.py
+<<<<<<< ORIGINAL
+def hello():
+    print("hello")
+=======
+def hello(name="world"):
+    print(f"hello {name}")
+>>>>>>> UPDATED
+```
+"""
+
+    result = apply(codeblock, content)
+    assert result.startswith(
+        """
+def hello(name="world"):
+    print(f"hello {name}")
+"""
+    )

tests/test-integration.sh

@@ -8,23 +8,46 @@ cd "$(dirname "$0")"
 mkdir -p output
 cd output
 
+# run interactive tests if not in CI (GITHUB_ACTIONS is set by github actions)
+interactive=${GITHUB_ACTIONS:-1}
+
 # set this to indicate tests are run (non-interactive)
 export PYTEST_CURRENT_TEST=1
 
 # test stdin and cli-provided prompt
 echo "The project mascot is a flying pig" | gptme "What is the project mascot?"
 
+# test load context from file
+echo "The project mascot is a flying pig" > mascot.txt
+gptme "What is the project mascot?" mascot.txt
+
 # test python command
 gptme "/python print('hello world')"
 
 # test shell command
 gptme "/shell echo 'hello world'"
 
-# interactive matplotlib
-gptme 'plot an x^2 graph'
+# test write small game
+gptme 'write a snake game with curses to snake.py'
+# works!
+
+# test implement game of life
+gptme 'write an implementation of the game of life with curses to life.py'
+# works? almost, needed to try-catch-pass an exception
+
+# test implement wireworld
+gptme 'write a implementation of wireworld with curses to wireworld.py'
+# works? almost, needed to try-catch-pass an exception, fix color setup, and build a proper circuit
+
+# test plot to file
+gptme 'plot up to the 5rd degree taylor expansion of sin(x), save to sin.png'
+# works!
 
-# matplotlib to file
-gptme 'plot up to the 3rd degree taylor expansion of sin(x), save to sin.png'
+# write C code and apply patch
+gptme 'write a hello world program in c to hello.c, then patch it to ask for your name and print it'
+# works!
 
-# interactive curses
-gptme 'write a snake game with curses'
+if [ "$interactive" = "1" ]; then
+    # interactive matplotlib
+    gptme 'plot an x^2 graph'
+fi