fix: improved how ToolUse examples are formatted

gptme/__init__.py

@@ -2,5 +2,6 @@
 from .logmanager import LogManager
 from .message import Message
 from .prompts import get_prompt
+from .codeblock import Codeblock
 
-__all__ = ["main", "chat", "LogManager", "Message", "get_prompt"]
+__all__ = ["main", "chat", "LogManager", "Message", "get_prompt", "Codeblock"]

gptme/tools/__init__.py

@@ -1,11 +1,9 @@
 import logging
 from collections.abc import Callable, Generator
-from dataclasses import dataclass
-from xml.etree import ElementTree
 
 from ..codeblock import Codeblock
 from ..message import Message
-from .base import ToolSpec
+from .base import ToolSpec, ToolUse
 from .browser import tool as browser_tool
 from .chats import tool as chats_tool
 from .gh import tool as gh_tool
@@ -45,64 +43,6 @@
 loaded_tools: list[ToolSpec] = []
 
 
-@dataclass
-class ToolUse:
-    tool: str
-    args: list[str]
-    content: str
-
-    def execute(self, ask: bool) -> Generator[Message, None, None]:
-        """Executes a tool-use tag and returns the output."""
-        tool = get_tool(self.tool)
-        if tool.execute:
-            yield from tool.execute(self.content, ask, self.args)
-
-    @classmethod
-    def from_codeblock(cls, codeblock: Codeblock) -> "ToolUse":
-        """Parses a codeblock into a ToolUse. Codeblock must be a supported type.
-
-        Example:
-          ```lang
-          content
-          ```
-        """
-        if tool := get_tool_for_langtag(codeblock.lang):
-            # NOTE: special case
-            args = (
-                codeblock.lang.split(" ")[1:]
-                if tool.name != "save"
-                else [codeblock.lang]
-            )
-            return ToolUse(tool.name, args, codeblock.content)
-        else:
-            raise ValueError(
-                f"Unknown codeblock type '{codeblock.lang}', neither supported language or filename."
-            )
-
-    @classmethod
-    def iter_from_xml(cls, content: str) -> Generator["ToolUse", None, None]:
-        """Returns all ToolUse in a message.
-
-        Example:
-          <tool-use>
-          <python>
-          print("Hello, world!")
-          </python>
-          </tool-use>
-        """
-        if "<tool-use>" not in content:
-            return
-
-        # TODO: this requires a strict format, should be more lenient
-        root = ElementTree.fromstring(content)
-        for tooluse in root.findall("tool-use"):
-            for child in tooluse:
-                # TODO: this child.attrib.values() thing wont really work
-                yield ToolUse(
-                    tooluse.tag, list(child.attrib.values()), child.text or ""
-                )
-
-
 def init_tools() -> None:
     """Runs initialization logic for tools."""
     for tool in all_tools:

gptme/tools/base.py

@@ -1,18 +1,27 @@
 from collections.abc import Callable, Generator
 from dataclasses import dataclass, field
-from typing import Any, Protocol, TypeAlias
+from typing import (
+    Any,
+    Literal,
+    Protocol,
+    TypeAlias,
+)
+from xml.etree import ElementTree
 
+from ..codeblock import Codeblock
 from ..message import Message
 from ..util import transform_examples_to_chat_directives
 
 InitFunc: TypeAlias = Callable[[], Any]
 
+# tooluse format mode
+mode: Literal["markdown", "xml"] = "markdown"
+
 
 class ExecuteFunc(Protocol):
     def __call__(
         self, code: str, ask: bool, args: list[str]
-    ) -> Generator[Message, None, None]:
-        ...
+    ) -> Generator[Message, None, None]: ...
 
 
 @dataclass
@@ -42,3 +51,78 @@ def get_doc(self, doc="") -> str:
                 f"# Examples\n\n{transform_examples_to_chat_directives(self.examples)}"
             )
         return doc
+
+
+@dataclass
+class ToolUse:
+    tool: str
+    args: list[str]
+    content: str
+
+    def execute(self, ask: bool) -> Generator[Message, None, None]:
+        """Executes a tool-use tag and returns the output."""
+        # noreorder
+        from . import get_tool  # fmt: skip
+        tool = get_tool(self.tool)
+        if tool.execute:
+            yield from tool.execute(self.content, ask, self.args)
+
+    @classmethod
+    def from_codeblock(cls, codeblock: Codeblock) -> "ToolUse":
+        """Parses a codeblock into a ToolUse. Codeblock must be a supported type.
+
+        Example:
+          ```lang
+          content
+          ```
+        """
+        # noreorder
+        from . import get_tool_for_langtag  # fmt: skip
+        if tool := get_tool_for_langtag(codeblock.lang):
+            # NOTE: special case
+            args = (
+                codeblock.lang.split(" ")[1:]
+                if tool.name != "save"
+                else [codeblock.lang]
+            )
+            return ToolUse(tool.name, args, codeblock.content)
+        else:
+            raise ValueError(
+                f"Unknown codeblock type '{codeblock.lang}', neither supported language or filename."
+            )
+
+    @classmethod
+    def iter_from_xml(cls, content: str) -> Generator["ToolUse", None, None]:
+        """Returns all ToolUse in a message.
+
+        Example:
+          <tool-use>
+          <python>
+          print("Hello, world!")
+          </python>
+          </tool-use>
+        """
+        if "<tool-use>" not in content:
+            return
+
+        # TODO: this requires a strict format, should be more lenient
+        root = ElementTree.fromstring(content)
+        for tooluse in root.findall("tool-use"):
+            for child in tooluse:
+                # TODO: this child.attrib.values() thing wont really work
+                yield ToolUse(
+                    tooluse.tag, list(child.attrib.values()), child.text or ""
+                )
+
+    def to_output(self) -> str:
+        if mode == "markdown":
+            return self.to_markdown()
+        elif mode == "xml":
+            return self.to_xml()
+
+    def to_markdown(self) -> str:
+        return f"```{self.tool} {"".join(self.args)}\n{self.content}\n```"
+
+    def to_xml(self) -> str:
+        args = " ".join(self.args)
+        return f"<{self.tool} args='{args}'>\n{self.content}\n</{self.tool}>"

