refactor: refactor tools, codeblock, and tooluse

gptme/tools/__init__.py

@@ -62,6 +62,52 @@ def execute(self, ask: bool) -> Generator[Message, None, None]:
         if tool.execute:
             yield from tool.execute(self.content, ask, self.args)
 
+    @classmethod
+    def from_codeblock(cls, codeblock) -> "ToolUse":
+        """Parses a codeblock into a ToolUse. Codeblock must be a supported type.
+
+        Example:
+          ```lang
+          content
+          ```
+        """
+        if tool := get_tool_for_codeblock(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:
+            assert not is_supported_codeblock_tool(codeblock.lang)
+            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."""
@@ -98,7 +144,8 @@ def execute_msg(msg: Message, ask: bool) -> Generator[Message, None, None]:
     for lang, content in extract_codeblocks(msg.content):
         try:
             if is_supported_codeblock_tool(lang):
-                yield from codeblock_to_tooluse(lang, content).execute(ask)
+                codeblock = Codeblock(lang, content)
+                yield from ToolUse.from_codeblock(codeblock).execute(ask)
             else:
                 logger.info(f"Codeblock not supported: {lang}")
         except Exception as e:
@@ -110,23 +157,10 @@ def execute_msg(msg: Message, ask: bool) -> Generator[Message, None, None]:
             break
 
     # TODO: execute them in order with codeblocks
-    for tooluse in get_tooluse_xml(msg.content):
+    for tooluse in ToolUse.iter_from_xml(msg.content):
         yield from tooluse.execute(ask)
 
 
-def codeblock_to_tooluse(lang: str, content: str) -> ToolUse:
-    """Parses a codeblock into a ToolUse. Codeblock must be a supported type."""
-    if tool := get_tool_for_codeblock(lang):
-        # NOTE: special case
-        args = lang.split(" ")[1:] if tool.name != "save" else [lang]
-        return ToolUse(tool.name, args, content)
-    else:
-        assert not is_supported_codeblock_tool(lang)
-        raise ValueError(
-            f"Unknown codeblock type '{lang}', neither supported language or filename."
-        )
-
-
 def execute_codeblock(
     lang: str, codeblock: str, ask: bool
 ) -> Generator[Message, None, None]:
@@ -142,67 +176,66 @@ def execute_codeblock(
 # TODO: use this instead of passing around codeblocks as strings (with or without ```)
 @dataclass
 class Codeblock:
-    lang_or_fn: str
+    lang: str
     content: str
+    path: str | None = None
+
+    # init path in __post_init__ if path is None and lang is pathy
+    def __post_init__(self):
+        if self.path is None and self.is_filename:
+            self.path = self.lang
 
     @classmethod
     def from_markdown(cls, content: str) -> "Codeblock":
-        if content.startswith("```"):
+        if content.strip().startswith("```"):
             content = content[3:]
-        if content.endswith("```"):
+        if content.strip().endswith("```"):
             content = content[:-3]
-        lang_or_fn = content.splitlines()[0].strip()
-        return cls(lang_or_fn, content[len(lang_or_fn) :])
+        lang = content.splitlines()[0].strip()
+        return cls(lang, content[len(lang) :])
+
+    @classmethod
+    def from_xml(cls, content: str) -> "Codeblock":
+        """
+        Example:
+          <codeblock lang="python" path="example.py">
+          print("Hello, world!")
+          </codeblock>
+        """
+        root = ElementTree.fromstring(content)
+        return cls(root.attrib["lang"], root.text or "", root.attrib.get("filename"))
 
     @property
     def is_filename(self) -> bool:
-        return "." in self.lang_or_fn or "/" in self.lang_or_fn
+        return "." in self.lang or "/" in self.lang
 
     @property
     def is_supported(self) -> bool:
-        return is_supported_codeblock_tool(self.lang_or_fn)
+        return is_supported_codeblock_tool(self.lang)
+
+    def execute(self, ask: bool) -> Generator[Message, None, None]:
+        return execute_codeblock(self.lang, self.content, ask)
 
 
-def get_tool_for_codeblock(lang_or_fn: str) -> ToolSpec | None:
-    block_type = lang_or_fn.split(" ")[0]
+def get_tool_for_codeblock(lang: str) -> ToolSpec | None:
+    block_type = lang.split(" ")[0]
     for tool in loaded_tools:
         if block_type in tool.block_types:
             return tool
-    is_filename = "." in lang_or_fn or "/" in lang_or_fn
+    is_filename = "." in lang or "/" in lang
     if is_filename:
         # NOTE: special case
         return tool_save
     return None
 
 
-def is_supported_codeblock_tool(lang_or_fn: str) -> bool:
-    if get_tool_for_codeblock(lang_or_fn):
+def is_supported_codeblock_tool(lang: str) -> bool:
+    if get_tool_for_codeblock(lang):
         return True
     else:
         return False
 
 
-def get_tooluse_xml(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 get_tool(tool_name: str) -> ToolSpec:
     """Returns a tool by name."""
     for tool in loaded_tools:

gptme/tools/base.py

@@ -3,14 +3,16 @@
 from typing import Any, Protocol, TypeAlias
 
 from ..message import Message
+from ..util import transform_examples_to_chat_directives
 
 InitFunc: TypeAlias = Callable[[], Any]
 
 
 class ExecuteFunc(Protocol):
     def __call__(
         self, code: str, ask: bool, args: list[str]
-    ) -> Generator[Message, None, None]: ...
+    ) -> Generator[Message, None, None]:
+        ...
 
 
 @dataclass
@@ -30,3 +32,16 @@ class ToolSpec:
     execute: ExecuteFunc | None = None
     block_types: list[str] = field(default_factory=list)
     available: bool = True
+
+    def get_doc(self, _doc="") -> str:
+        """Returns a string about the tool to be appended to the __doc__ string of the module."""
+        if _doc:
+            _doc += "\n\n"
+        _doc += f"ToolSpec: {self.name}\n{self.desc}"
+        if self.instructions:
+            _doc += f"\n# Instructions: {self.instructions}"
+        if self.examples:
+            _doc += (
+                f"\n# Examples: {transform_examples_to_chat_directives(self.examples)}"
+            )
+        return _doc

gptme/tools/browser.py

@@ -19,7 +19,6 @@
 import tempfile
 from typing import Literal
 
-from ..util import transform_examples_to_chat_directives
 from .base import ToolSpec
 
 has_playwright = importlib.util.find_spec("playwright") is not None
@@ -190,8 +189,6 @@ def html_to_markdown(html):
     return markdown
 
 
-__doc__ += transform_examples_to_chat_directives(examples)
-
 tool = ToolSpec(
     name="browser",
     desc="Browse the web",
@@ -200,3 +197,4 @@ def html_to_markdown(html):
     functions=[read_url, search, screenshot_url],
     available=has_browser_tool(),
 )
+__doc__ += tool.get_doc(__doc__)

gptme/tools/chats.py

@@ -156,12 +156,12 @@ def search_chats(query: str, max_results: int = 5) -> None:
 ```
 """
 
-__doc__ += transform_examples_to_chat_directives(examples)
-
 tool = ToolSpec(
     name="chats",
     desc="List, search, and summarize past conversation logs",
     instructions=instructions,
     examples=examples,
     functions=[list_chats, search_chats],
 )
+
+__doc__ += tool.get_doc(__doc__)

gptme/tools/patch.py

@@ -1,29 +1,44 @@
 """
-Gives the LLM agent the ability to patch files, by using a adapted version git conflict markers.
+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: patch the file `hello.py` to ask for the name of the user
+    User: edit the file `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}")
+    print(f"Hello {name}")
     >>>>>>> UPDATED
     ```
     System: Patch applied
-
-Inspired by aider.
 """
 
 # 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
+from typing import Literal
 
 from ..message import Message
 from ..util import ask_execute
@@ -32,33 +47,70 @@
 instructions = """
 To patch/modify files, we can use an adapted version of 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.
+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.
+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`."""
 
-examples = """
+ORIGINAL = "<<<<<<< ORIGINAL\n"
+DIVIDER = "\n=======\n"
+UPDATED = "\n>>>>>>> UPDATED"
+
+mode: Literal["markdown", "xml"] = "markdown"
+
+
+def patch_to_markdown(patch: str, filename: str, append: bool = False) -> str:
+    _tool = "patch" if not append else "append"
+    patch = patch.replace("```", "\\`\\`\\`")
+    return f"```{_tool} {filename}\n{patch}\n```"
+
+
+def patch_to_xml(patch: str, filename: 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(patch: str, filename: str, append: bool = False) -> str:
+    if mode == "markdown":
+        return patch_to_markdown(patch, filename, append)
+    elif mode == "xml":
+        return patch_to_xml(patch, filename, append)
+    else:
+        raise ValueError(f"Invalid mode: {mode}")
+
+
+examples = f"""
 > User: patch the file `hello.py` to ask for the name of the user
-```patch hello.py
+> Assistant: {patch_to_output("hello.py", '''
 <<<<<<< ORIGINAL
     print("Hello world")
 =======
     name = input("What is your name? ")
-    print(f"hello {name}")
+    print(f"Hello {name}")
 >>>>>>> UPDATED
-```
+''')}
 
-> User: run the function when the script is run
+> User: change the codeblock to append to the file
+> Assistant: {patch_to_output("patch.py", '''
+<<<<<<< ORIGINAL
+```save hello.py
+=======
 ```append hello.py
-if __name__ == "__main__":
-    hello()
-```
-""".strip()
+>>>>>>> UPDATED
+''')}
 
-ORIGINAL = "<<<<<<< ORIGINAL\n"
-DIVIDER = "\n=======\n"
-UPDATED = "\n>>>>>>> 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,3 +189,4 @@ def execute_patch(
     execute=execute_patch,
     block_types=["patch"],
 )
+__doc__ += tool.get_doc(__doc__)