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/cli.py b/gptme/cli.py index 6d98c088..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, @@ -273,13 +272,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 new file mode 100644 index 00000000..b73909c6 --- /dev/null +++ b/gptme/codeblock.py @@ -0,0 +1,81 @@ +from collections.abc import Generator +from dataclasses import dataclass +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 + + 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("```"): + 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 + + @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 + + 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/commands.py b/gptme/commands.py index 497057f7..5da97e69 100644 --- a/gptme/commands.py +++ b/gptme/commands.py @@ -2,11 +2,11 @@ import re import sys from collections.abc import Generator -from pathlib import Path from time import sleep from typing import Literal from . import llm +from .codeblock import Codeblock from .constants import CMDFIX from .logmanager import LogManager from .message import ( @@ -18,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 @@ -35,9 +35,6 @@ "fork", "summarize", "context", - "save", - "shell", - "python", "replay", "undo", "impersonate", @@ -54,9 +51,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", "impersonate": "Impersonate the assistant", "tokens": "Show the number of tokens used", @@ -89,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) @@ -124,11 +113,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": @@ -159,17 +143,23 @@ 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 _: - 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 @@ -193,22 +183,6 @@ def edit(log: LogManager) -> Generator[Message, None, None]: # pragma: no cover print("Applied edited messages, write /log to see the result") -def save(log: LogManager, filename: str): - # save the most recent code block to a file - codeblock = log.get_last_code_block() - 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) - 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()) @@ -225,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 echo hello" + yield " /python print('hello')" + 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/llm.py b/gptme/llm.py index 4bb4a845..873974dd 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,13 +93,13 @@ def print_clear(): sys.stdout.flush() # pause inference on finished code-block, letting user run the command before continuing - if codeblocks := 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/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/reduce.py b/gptme/reduce.py index 0608fe08..e1ae9692 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/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/__init__.py b/gptme/tools/__init__.py index a0e0a989..214ed918 100644 --- a/gptme/tools/__init__.py +++ b/gptme/tools/__init__.py @@ -3,19 +3,17 @@ 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 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", @@ -62,6 +57,51 @@ 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: 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: + + + 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.""" @@ -95,12 +135,9 @@ 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): - yield from codeblock_to_tooluse(lang, content).execute(ask) - else: - logger.info(f"Codeblock not supported: {lang}") + yield from execute_codeblock(codeblock, ask) except Exception as e: logger.exception(e) yield Message( @@ -110,97 +147,37 @@ 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 + codeblock: Codeblock, ask: bool ) -> Generator[Message, None, None]: """Executes a codeblock and returns the output.""" - if tool := get_tool_for_codeblock(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) - 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_or_fn: str - content: str - - @classmethod - def from_markdown(cls, content: str) -> "Codeblock": - if content.startswith("```"): - content = content[3:] - if content.endswith("```"): - content = content[:-3] - lang_or_fn = content.splitlines()[0].strip() - return cls(lang_or_fn, content[len(lang_or_fn) :]) - - @property - def is_filename(self) -> bool: - return "." in self.lang_or_fn or "/" in self.lang_or_fn - - @property - def is_supported(self) -> bool: - return is_supported_codeblock_tool(self.lang_or_fn) + 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_codeblock(lang_or_fn: str) -> ToolSpec | None: - block_type = lang_or_fn.split(" ")[0] +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: 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): - 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 is_supported_langtag(lang: str) -> bool: + return bool(get_tool_for_langtag(lang)) def get_tool(tool_name: str) -> ToolSpec: diff --git a/gptme/tools/base.py b/gptme/tools/base.py index 081805ec..1722f77d 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,13 @@ class ToolSpec: execute: ExecuteFunc | None = None block_types: list[str] = field(default_factory=list) available: bool = True + + def get_doc(self, doc="") -> str: + """Returns an updated docstring with examples.""" + if doc: + doc += "\n\n" + if self.examples: + doc += ( + f"# Examples\n\n{transform_examples_to_chat_directives(self.examples)}" + ) + return doc diff --git a/gptme/tools/browser.py b/gptme/tools/browser.py index d746f76f..ae69ae2f 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..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__) @@ -156,8 +155,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 +162,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..c7affb6b 100644 --- a/gptme/tools/patch.py +++ b/gptme/tools/patch.py @@ -1,29 +1,12 @@ """ -Gives the LLM agent the ability to patch files, by using a adapted version git conflict markers. - -Example: - -.. chat:: - - User: patch the file `hello.py` to ask for the name of the user - Assistant: - ```patch hello.py - <<<<<<< ORIGINAL - print("Hello world") - ======= - name = input("What is your name? ") - print(f"hello {name}") - >>>>>>> UPDATED - ``` - System: Patch applied - -Inspired by aider. +Gives the LLM agent the ability to patch text files, by using a adapted version git conflict markers. """ # 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 +15,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" + 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 -``` +''')} +> System: Patch applied -> 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 +157,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..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 @@ -12,7 +13,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 +189,35 @@ 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. +""" + + +# 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) - 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} @@ -206,15 +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 - ) + # create a copy with the updated instructions + return dataclasses.replace(tool, instructions=_instructions) 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/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__) 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__) diff --git a/gptme/tools/tmux.py b/gptme/tools/tmux.py index d7b6bea2..52ebd2e2 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__) 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") 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" diff --git a/tests/test_codeblock.py b/tests/test_codeblock.py new file mode 100644 index 00000000..aba2f98e --- /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) == [ + Codeblock("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) == [ + Codeblock( + "java", + 'public class Main {\n public static void main(String[] args) {\n System.out.println("Hello, Java!");\n }\n}', + ), + Codeblock("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) == [ + Codeblock( + "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) == [ + Codeblock("", 'def hello():\n print("Hello, World!")') + ] diff --git a/tests/test_logmanager.py b/tests/test_logmanager.py index da0229e4..2b1f59a3 100644 --- a/tests/test_logmanager.py +++ b/tests/test_logmanager.py @@ -1,7 +1,8 @@ +from gptme.codeblock import Codeblock 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 +19,7 @@ def test_get_last_code_block(): """, ) ) - assert ("python", "print('world')") == log.get_last_code_block() + assert Codeblock("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!")') - ]