gptme/tools/patch.py

@@ -6,79 +6,42 @@
 import re
 from collections.abc import Generator
 from pathlib import Path
-from typing import Literal
 
 from ..message import Message
 from ..util import ask_execute
-from .base import ToolSpec
+from .base import ToolSpec, ToolUse
 
 instructions = """
 To patch/modify files, we can use an adapted version of git conflict markers.
 
 This can be used to make changes to files, without having to rewrite the whole file.
 Only one patch block can be written per codeblock. Extra ORIGINAL/UPDATED blocks will be ignored.
 Try to keep the patch as small as possible. Do not use placeholders, as they will make the patch fail.
-
-We can also append to files by prefixing the filename with `append`."""
+"""
 
 ORIGINAL = "<<<<<<< ORIGINAL\n"
 DIVIDER = "\n=======\n"
 UPDATED = "\n>>>>>>> UPDATED"
 
-mode: Literal["markdown", "xml"] = "markdown"
-
-
-def patch_to_markdown(filename: str, patch: str, append: bool = False) -> str:
-    _tool = "patch" if not append else "append"
-    return f"```{_tool} {filename}\n{patch}\n```"
 
-
-def patch_to_xml(filename: str, patch: str, append: bool = False) -> str:
-    _tool = "patch" if not append else "append"
-    return f"<{_tool} filename='{filename}'>\n{patch}\n</patch>"
-
-
-def patch_to_output(filename: str, patch: str, append: bool = False) -> str:
-    if mode == "markdown":
-        return patch_to_markdown(filename, patch, append)
-    elif mode == "xml":
-        return patch_to_xml(filename, patch, append)
-    else:
-        raise ValueError(f"Invalid mode: {mode}")
+def patch_to_output(filename: str, patch: str) -> str:
+    return ToolUse("patch", [filename], patch).to_output()
 
 
 examples = f"""
 > User: patch the file `hello.py` to ask for the name of the user
 > Assistant: {patch_to_output("hello.py", '''
 <<<<<<< ORIGINAL
+def hello():
     print("Hello world")
 =======
+def hello():
     name = input("What is your name? ")
     print(f"Hello {name}")
 >>>>>>> UPDATED
 ''')}
 > System: Patch applied
-
-> User: change the codeblock to append to the file
-> Assistant: {patch_to_output("patch.py", '''
-<<<<<<< ORIGINAL
-```save hello.py
-=======
-```append hello.py
->>>>>>> UPDATED
-''')}
-
-
-> User: run the function when the script is run
-> Assistant: {patch_to_output("hello.py", '''
-<<<<<<< ORIGINAL
-    print("Hello world")
-=======
-    name = input("What is your name? ")
-    print(f"Hello {name}")
->>>>>>> UPDATED
-''', append=True)}
-""".strip()
+"""
 
 
 def apply(codeblock: str, content: str) -> str:
@@ -137,16 +100,13 @@ def apply_file(codeblock, filename):
     if not Path(filename).exists():
         raise FileNotFoundError(filename)
 
-    with open(filename) as f:
+    with open(filename, "r+") as f:
         content = f.read()
-
-    result = apply(codeblock, content)
-
-    with open(filename, "w") as f:
+        result = apply(codeblock, content)
+        f.seek(0)
+        f.truncate()
         f.write(result)
 
-    print(f"Applied patch to {filename}")
-
 
 def execute_patch(
     code: str, ask: bool, args: list[str]
@@ -157,14 +117,14 @@ def execute_patch(
     fn = " ".join(args)
     assert fn, "No filename provided"
     if ask:
-        confirm = ask_execute("Apply patch?")
+        confirm = ask_execute(f"Apply patch to {fn}?")
         if not confirm:
             print("Patch not applied")
             return
 
     try:
         apply_file(code, fn)
-        yield Message("system", "Patch applied")
+        yield Message("system", f"Patch applied to {fn}")
     except (ValueError, FileNotFoundError) as e:
         yield Message("system", f"Patch failed: {e.args[0]}")