From 68700a8cb1d9f3c89defc0f3cac9a4a6ddc2d88b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 16:10:15 +0200 Subject: [PATCH 01/17] refactor: refactor tools, streamline amending __doc__ for tools, refactor type_o_tooluse to ToolUse.from_type --- gptme/tools/__init__.py | 129 +++++++++++++++++++++++++--------------- gptme/tools/base.py | 17 +++++- gptme/tools/browser.py | 4 +- gptme/tools/chats.py | 4 +- gptme/tools/patch.py | 91 ++++++++++++++++++++++------ gptme/tools/python.py | 24 ++++++-- gptme/tools/tmux.py | 7 +-- 7 files changed, 192 insertions(+), 84 deletions(-) diff --git a/gptme/tools/__init__.py b/gptme/tools/__init__.py index a0e0a989..d92beb93 100644 --- a/gptme/tools/__init__.py +++ b/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: + + + print("Hello, world!") + + + """ + if "" 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: + + print("Hello, world!") + + """ + 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: - - - print("Hello, world!") - - - """ - if "" 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: diff --git a/gptme/tools/base.py b/gptme/tools/base.py index 081805ec..307937a0 100644 --- a/gptme/tools/base.py +++ b/gptme/tools/base.py @@ -3,6 +3,7 @@ from typing import Any, Protocol, TypeAlias from ..message import Message +from ..util import transform_examples_to_chat_directives InitFunc: TypeAlias = Callable[[], Any] @@ -10,7 +11,8 @@ 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 diff --git a/gptme/tools/browser.py b/gptme/tools/browser.py index d746f76f..dff60f84 100644 --- a/gptme/tools/browser.py +++ b/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__) diff --git a/gptme/tools/chats.py b/gptme/tools/chats.py index 6fef2a13..dbabc845 100644 --- a/gptme/tools/chats.py +++ b/gptme/tools/chats.py @@ -156,8 +156,6 @@ 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", @@ -165,3 +163,5 @@ def search_chats(query: str, max_results: int = 5) -> None: examples=examples, functions=[list_chats, search_chats], ) + +__doc__ += tool.get_doc(__doc__) diff --git a/gptme/tools/patch.py b/gptme/tools/patch.py index dc2ffd56..aec340f4 100644 --- a/gptme/tools/patch.py +++ b/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 + <<<<<<< 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" + + +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__) diff --git a/gptme/tools/python.py b/gptme/tools/python.py index 85981e8e..d4cda8e8 100644 --- a/gptme/tools/python.py +++ b/gptme/tools/python.py @@ -12,7 +12,7 @@ from typing import Literal, TypeVar, get_origin from ..message import Message -from ..util import ask_execute, print_preview, transform_examples_to_chat_directives +from ..util import ask_execute, print_preview from .base import ToolSpec logger = getLogger(__name__) @@ -188,16 +188,18 @@ def check_available_packages(): ``` """.strip() -__doc__ += transform_examples_to_chat_directives(examples) + +instructions = """ +When you send a message containing Python code (and is not a file block), it will be executed in a stateful environment. +Python will respond with the output of the execution. +""" def get_tool() -> ToolSpec: python_libraries = get_installed_python_libraries() python_libraries_str = "\n".join(f"- {lib}" for lib in python_libraries) - instructions = f""" -When you send a message containing Python code (and is not a file block), it will be executed in a stateful environment. -Python will respond with the output of the execution. + _instructions = f"""{instructions} The following libraries are available: {python_libraries_str} @@ -209,7 +211,7 @@ def get_tool() -> ToolSpec: return ToolSpec( name="python", desc="Execute Python code", - instructions=instructions, + instructions=_instructions, examples=examples, init=init_python, execute=execute_python, @@ -218,3 +220,13 @@ def get_tool() -> ToolSpec: "ipython", ], # ideally, models should use `ipython` and not `python`, but they don't ) + + +# only used for doc generation, use get_tool() in the code +tool_placeholder = ToolSpec( + name="python", + desc="Execute Python code", + instructions=instructions, + examples=examples, +) +__doc__ += tool_placeholder.get_doc() diff --git a/gptme/tools/tmux.py b/gptme/tools/tmux.py index d7b6bea2..c11f61b1 100644 --- a/gptme/tools/tmux.py +++ b/gptme/tools/tmux.py @@ -13,7 +13,7 @@ from time import sleep from ..message import Message -from ..util import ask_execute, print_preview, transform_examples_to_chat_directives +from ..util import ask_execute, print_preview from .base import ToolSpec logger = logging.getLogger(__name__) @@ -272,10 +272,6 @@ def execute_tmux( """ -new_examples = transform_examples_to_chat_directives(examples) -__doc__ += new_examples - - tool = ToolSpec( name="tmux", desc="Executes shell commands in a tmux session", @@ -286,3 +282,4 @@ def execute_tmux( block_types=["tmux"], available=shutil.which("tmux") is not None, ) +__doc__ += tool.get_doc(__doc__) From 966cbfa8049ec61576f81b079c2aff553615e07c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 17:39:41 +0200 Subject: [PATCH 02/17] Apply suggestions from code review --- gptme/tools/__init__.py | 2 +- gptme/tools/chats.py | 2 +- gptme/tools/python.py | 2 +- gptme/tools/tmux.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/gptme/tools/__init__.py b/gptme/tools/__init__.py index d92beb93..dc374fb1 100644 --- a/gptme/tools/__init__.py +++ b/gptme/tools/__init__.py @@ -203,7 +203,7 @@ def from_xml(cls, content: str) -> "Codeblock": """ root = ElementTree.fromstring(content) - return cls(root.attrib["lang"], root.text or "", root.attrib.get("filename")) + return cls(root.attrib["lang"], root.text or "", root.attrib.get("path")) @property def is_filename(self) -> bool: diff --git a/gptme/tools/chats.py b/gptme/tools/chats.py index dbabc845..6a136778 100644 --- a/gptme/tools/chats.py +++ b/gptme/tools/chats.py @@ -164,4 +164,4 @@ def search_chats(query: str, max_results: int = 5) -> None: functions=[list_chats, search_chats], ) -__doc__ += tool.get_doc(__doc__) +__doc__ = tool.get_doc(__doc__) diff --git a/gptme/tools/python.py b/gptme/tools/python.py index d4cda8e8..04702259 100644 --- a/gptme/tools/python.py +++ b/gptme/tools/python.py @@ -229,4 +229,4 @@ def get_tool() -> ToolSpec: instructions=instructions, examples=examples, ) -__doc__ += tool_placeholder.get_doc() +__doc__ = tool_placeholder.get_doc(__doc__) diff --git a/gptme/tools/tmux.py b/gptme/tools/tmux.py index c11f61b1..52ebd2e2 100644 --- a/gptme/tools/tmux.py +++ b/gptme/tools/tmux.py @@ -282,4 +282,4 @@ def execute_tmux( block_types=["tmux"], available=shutil.which("tmux") is not None, ) -__doc__ += tool.get_doc(__doc__) +__doc__ = tool.get_doc(__doc__) From 06882d3c3563e7b6ae08967329763faf4b042649 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 17:43:14 +0200 Subject: [PATCH 03/17] refactor: move Codeblock to new file codeblock.py --- gptme/codeblock.py | 78 +++++++++++++++++++++++++++++++++++++++++ gptme/llm.py | 4 +-- gptme/tools/__init__.py | 70 +++++------------------------------- gptme/util.py | 35 ------------------ 4 files changed, 88 insertions(+), 99 deletions(-) create mode 100644 gptme/codeblock.py diff --git a/gptme/codeblock.py b/gptme/codeblock.py new file mode 100644 index 00000000..178d5788 --- /dev/null +++ b/gptme/codeblock.py @@ -0,0 +1,78 @@ +from dataclasses import dataclass +from typing import Generator +from xml.etree import ElementTree + + +@dataclass +class Codeblock: + 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.strip().startswith("```"): + content = content[3:] + if content.strip().endswith("```"): + content = content[:-3] + lang = content.splitlines()[0].strip() + return cls(lang, content[len(lang) :]) + + @classmethod + def from_xml(cls, content: str) -> "Codeblock": + """ + Example: + + print("Hello, world!") + + """ + 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 "/" in self.lang + + @classmethod + def iter_from_markdown(cls, markdown: str) -> list["Codeblock"]: + return list(_extract_codeblocks(markdown)) + + +def _extract_codeblocks(markdown: str) -> Generator[Codeblock, None, None]: + # speed check (early exit): check if message contains a code block + backtick_count = markdown.count("```") + if backtick_count < 2: + return + # check that end tick has a newline (no language tag) + # if not markdown.count("```\n"): + # return + + lines = markdown.split("\n") + stack: list[str] = [] + current_block = [] + current_lang = "" + + for line in lines: + stripped_line = line.strip() + if stripped_line.startswith("```"): + if not stack: # Start of a new block + stack.append(stripped_line[3:]) + current_lang = stripped_line[3:] + elif stripped_line[3:] and stack[-1] != stripped_line[3:]: # Nested start + current_block.append(line) + stack.append(stripped_line[3:]) + else: # End of a block + if len(stack) == 1: # Outermost block + yield Codeblock(current_lang, "\n".join(current_block)) + current_block = [] + current_lang = "" + else: # Nested end + current_block.append(line) + stack.pop() + elif stack: + current_block.append(line) diff --git a/gptme/llm.py b/gptme/llm.py index 4bb4a845..663d1603 100644 --- a/gptme/llm.py +++ b/gptme/llm.py @@ -7,6 +7,7 @@ from rich import print +from .codeblock import Codeblock from .config import get_config from .constants import PROMPT_ASSISTANT from .llm_anthropic import chat as chat_anthropic @@ -19,7 +20,6 @@ from .llm_openai import stream as stream_openai from .message import Message, format_msgs, len_tokens from .models import MODELS, get_summary_model -from .util import extract_codeblocks logger = logging.getLogger(__name__) @@ -93,7 +93,7 @@ def print_clear(): sys.stdout.flush() # pause inference on finished code-block, letting user run the command before continuing - if codeblocks := extract_codeblocks(output): + if codeblocks := Codeblock.extract_codeblocks(output): lang, _ = codeblocks[0] # noreorder from .tools import is_supported_codeblock_tool # fmt: skip diff --git a/gptme/tools/__init__.py b/gptme/tools/__init__.py index dc374fb1..907b55ff 100644 --- a/gptme/tools/__init__.py +++ b/gptme/tools/__init__.py @@ -3,8 +3,8 @@ from dataclasses import dataclass from xml.etree import ElementTree +from ..codeblock import Codeblock from ..message import Message -from ..util import extract_codeblocks from .base import ToolSpec from .browser import tool as browser_tool from .chats import tool as chats_tool @@ -63,7 +63,7 @@ def execute(self, ask: bool) -> Generator[Message, None, None]: yield from tool.execute(self.content, ask, self.args) @classmethod - def from_codeblock(cls, codeblock) -> "ToolUse": + def from_codeblock(cls, codeblock: Codeblock) -> "ToolUse": """Parses a codeblock into a ToolUse. Codeblock must be a supported type. Example: @@ -71,7 +71,7 @@ def from_codeblock(cls, codeblock) -> "ToolUse": content ``` """ - if tool := get_tool_for_codeblock(codeblock.lang): + if tool := get_tool_for_langtag(codeblock.lang): # NOTE: special case args = ( codeblock.lang.split(" ")[1:] @@ -80,7 +80,6 @@ def from_codeblock(cls, codeblock) -> "ToolUse": ) 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." ) @@ -141,13 +140,12 @@ def execute_msg(msg: Message, ask: bool) -> Generator[Message, None, None]: assert msg.role == "assistant", "Only assistant messages can be executed" # get all markdown code blocks - for lang, content in extract_codeblocks(msg.content): + for codeblock in Codeblock.iter_from_markdown(msg.content): try: - if is_supported_codeblock_tool(lang): - codeblock = Codeblock(lang, content) + if get_tool_for_langtag(codeblock.lang): yield from ToolUse.from_codeblock(codeblock).execute(ask) else: - logger.info(f"Codeblock not supported: {lang}") + logger.info(f"Codeblock not supported: {codeblock.lang}") except Exception as e: logger.exception(e) yield Message( @@ -165,59 +163,14 @@ def execute_codeblock( lang: str, codeblock: str, ask: bool ) -> Generator[Message, None, None]: """Executes a codeblock and returns the output.""" - if tool := get_tool_for_codeblock(lang): + if tool := get_tool_for_langtag(lang): if tool.execute: args = lang.split(" ")[1:] yield from tool.execute(codeblock, ask, args) - assert not is_supported_codeblock_tool(codeblock) logger.debug("Unknown codeblock, neither supported language or filename.") -# TODO: use this instead of passing around codeblocks as strings (with or without ```) -@dataclass -class Codeblock: - 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.strip().startswith("```"): - content = content[3:] - if content.strip().endswith("```"): - content = content[:-3] - lang = content.splitlines()[0].strip() - return cls(lang, content[len(lang) :]) - - @classmethod - def from_xml(cls, content: str) -> "Codeblock": - """ - Example: - - print("Hello, world!") - - """ - root = ElementTree.fromstring(content) - return cls(root.attrib["lang"], root.text or "", root.attrib.get("path")) - - @property - def is_filename(self) -> bool: - return "." in self.lang or "/" in self.lang - - @property - def is_supported(self) -> bool: - 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: str) -> ToolSpec | None: +def get_tool_for_langtag(lang: str) -> ToolSpec | None: block_type = lang.split(" ")[0] for tool in loaded_tools: if block_type in tool.block_types: @@ -229,13 +182,6 @@ def get_tool_for_codeblock(lang: str) -> ToolSpec | None: return None -def is_supported_codeblock_tool(lang: str) -> bool: - if get_tool_for_codeblock(lang): - return True - else: - return False - - def get_tool(tool_name: str) -> ToolSpec: """Returns a tool by name.""" for tool in loaded_tools: diff --git a/gptme/util.py b/gptme/util.py index 60e62e7a..62bae031 100644 --- a/gptme/util.py +++ b/gptme/util.py @@ -163,41 +163,6 @@ def transform_examples_to_chat_directives(s: str, strict=False) -> str: return s -def extract_codeblocks(markdown: str) -> list[tuple[str, str]]: - # speed check (early exit): check if message contains a code block - backtick_count = markdown.count("```") - if backtick_count < 2: - return [] - - codeblocks = [] - lines = markdown.split("\n") - stack: list[str] = [] - current_block = [] - current_lang = "" - - for line in lines: - stripped_line = line.strip() - if stripped_line.startswith("```"): - if not stack: # Start of a new block - stack.append(stripped_line[3:]) - current_lang = stripped_line[3:] - elif stripped_line[3:] and stack[-1] != stripped_line[3:]: # Nested start - current_block.append(line) - stack.append(stripped_line[3:]) - else: # End of a block - if len(stack) == 1: # Outermost block - codeblocks.append((current_lang, "\n".join(current_block))) - current_block = [] - current_lang = "" - else: # Nested end - current_block.append(line) - stack.pop() - elif stack: - current_block.append(line) - - return codeblocks - - def print_bell(): """Ring the terminal bell.""" sys.stdout.write("\a") From 58a7ca88e2f57ec32159dfcb188bcdd13e1bb8b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 17:45:44 +0200 Subject: [PATCH 04/17] Apply suggestions from code review --- gptme/tools/browser.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gptme/tools/browser.py b/gptme/tools/browser.py index dff60f84..ae69ae2f 100644 --- a/gptme/tools/browser.py +++ b/gptme/tools/browser.py @@ -197,4 +197,4 @@ def html_to_markdown(html): functions=[read_url, search, screenshot_url], available=has_browser_tool(), ) -__doc__ += tool.get_doc(__doc__) +__doc__ = tool.get_doc(__doc__) From d2af905575614ea318a1f72a3d7981910238f32d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 17:51:14 +0200 Subject: [PATCH 05/17] misc fixes --- gptme/message.py | 9 +++++---- gptme/tools/shell.py | 6 ++---- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/gptme/message.py b/gptme/message.py index 0a20be76..eddb274b 100644 --- a/gptme/message.py +++ b/gptme/message.py @@ -16,8 +16,9 @@ from tomlkit._utils import escape_string from typing_extensions import Self +from .codeblock import Codeblock from .constants import ROLE_COLOR -from .util import extract_codeblocks, get_tokenizer +from .util import get_tokenizer logger = logging.getLogger(__name__) @@ -188,9 +189,9 @@ def from_toml(cls, toml: str) -> Self: timestamp=datetime.fromisoformat(msg["timestamp"]), ) - def get_codeblocks(self) -> list[tuple[str, str]]: + def get_codeblocks(self) -> list[Codeblock]: """ - Get all codeblocks from the message content, as a list of tuples (lang, content). + Get all codeblocks from the message content. """ content_str = self.content @@ -203,7 +204,7 @@ def get_codeblocks(self) -> list[tuple[str, str]]: if backtick_count < 2: return [] - return extract_codeblocks(content_str) + return Codeblock.iter_from_markdown(content_str) def format_msgs( diff --git a/gptme/tools/shell.py b/gptme/tools/shell.py index b10bfe59..45cdda5b 100644 --- a/gptme/tools/shell.py +++ b/gptme/tools/shell.py @@ -16,7 +16,7 @@ import bashlex from ..message import Message -from ..util import ask_execute, print_preview, transform_examples_to_chat_directives +from ..util import ask_execute, print_preview from .base import ToolSpec logger = logging.getLogger(__name__) @@ -369,9 +369,6 @@ def split_commands(script: str) -> list[str]: return commands -__doc__ += transform_examples_to_chat_directives(examples) - - tool = ToolSpec( name="shell", desc="Executes shell commands.", @@ -380,3 +377,4 @@ def split_commands(script: str) -> list[str]: execute=execute_shell, block_types=["bash", "sh", "shell"], ) +__doc__ = tool.get_doc(__doc__) From 6d118dd81bfe0e059d79aef7d36f96e79f5e9b33 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 18:08:57 +0200 Subject: [PATCH 06/17] fix: more codeblock refactoring --- gptme/cli.py | 7 ++-- gptme/codeblock.py | 10 ++++- gptme/commands.py | 6 +-- gptme/llm.py | 8 ++-- gptme/logmanager.py | 5 ++- gptme/reduce.py | 9 +++-- gptme/tools/__init__.py | 21 ++++++----- gptme/tools/base.py | 3 +- gptme/tools/chats.py | 1 - tests/test_codeblock.py | 81 ++++++++++++++++++++++++++++++++++++++++ tests/test_logmanager.py | 4 +- tests/test_util.py | 78 -------------------------------------- 12 files changed, 122 insertions(+), 111 deletions(-) create mode 100644 tests/test_codeblock.py diff --git a/gptme/cli.py b/gptme/cli.py index 6d98c088..3db282c9 100644 --- a/gptme/cli.py +++ b/gptme/cli.py @@ -273,13 +273,12 @@ def chat( # then exit elif not interactive: # noreorder - from .tools import is_supported_codeblock_tool # fmt: skip + from .tools import is_supported_langtag # fmt: skip # continue if we can run tools on the last message runnable = False - if codeblock := log.get_last_code_block("assistant", history=1): - lang, _ = codeblock - if is_supported_codeblock_tool(lang): + if codeblock := log.get_last_codeblock("assistant", history=1): + if is_supported_langtag(codeblock.lang): runnable = True if not runnable: logger.info("Non-interactive and exhausted prompts, exiting") diff --git a/gptme/codeblock.py b/gptme/codeblock.py index 178d5788..bb82e916 100644 --- a/gptme/codeblock.py +++ b/gptme/codeblock.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import Generator +from collections.abc import Generator from xml.etree import ElementTree @@ -14,6 +14,12 @@ def __post_init__(self): if self.path is None and self.is_filename: self.path = self.lang + def to_markdown(self) -> str: + return f"```{self.lang}\n{self.content}\n```" + + def to_xml(self) -> str: + return f'\n{self.content}\n' + @classmethod def from_markdown(cls, content: str) -> "Codeblock": if content.strip().startswith("```"): @@ -32,7 +38,7 @@ def from_xml(cls, content: str) -> "Codeblock": """ root = ElementTree.fromstring(content) - return cls(root.attrib["lang"], root.text or "", root.attrib.get("filename")) + return cls(root.attrib["lang"], root.text or "", root.attrib.get("path")) @property def is_filename(self) -> bool: diff --git a/gptme/commands.py b/gptme/commands.py index 497057f7..76be3168 100644 --- a/gptme/commands.py +++ b/gptme/commands.py @@ -193,19 +193,19 @@ def edit(log: LogManager) -> Generator[Message, None, None]: # pragma: no cover print("Applied edited messages, write /log to see the result") +# TODO: remove? def save(log: LogManager, filename: str): # save the most recent code block to a file - codeblock = log.get_last_code_block() + codeblock = log.get_last_codeblock() if not codeblock: print("No code block found") return - _, content = codeblock if Path(filename).exists(): confirm = ask_execute("File already exists, overwrite?", default=False) if not confirm: return with open(filename, "w") as f: - f.write(content) + f.write(codeblock.content) print(f"Saved code block to {filename}") diff --git a/gptme/llm.py b/gptme/llm.py index 663d1603..873974dd 100644 --- a/gptme/llm.py +++ b/gptme/llm.py @@ -93,13 +93,13 @@ def print_clear(): sys.stdout.flush() # pause inference on finished code-block, letting user run the command before continuing - if codeblocks := Codeblock.extract_codeblocks(output): - lang, _ = codeblocks[0] + if codeblocks := Codeblock.iter_from_markdown(output): + codeblock = codeblocks[0] # noreorder - from .tools import is_supported_codeblock_tool # fmt: skip + from .tools import is_supported_langtag # fmt: skip # if closing a code block supported by tools, abort generation to let them run - if is_supported_codeblock_tool(lang): + if is_supported_langtag(codeblock.lang): print("\nFound codeblock, breaking") break except KeyboardInterrupt: diff --git a/gptme/logmanager.py b/gptme/logmanager.py index df282e8b..ac896315 100644 --- a/gptme/logmanager.py +++ b/gptme/logmanager.py @@ -11,6 +11,7 @@ from rich import print +from .codeblock import Codeblock from .constants import CMDFIX from .dirs import get_logs_dir from .message import Message, len_tokens, print_msg @@ -215,11 +216,11 @@ def load( msgs = initial_msgs return cls(msgs, logdir=logdir, branch=branch, **kwargs) - def get_last_code_block( + def get_last_codeblock( self, role: RoleLiteral | None = None, history: int | None = None, - ) -> tuple[str, str] | None: + ) -> Codeblock | None: """Returns the last code block in the log, if any. If `role` set, only check that role. diff --git a/gptme/reduce.py b/gptme/reduce.py index 0608fe08..14f71453 100644 --- a/gptme/reduce.py +++ b/gptme/reduce.py @@ -8,6 +8,7 @@ from collections.abc import Generator from copy import copy +from ..codeblock import Codeblock from .message import Message, len_tokens from .models import get_model @@ -72,13 +73,13 @@ def truncate_msg(msg: Message, lines_pre=10, lines_post=10) -> Message | None: content_staged = msg.content # Truncate long codeblocks - for lang, content in msg.get_codeblocks(): + for codeblock in msg.get_codeblocks(): # check that the reformatted codeblock is in the content - full_block = f"```{lang}\n{content}\n```" + full_block = codeblock.to_markdown() assert full_block in content_staged, f"{full_block} not in {content_staged}" # truncate the middle part of the codeblock, keeping the first and last n lines - lines = content.split("\n") + lines = codeblock.content.split("\n") if len(lines) > lines_pre + lines_post + 1: content = "\n".join([*lines[:lines_pre], "[...]", *lines[-lines_post:]]) else: @@ -88,7 +89,7 @@ def truncate_msg(msg: Message, lines_pre=10, lines_post=10) -> Message | None: # replace the codeblock with the truncated version content_staged_prev = content_staged content_staged = content_staged.replace( - full_block, f"```{lang}\n{content}\n```" + full_block, Codeblock(codeblock.lang, content).to_markdown() ) assert content_staged != content_staged_prev assert full_block not in content_staged diff --git a/gptme/tools/__init__.py b/gptme/tools/__init__.py index 907b55ff..a60cb153 100644 --- a/gptme/tools/__init__.py +++ b/gptme/tools/__init__.py @@ -142,10 +142,7 @@ def execute_msg(msg: Message, ask: bool) -> Generator[Message, None, None]: # get all markdown code blocks for codeblock in Codeblock.iter_from_markdown(msg.content): try: - if get_tool_for_langtag(codeblock.lang): - yield from ToolUse.from_codeblock(codeblock).execute(ask) - else: - logger.info(f"Codeblock not supported: {codeblock.lang}") + yield from execute_codeblock(codeblock, ask) except Exception as e: logger.exception(e) yield Message( @@ -160,14 +157,16 @@ def execute_msg(msg: Message, ask: bool) -> Generator[Message, None, None]: def execute_codeblock( - lang: str, codeblock: str, ask: bool + codeblock: Codeblock, ask: bool ) -> Generator[Message, None, None]: """Executes a codeblock and returns the output.""" - if tool := get_tool_for_langtag(lang): + ToolUse.from_codeblock(codeblock) + if tool := get_tool_for_langtag(codeblock.lang): if tool.execute: - args = lang.split(" ")[1:] - yield from tool.execute(codeblock, ask, args) - logger.debug("Unknown codeblock, neither supported language or filename.") + args = codeblock.lang.split(" ")[1:] + yield from tool.execute(codeblock.content, ask, args) + else: + logger.info(f"Codeblock not supported: {codeblock.lang}") def get_tool_for_langtag(lang: str) -> ToolSpec | None: @@ -182,6 +181,10 @@ def get_tool_for_langtag(lang: str) -> ToolSpec | None: return None +def is_supported_langtag(lang: str) -> bool: + return bool(get_tool_for_langtag(lang)) + + def get_tool(tool_name: str) -> ToolSpec: """Returns a tool by name.""" for tool in loaded_tools: diff --git a/gptme/tools/base.py b/gptme/tools/base.py index 307937a0..3fac1d92 100644 --- a/gptme/tools/base.py +++ b/gptme/tools/base.py @@ -11,8 +11,7 @@ class ExecuteFunc(Protocol): def __call__( self, code: str, ask: bool, args: list[str] - ) -> Generator[Message, None, None]: - ... + ) -> Generator[Message, None, None]: ... @dataclass diff --git a/gptme/tools/chats.py b/gptme/tools/chats.py index 6a136778..bde0f460 100644 --- a/gptme/tools/chats.py +++ b/gptme/tools/chats.py @@ -8,7 +8,6 @@ from ..llm import summarize as llm_summarize from ..message import Message -from ..util import transform_examples_to_chat_directives from .base import ToolSpec logger = logging.getLogger(__name__) diff --git a/tests/test_codeblock.py b/tests/test_codeblock.py new file mode 100644 index 00000000..bcef387d --- /dev/null +++ b/tests/test_codeblock.py @@ -0,0 +1,81 @@ +from gptme.codeblock import Codeblock + + +def test_extract_codeblocks_basic(): + markdown = """ +Some text +```python +def hello(): + print("Hello, World!") +``` +More text +""" + assert Codeblock.iter_from_markdown(markdown) == [ + ("python", 'def hello():\n print("Hello, World!")') + ] + + +def test_extract_codeblocks_multiple(): + markdown = """ +```java +public class Main { + public static void main(String[] args) { + System.out.println("Hello, Java!"); + } +} +``` +Some text +```python +def greet(name): + return f"Hello, {name}!" +``` +""" + assert Codeblock.iter_from_markdown(markdown) == [ + ( + "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}!"'), + ] + + +def test_extract_codeblocks_nested(): + markdown = """ +```python +def print_readme(): + print('''Usage: +```javascript +callme() +``` +''') +``` +""" + assert Codeblock.iter_from_markdown(markdown) == [ + ( + "python", + "def print_readme():\n print('''Usage:\n```javascript\ncallme()\n```\n''')", + ) + ] + + +def test_extract_codeblocks_empty(): + assert Codeblock.iter_from_markdown("") == [] + + +def test_extract_codeblocks_text_only(): + assert ( + Codeblock.iter_from_markdown("Just some regular text\nwithout any code blocks.") + == [] + ) + + +def test_extract_codeblocks_no_language(): + markdown = """ +``` +def hello(): + print("Hello, World!") +``` +""" + assert Codeblock.iter_from_markdown(markdown) == [ + ("", 'def hello():\n print("Hello, World!")') + ] diff --git a/tests/test_logmanager.py b/tests/test_logmanager.py index da0229e4..9ad8789e 100644 --- a/tests/test_logmanager.py +++ b/tests/test_logmanager.py @@ -1,7 +1,7 @@ from gptme.logmanager import LogManager, Message -def test_get_last_code_block(): +def test_get_last_codeblock(): # tests that the last code block is indeed returned, with the correct formatting log = LogManager() log.append( @@ -18,7 +18,7 @@ def test_get_last_code_block(): """, ) ) - assert ("python", "print('world')") == log.get_last_code_block() + assert ("python", "print('world')") == log.get_last_codeblock() def test_branch(): diff --git a/tests/test_util.py b/tests/test_util.py index 257c67b0..bc783c13 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -2,7 +2,6 @@ from gptme.util import ( epoch_to_age, - extract_codeblocks, generate_name, is_generated_name, transform_examples_to_chat_directives, @@ -57,80 +56,3 @@ def test_transform_examples_to_chat_directives_tricky(): Assistant: lol""" assert transform_examples_to_chat_directives(src, strict=True) == expected - - -def test_extract_codeblocks_basic(): - markdown = """ -Some text -```python -def hello(): - print("Hello, World!") -``` -More text -""" - assert extract_codeblocks(markdown) == [ - ("python", 'def hello():\n print("Hello, World!")') - ] - - -def test_extract_codeblocks_multiple(): - markdown = """ -```java -public class Main { - public static void main(String[] args) { - System.out.println("Hello, Java!"); - } -} -``` -Some text -```python -def greet(name): - return f"Hello, {name}!" -``` -""" - assert extract_codeblocks(markdown) == [ - ( - "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}!"'), - ] - - -def test_extract_codeblocks_nested(): - markdown = """ -```python -def print_readme(): - print('''Usage: -```javascript -callme() -``` -''') -``` -""" - assert extract_codeblocks(markdown) == [ - ( - "python", - "def print_readme():\n print('''Usage:\n```javascript\ncallme()\n```\n''')", - ) - ] - - -def test_extract_codeblocks_empty(): - assert extract_codeblocks("") == [] - - -def test_extract_codeblocks_text_only(): - assert extract_codeblocks("Just some regular text\nwithout any code blocks.") == [] - - -def test_extract_codeblocks_no_language(): - markdown = """ -``` -def hello(): - print("Hello, World!") -``` -""" - assert extract_codeblocks(markdown) == [ - ("", 'def hello():\n print("Hello, World!")') - ] From 9d664016106ffec5f0a2f2bf0bd47fae44301a27 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 18:11:37 +0200 Subject: [PATCH 07/17] Apply suggestions from code review --- gptme/tools/patch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gptme/tools/patch.py b/gptme/tools/patch.py index aec340f4..c8995009 100644 --- a/gptme/tools/patch.py +++ b/gptme/tools/patch.py @@ -20,7 +20,7 @@ .. chat:: - User: edit the file `hello.py` to ask for the name to greet + User: edit `hello.py` to ask for the name to greet Assistant: ```patch hello.py <<<<<<< ORIGINAL @@ -189,4 +189,4 @@ def execute_patch( execute=execute_patch, block_types=["patch"], ) -__doc__ += tool.get_doc(__doc__) +__doc__ = tool.get_doc(__doc__) From 3c98d5d3a403c1c01802d4ab1d1d802861f1737a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 18:19:18 +0200 Subject: [PATCH 08/17] 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 - <<<<<<< 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(): From 38eadcb58c8f7e4f42e457664bd3e9896e26237d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 18:38:52 +0200 Subject: [PATCH 09/17] more fixes --- gptme/tools/read.py | 5 +++++ gptme/tools/save.py | 13 ++----------- gptme/tools/subagent.py | 5 ++--- 3 files changed, 9 insertions(+), 14 deletions(-) diff --git a/gptme/tools/read.py b/gptme/tools/read.py index 478caa82..2474e111 100644 --- a/gptme/tools/read.py +++ b/gptme/tools/read.py @@ -1,3 +1,7 @@ +""" +Read the contents of a file. +""" + from gptme.tools.base import ToolSpec # Note: this isn't actually a tool, it only serves prompting purposes @@ -9,3 +13,4 @@ cat file.txt ```""", ) +__doc__ = tool.get_doc(__doc__) diff --git a/gptme/tools/save.py b/gptme/tools/save.py index 7afced91..8f706027 100644 --- a/gptme/tools/save.py +++ b/gptme/tools/save.py @@ -1,16 +1,5 @@ """ Gives the assistant the ability to save code to a file. - -Example: - -.. chat:: - - User: write hello world to hello.py - Assistant: - ```save hello.py - print("hello world") - ``` - System: Saved to hello.py """ from collections.abc import Generator @@ -138,6 +127,7 @@ def execute_append( execute=execute_save, block_types=["save"], ) +__doc__ = tool_save.get_doc(__doc__) instructions_append = """ To append code to a file, use a code block with the language: append @@ -160,3 +150,4 @@ def execute_append( execute=execute_append, block_types=["append"], ) +__doc__ = tool_append.get_doc(__doc__) diff --git a/gptme/tools/subagent.py b/gptme/tools/subagent.py index ee94827c..1eb41572 100644 --- a/gptme/tools/subagent.py +++ b/gptme/tools/subagent.py @@ -12,7 +12,6 @@ from typing import TYPE_CHECKING, Literal from ..message import Message -from ..util import transform_examples_to_chat_directives from .base import ToolSpec from .python import register_function @@ -20,6 +19,7 @@ # noreorder from ..logmanager import LogManager # fmt: skip + logger = logging.getLogger(__name__) Status = Literal["running", "success", "failure"] @@ -163,8 +163,6 @@ def subagent_wait(agent_id: str) -> dict: ``` """ -__doc__ += transform_examples_to_chat_directives(examples) - tool = ToolSpec( name="subagent", @@ -172,3 +170,4 @@ def subagent_wait(agent_id: str) -> dict: examples=examples, functions=[subagent, subagent_status, subagent_wait], ) +__doc__ = tool.get_doc(__doc__) From d3d3534bee76a4e7091ab77c370f3ae11b8e1290 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 19:35:50 +0200 Subject: [PATCH 10/17] fix: removed save command --- gptme/commands.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/gptme/commands.py b/gptme/commands.py index 76be3168..297875ef 100644 --- a/gptme/commands.py +++ b/gptme/commands.py @@ -2,7 +2,6 @@ import re import sys from collections.abc import Generator -from pathlib import Path from time import sleep from typing import Literal @@ -35,7 +34,6 @@ "fork", "summarize", "context", - "save", "shell", "python", "replay", @@ -54,7 +52,6 @@ "rename": "Rename the conversation", "fork": "Create a copy of the conversation with a new name", "summarize": "Summarize the conversation", - "save": "Save the last code block to a file", "shell": "Execute shell code", "python": "Execute Python code", "replay": "Re-execute codeblocks in the conversation, wont store output in log", @@ -124,11 +121,6 @@ def handle_cmd( # if int, undo n messages n = int(args[0]) if args and args[0].isdigit() else 1 log.undo(n) - case "save": - # undo - log.undo(1, quiet=True) - filename = args[0] if args else input("Filename: ") - save(log, filename) case "exit": sys.exit(0) case "replay": @@ -193,22 +185,6 @@ def edit(log: LogManager) -> Generator[Message, None, None]: # pragma: no cover print("Applied edited messages, write /log to see the result") -# TODO: remove? -def save(log: LogManager, filename: str): - # save the most recent code block to a file - codeblock = log.get_last_codeblock() - if not codeblock: - print("No code block found") - return - if Path(filename).exists(): - confirm = ask_execute("File already exists, overwrite?", default=False) - if not confirm: - return - with open(filename, "w") as f: - f.write(codeblock.content) - print(f"Saved code block to {filename}") - - def rename(log: LogManager, new_name: str, ask: bool = True): if new_name in ["", "auto"]: new_name = llm.generate_name(log.prepare_messages()) From 8ef05f8e315f3ae2962dab03f6b8974615127a2e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 19:36:52 +0200 Subject: [PATCH 11/17] Apply suggestions from code review --- gptme/tools/base.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/gptme/tools/base.py b/gptme/tools/base.py index cf85e063..3e80a8e0 100644 --- a/gptme/tools/base.py +++ b/gptme/tools/base.py @@ -36,9 +36,6 @@ 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\n# Instructions\n\n{self.instructions}" if self.examples: _doc += f"\n\n# Examples\n\n{transform_examples_to_chat_directives(self.examples)}" return _doc From 6e71fa9bb832d8a9c3055feba4cf3c334512ad01 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 19:41:05 +0200 Subject: [PATCH 12/17] removed unused test --- tests/test_cli.py | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index 76933cdd..a77c74d9 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -8,7 +8,7 @@ import gptme.constants import pytest from click.testing import CliRunner -from gptme.constants import CMDFIX, MULTIPROMPT_SEPARATOR +from gptme.constants import CMDFIX project_root = Path(__file__).parent.parent logo = project_root / "media" / "logo.png" @@ -85,22 +85,6 @@ def test_command_summarize(args: list[str], runner: CliRunner): assert result.exit_code == 0 -@pytest.mark.slow -def test_command_save(args: list[str], runner: CliRunner): - # tests the /save command - args.append(f"{CMDFIX}impersonate ```ipython\nprint('hello')\n```") - args.append(MULTIPROMPT_SEPARATOR) - args.append(f"{CMDFIX}save output.txt") - print(f"running: gptme {' '.join(args)}") - result = runner.invoke(gptme.cli.main, args) - assert result.exit_code == 0 - - # read the file - with open("output.txt") as f: - content = f.read() - assert content == "hello" - - def test_command_fork(args: list[str], runner: CliRunner, name: str): # tests the /fork command name += "-fork" From 55947547a46d1e7269b33d90734d3c5e51213669 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 19:43:07 +0200 Subject: [PATCH 13/17] Apply suggestions from code review --- gptme/tools/patch.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gptme/tools/patch.py b/gptme/tools/patch.py index ced02b79..c7affb6b 100644 --- a/gptme/tools/patch.py +++ b/gptme/tools/patch.py @@ -30,7 +30,6 @@ 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```" From 98ad29e0af597b06593f9e154d1b268f3fc0bb6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 21:21:37 +0200 Subject: [PATCH 14/17] fix: more refactoring, use block_types to register commands for python, shell, etc --- gptme/cli.py | 19 ++++++------- gptme/commands.py | 63 +++++++++++++++++++++++++++-------------- gptme/tools/__init__.py | 7 +---- gptme/tools/python.py | 42 +++++++++++++-------------- 4 files changed, 72 insertions(+), 59 deletions(-) diff --git a/gptme/cli.py b/gptme/cli.py index 3db282c9..809ae329 100644 --- a/gptme/cli.py +++ b/gptme/cli.py @@ -17,7 +17,12 @@ from rich import print # noqa: F401 from rich.console import Console -from .commands import CMDFIX, action_descriptions, execute_cmd +from .commands import ( + CMDFIX, + _gen_help, + action_descriptions, + execute_cmd, +) from .config import get_workspace_prompt from .constants import MULTIPROMPT_SEPARATOR, PROMPT_USER from .dirs import get_logs_dir @@ -34,15 +39,9 @@ logger = logging.getLogger(__name__) print_builtin = __builtins__["print"] # type: ignore -# TODO: these are a bit redundant/incorrect -LLMChoice = Literal["openai", "anthropic", "local"] -ModelChoice = Literal["gpt-3.5-turbo", "gpt-4", "gpt-4-1106-preview"] - script_path = Path(os.path.realpath(__file__)) -action_readme = "\n".join( - f" {CMDFIX}{cmd:11s} {desc}." for cmd, desc in action_descriptions.items() -) +commands_help = "\n".join(_gen_help(incl_langtags=False)) docstring = f""" @@ -55,7 +54,7 @@ The chat offers some commands that can be used to interact with the system: \b -{action_readme}""" +{commands_help}""" @click.command(help=docstring) @@ -121,7 +120,7 @@ def main( prompts: list[str], prompt_system: str, name: str, - model: ModelChoice, + model: str, stream: bool, verbose: bool, no_confirm: bool, diff --git a/gptme/commands.py b/gptme/commands.py index 297875ef..c4c9d1c3 100644 --- a/gptme/commands.py +++ b/gptme/commands.py @@ -6,6 +6,7 @@ from typing import Literal from . import llm +from .codeblock import Codeblock from .constants import CMDFIX from .logmanager import LogManager from .message import ( @@ -17,9 +18,9 @@ ) from .models import get_model from .tools import ( + execute_codeblock, execute_msg, - execute_python, - execute_shell, + is_supported_langtag, loaded_tools, ) from .useredit import edit_text_with_editor @@ -34,8 +35,6 @@ "fork", "summarize", "context", - "shell", - "python", "replay", "undo", "impersonate", @@ -52,8 +51,6 @@ "rename": "Rename the conversation", "fork": "Create a copy of the conversation with a new name", "summarize": "Summarize the conversation", - "shell": "Execute shell code", - "python": "Execute Python code", "replay": "Re-execute codeblocks in the conversation, wont store output in log", "impersonate": "Impersonate the assistant", "tokens": "Show the number of tokens used", @@ -86,11 +83,6 @@ def handle_cmd( name, *args = re.split(r"[\n\s]", cmd) full_args = cmd.split(" ", 1)[1] if " " in cmd else "" match name: - # TODO: rewrite to auto-register tools using block_types - case "bash" | "sh" | "shell": - yield from execute_shell(full_args, ask=not no_confirm, args=[]) - case "python" | "py": - yield from execute_python(full_args, ask=not no_confirm, args=[]) case "log": log.undo(1, quiet=True) log.print(show_hidden="--hidden" in args) @@ -156,12 +148,18 @@ def handle_cmd( """.strip() ) case _: - if log.log[-1].content != f"{CMDFIX}help": - print("Unknown command") - # undo the '/help' command itself - log.undo(1, quiet=True) - log.write() - help() + # the case for python, shell, and other block_types supported by tools + if is_supported_langtag(name): + yield from execute_codeblock( + Codeblock(name, full_args), ask=not no_confirm + ) + else: + if log.log[-1].content != f"{CMDFIX}help": + print("Unknown command") + # undo the '/help' command itself + log.undo(1, quiet=True) + log.write() + help() def edit(log: LogManager) -> Generator[Message, None, None]: # pragma: no cover @@ -201,8 +199,31 @@ def rename(log: LogManager, new_name: str, ask: bool = True): print(f"Renamed conversation to {log.logfile.parent}") -def help(): - longest_cmd = max(len(cmd) for cmd in COMMANDS) - print("Available commands:") +def _gen_help(incl_langtags: bool = True) -> Generator[str, None, None]: + yield "Available commands:" + max_cmdlen = max(len(cmd) for cmd in COMMANDS) for cmd, desc in action_descriptions.items(): - print(f" /{cmd.ljust(longest_cmd)} {desc}") + yield f" /{cmd.ljust(max_cmdlen)} {desc}" + + if incl_langtags: + yield "" + yield "To execute code with supported tools, use the following syntax:" + yield " / " + yield "" + yield "Example:" + yield " /sh " + yield " /python " + yield "" + yield "Supported langtags:" + for tool in loaded_tools: + if tool.block_types: + yield f" - {tool.block_types[0]}" + ( + f" (alias: {', '.join(tool.block_types[1:])})" + if len(tool.block_types) > 1 + else "" + ) + + +def help(): + for line in _gen_help(): + print(line) diff --git a/gptme/tools/__init__.py b/gptme/tools/__init__.py index a60cb153..214ed918 100644 --- a/gptme/tools/__init__.py +++ b/gptme/tools/__init__.py @@ -10,12 +10,10 @@ from .chats import tool as chats_tool from .gh import tool as gh_tool from .patch import tool as patch_tool -from .python import execute_python from .python import get_tool as get_python_tool from .python import register_function from .read import tool as tool_read -from .save import execute_save, tool_append, tool_save -from .shell import execute_shell +from .save import tool_append, tool_save from .shell import tool as shell_tool from .subagent import tool as subagent_tool from .tmux import tool as tmux_tool @@ -25,9 +23,6 @@ __all__ = [ "execute_codeblock", - "execute_python", - "execute_shell", - "execute_save", "ToolSpec", "ToolUse", "all_tools", diff --git a/gptme/tools/python.py b/gptme/tools/python.py index 04702259..844d4bf0 100644 --- a/gptme/tools/python.py +++ b/gptme/tools/python.py @@ -4,6 +4,7 @@ It uses IPython to do so, and persists the IPython instance between calls to give a REPL-like experience. """ +import dataclasses import functools import re import types @@ -195,6 +196,23 @@ def check_available_packages(): """ +# only used for doc generation, use get_tool() in the code +tool = ToolSpec( + name="python", + desc="Execute Python code", + instructions=instructions, + examples=examples, + init=init_python, + execute=execute_python, + block_types=[ + "python", + "ipython", + "py", + ], +) +__doc__ = tool.get_doc(__doc__) + + def get_tool() -> ToolSpec: python_libraries = get_installed_python_libraries() python_libraries_str = "\n".join(f"- {lib}" for lib in python_libraries) @@ -208,25 +226,5 @@ def get_tool() -> ToolSpec: {get_functions_prompt()} """.strip() - return ToolSpec( - name="python", - desc="Execute Python code", - instructions=_instructions, - examples=examples, - init=init_python, - execute=execute_python, - block_types=[ - "python", - "ipython", - ], # ideally, models should use `ipython` and not `python`, but they don't - ) - - -# only used for doc generation, use get_tool() in the code -tool_placeholder = ToolSpec( - name="python", - desc="Execute Python code", - instructions=instructions, - examples=examples, -) -__doc__ = tool_placeholder.get_doc(__doc__) + # create a copy with the updated instructions + return dataclasses.replace(tool, instructions=_instructions) From 7fb8aa77b9014de3896ae45381332daef93ab83f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 21:36:50 +0200 Subject: [PATCH 15/17] fix: minor improvements to /help and /tools formatting --- gptme/commands.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/gptme/commands.py b/gptme/commands.py index c4c9d1c3..5da97e69 100644 --- a/gptme/commands.py +++ b/gptme/commands.py @@ -143,9 +143,9 @@ def handle_cmd( for tool in loaded_tools: print( f""" -- {tool.name} ({tool.desc.rstrip(".")}) - tokens (example): {len_tokens(tool.examples)} - """.strip() + # {tool.name} + {tool.desc.rstrip(".")} + tokens (example): {len_tokens(tool.examples)}""" ) case _: # the case for python, shell, and other block_types supported by tools @@ -211,8 +211,8 @@ def _gen_help(incl_langtags: bool = True) -> Generator[str, None, None]: yield " / " yield "" yield "Example:" - yield " /sh " - yield " /python " + yield " /sh echo hello" + yield " /python print('hello')" yield "" yield "Supported langtags:" for tool in loaded_tools: From 3ae46497f51d5dae59ccedd484d6d44ae1d21277 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 21:55:43 +0200 Subject: [PATCH 16/17] fix: misc improvements to docs --- docs/api.rst | 58 ++++++++++----------------------------------- gptme/server/api.py | 6 +++-- gptme/tools/base.py | 17 +++++++------ 3 files changed, 26 insertions(+), 55 deletions(-) diff --git a/docs/api.rst b/docs/api.rst index eb46611b..72325fc9 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -6,66 +6,32 @@ Here is the API reference for ``gptme``. .. toctree:: :maxdepth: 2 -gptme package -^^^^^^^^^^^^^ +core +---- -Some of the main classes in ``gptme``. +Some of the core classes in ``gptme``. -gptme.message -------------- - -.. automodule:: gptme.message +.. autoclass:: gptme.message.Message :members: -gptme.logmanager ----------------- - -.. automodule:: gptme.logmanager +.. automodule:: gptme.codeblock :members: -gptme.server -^^^^^^^^^^^^ - -Endpoint functions for the server. - -.. automodule:: gptme.server +.. automodule:: gptme.logmanager :members: - :undoc-members: -gptme.tools -^^^^^^^^^^^ +tools +----- -Tools available to gptme. +Supporting classes and functions for creating and using tools. .. automodule:: gptme.tools :members: -gptme.tools.shell ------------------ - -.. automodule:: gptme.tools.shell - :members: - -Python +server ------ -.. automodule:: gptme.tools.python - :members: - -gptme.tools.context -------------------- - -.. automodule:: gptme.tools.context - :members: - -gptme.tools.save ----------------- +See `Server `_ for more information. -.. automodule:: gptme.tools.save - :members: - -gptme.tools.patch ------------------ - -.. automodule:: gptme.tools.patch +.. automodule:: gptme.server :members: diff --git a/gptme/server/api.py b/gptme/server/api.py index 2d97e2cb..86bcb675 100644 --- a/gptme/server/api.py +++ b/gptme/server/api.py @@ -144,12 +144,14 @@ def favicon(): return flask.send_from_directory(static_path.parent / "media", "logo.png") -def create_app(): +def create_app() -> flask.Flask: + """Create the Flask app.""" app = flask.Flask(__name__, static_folder=static_path) app.register_blueprint(api) return app -def main(): +def main() -> None: + """Run the Flask app.""" app = create_app() app.run(debug=True) diff --git a/gptme/tools/base.py b/gptme/tools/base.py index 3e80a8e0..1722f77d 100644 --- a/gptme/tools/base.py +++ b/gptme/tools/base.py @@ -11,7 +11,8 @@ class ExecuteFunc(Protocol): def __call__( self, code: str, ask: bool, args: list[str] - ) -> Generator[Message, None, None]: ... + ) -> Generator[Message, None, None]: + ... @dataclass @@ -32,10 +33,12 @@ class ToolSpec: 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" + def get_doc(self, doc="") -> str: + """Returns an updated docstring with examples.""" + if doc: + doc += "\n\n" if self.examples: - _doc += f"\n\n# Examples\n\n{transform_examples_to_chat_directives(self.examples)}" - return _doc + doc += ( + f"# Examples\n\n{transform_examples_to_chat_directives(self.examples)}" + ) + return doc From f8b7bd261fff1d450404c192b66c2c29d933e859 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Mon, 9 Sep 2024 22:16:18 +0200 Subject: [PATCH 17/17] Apply suggestions from code review --- gptme/codeblock.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/gptme/codeblock.py b/gptme/codeblock.py index 6503143c..b73909c6 100644 --- a/gptme/codeblock.py +++ b/gptme/codeblock.py @@ -54,9 +54,6 @@ def _extract_codeblocks(markdown: str) -> Generator[Codeblock, None, None]: backtick_count = markdown.count("```") if backtick_count < 2: return - # check that end tick has a newline (no language tag) - # if not markdown.count("```\n"): - # return lines = markdown.split("\n") stack: list[str] = []