From 6eb763c6aeb4e127d27a02d65898c8cf37d9c2a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Wed, 11 Oct 2023 19:50:41 +0200 Subject: [PATCH] feat: added rename and fork commands, refactor commands and tabcomplete --- gptme/cli.py | 116 ++++--------------------------------------- gptme/commands.py | 42 ++++++++++++++++ gptme/logmanager.py | 13 ++++- gptme/tabcomplete.py | 75 ++++++++++++++++++++++++++++ 4 files changed, 140 insertions(+), 106 deletions(-) create mode 100644 gptme/commands.py create mode 100644 gptme/tabcomplete.py diff --git a/gptme/cli.py b/gptme/cli.py index 17e75bd5..2a4f8b15 100644 --- a/gptme/cli.py +++ b/gptme/cli.py @@ -20,7 +20,6 @@ to do so, it needs to be able to store and query past conversations in a database. """ # The above may be used as a prompt for the agent. - import atexit import importlib.metadata import io @@ -29,7 +28,6 @@ import readline # noqa: F401 import sys from datetime import datetime -from functools import lru_cache from pathlib import Path from time import sleep from typing import Generator, Literal @@ -40,6 +38,7 @@ from rich import print # noqa: F401 from rich.console import Console +from .commands import CMDFIX, action_descriptions from .constants import HISTORY_FILE, LOGSDIR, PROMPT_USER from .llm import init_llm, reply from .logmanager import LogManager @@ -50,6 +49,7 @@ toml_to_msgs, ) from .prompts import initial_prompt_single_message +from .tabcomplete import register_tabcomplete from .tools import execute_msg, execute_python, execute_shell from .tools.shell import get_shell from .tools.summarize import summarize @@ -62,43 +62,6 @@ LLMChoice = Literal["openai", "llama"] ModelChoice = Literal["gpt-3.5-turbo", "gpt4"] -CMDFIX = "/" # prefix for commands, e.g. /help - -Actions = Literal[ - "continue", - "summarize", - "log", - "edit", - "summarize", - "context", - "load", - "save", - "shell", - "python", - "replay", - "undo", - "impersonate", - "help", - "exit", -] - -action_descriptions: dict[Actions, str] = { - "continue": "Continue response", - "undo": "Undo the last action", - "log": "Show the conversation log", - "edit": "Edit previous messages", - "summarize": "Summarize the conversation so far", - "load": "Load a file", - "save": "Save the most recent code block to a file", - "shell": "Execute a shell command", - "python": "Execute a Python command", - "replay": "Re-execute past commands in the conversation (does not store output in log)", - "impersonate": "Impersonate the assistant", - "help": "Show this help message", - "exit": "Exit the program", -} -COMMANDS = list(action_descriptions.keys()) - def handle_cmd( cmd: str, log: LogManager, no_confirm: bool @@ -118,6 +81,14 @@ def handle_cmd( case "log": log.undo(1, quiet=True) log.print(show_hidden="--hidden" in args) + case "rename": + # rename the conversation + new_name = args[0] if args else input("New name: ") + log.rename(new_name) + case "fork": + # fork the conversation + new_name = args[0] if args else input("New name: ") + log.fork(new_name) case "summarize": msgs = log.prepare_messages() msgs = [m for m in msgs if not m.hide] @@ -293,6 +264,7 @@ def main( # init logger.debug("Started") load_dotenv() + register_tabcomplete() _load_readline_history() init_llm(llm) # set up API_KEY and API_BASE @@ -486,58 +458,6 @@ def _gen_context_msg() -> Message: ] -def _completer(text: str, state: int) -> str | None: - """ - Tab completion for readline. - - Completes /commands and paths in arguments. - - The completer function is called as function(text, state), for state in 0, 1, 2, …, until it returns a non-string value. - It should return the next possible completion starting with text. - """ - return _matches(text)[state] - - -@lru_cache(maxsize=1) -def _matches(text: str) -> list[str]: - """Returns a list of matches for text to complete.""" - - # if text starts with /, complete with commands or files as absolute paths - if text.startswith("/"): - # if no text, list all commands - all_commands = [f"{CMDFIX}{cmd}" for cmd in COMMANDS if cmd != "help"] - if not text[1:]: - return all_commands - # else, filter commands with text - else: - matching_files = [str(p) for p in Path("/").glob(text[1:] + "*")] - return [ - cmd for cmd in all_commands if cmd.startswith(text) - ] + matching_files - - # if text starts with ., complete with current dir - elif text.startswith("."): - if not text[1:]: - return [str(Path.cwd())] - else: - all_files = [str(p) for p in Path.cwd().glob("*")] - return [f for f in all_files if f.startswith(text)] - - # if text starts with ../, complete with parent dir - elif text.startswith(".."): - if not text[2:]: - return [str(Path.cwd().parent)] - else: - return [str(p) for p in Path.cwd().parent.glob(text[2:] + "*")] - - # else, complete with files in current dir - else: - if not text: - return [str(Path.cwd())] - else: - return [str(p) for p in Path.cwd().glob(text + "*")] - - def _load_readline_history() -> None: logger.debug("Loading history") # enabled by default in CPython, make it explicit @@ -550,20 +470,6 @@ def _load_readline_history() -> None: for line in history_examples: readline.add_history(line) - # set up tab completion - print("Setting up tab completion") - readline.set_completer(_completer) - readline.set_completer_delims(" ") - readline.parse_and_bind("tab: complete") - - # https://github.com/python/cpython/issues/102130#issuecomment-1439242363 - if "libedit" in readline.__doc__: # type: ignore - print("Found libedit readline") - readline.parse_and_bind("bind ^I rl_complete") - else: - print("Found gnu readline") - readline.parse_and_bind("tab: complete") - atexit.register(readline.write_history_file, HISTORY_FILE) diff --git a/gptme/commands.py b/gptme/commands.py new file mode 100644 index 00000000..36503509 --- /dev/null +++ b/gptme/commands.py @@ -0,0 +1,42 @@ +from typing import Literal + +CMDFIX = "/" # prefix for commands, e.g. /help + +Actions = Literal[ + "continue", + "summarize", + "log", + "edit", + "rename", + "fork", + "summarize", + "context", + "load", + "save", + "shell", + "python", + "replay", + "undo", + "impersonate", + "help", + "exit", +] + +action_descriptions: dict[Actions, str] = { + "continue": "Continue response", + "undo": "Undo the last action", + "log": "Show the conversation log", + "edit": "Edit previous messages", + "rename": "Rename the conversation", + "fork": "Create a copy of the conversation with a new name", + "summarize": "Summarize the conversation so far", + "load": "Load a file", + "save": "Save the most recent code block to a file", + "shell": "Execute a shell command", + "python": "Execute a Python command", + "replay": "Re-execute past commands in the conversation (does not store output in log)", + "impersonate": "Impersonate the assistant", + "help": "Show this help message", + "exit": "Exit the program", +} +COMMANDS = list(action_descriptions.keys()) diff --git a/gptme/logmanager.py b/gptme/logmanager.py index 07bbd98b..6f3c5108 100644 --- a/gptme/logmanager.py +++ b/gptme/logmanager.py @@ -30,7 +30,7 @@ def __init__( fpath = NamedTemporaryFile(delete=False).name print(f"[yellow]No logfile specified, using tmpfile {fpath}.[/]") logfile = Path(fpath) - self.logfile = logfile + self.logfile = logfile if isinstance(logfile, Path) else Path(logfile) self.show_hidden = show_hidden # TODO: Check if logfile has contents, then maybe load, or should it overwrite? @@ -122,6 +122,17 @@ def get_last_code_block(self) -> str | None: return msg.content.split("```")[-2].split("\n", 1)[-1] return None + def rename(self, name: str) -> None: + # rename the conversation and log file + # if you want to keep the old log, use fork() + self.logfile.rename(self.logfile.parent / f"{name}.log") + self.logfile = self.logfile.parent / f"{name}.log" + + def fork(self, name: str) -> None: + # save and switch to a new log file without renaming the old one + self.write() + self.logfile = self.logfile.parent / f"{name}.log" + def write_log(msg_or_log: Message | list[Message], logfile: PathLike) -> None: """ diff --git a/gptme/tabcomplete.py b/gptme/tabcomplete.py new file mode 100644 index 00000000..e8dcf43f --- /dev/null +++ b/gptme/tabcomplete.py @@ -0,0 +1,75 @@ +import readline +from functools import lru_cache +from pathlib import Path + +from .commands import CMDFIX, COMMANDS + + +def register_tabcomplete() -> None: + """Register tab completion for readline.""" + + # set up tab completion + print("Setting up tab completion") + readline.set_completer(_completer) + readline.set_completer_delims(" ") + readline.parse_and_bind("tab: complete") + + # https://github.com/python/cpython/issues/102130#issuecomment-1439242363 + if "libedit" in readline.__doc__: # type: ignore + print("Found libedit readline") + readline.parse_and_bind("bind ^I rl_complete") + else: + print("Found gnu readline") + readline.parse_and_bind("tab: complete") + + +def _completer(text: str, state: int) -> str | None: + """ + Tab completion for readline. + + Completes /commands and paths in arguments. + + The completer function is called as function(text, state), for state in 0, 1, 2, …, until it returns a non-string value. + It should return the next possible completion starting with text. + """ + return _matches(text)[state] + + +@lru_cache(maxsize=1) +def _matches(text: str) -> list[str]: + """Returns a list of matches for text to complete.""" + + # if text starts with /, complete with commands or files as absolute paths + if text.startswith("/"): + # if no text, list all commands + all_commands = [f"{CMDFIX}{cmd}" for cmd in COMMANDS if cmd != "help"] + if not text[1:]: + return all_commands + # else, filter commands with text + else: + matching_files = [str(p) for p in Path("/").glob(text[1:] + "*")] + return [ + cmd for cmd in all_commands if cmd.startswith(text) + ] + matching_files + + # if text starts with ., complete with current dir + elif text.startswith("."): + if not text[1:]: + return [str(Path.cwd())] + else: + all_files = [str(p) for p in Path.cwd().glob("*")] + return [f for f in all_files if f.startswith(text)] + + # if text starts with ../, complete with parent dir + elif text.startswith(".."): + if not text[2:]: + return [str(Path.cwd().parent)] + else: + return [str(p) for p in Path.cwd().parent.glob(text[2:] + "*")] + + # else, complete with files in current dir + else: + if not text: + return [str(Path.cwd())] + else: + return [str(p) for p in Path.cwd().glob(text + "*")]