feat: added branching to logmanager (#33)

* feat: added branching to logmanager, to allow tree-structured conversations

* feat: more progress on branching, added support for switching branches in web UI

* fix: added dirs.py to better handle directories and paths, misc fixes

* build: fixed Makefile EXCLUDES for SRCFILES

* fix: fixed bug in commands.py

* feat: added ability to send messages to non-main branch in web UI

Makefile

@@ -9,7 +9,7 @@ SRCFILES = $(shell find ${SRCDIRS} -name '*.py')
 
 # exclude files
 EXCLUDES = tests/output scripts/build_changelog.py
-SRCFILES = $(shell find ${SRCDIRS} -name '*.py' -not -path ${EXCLUDES})
+SRCFILES = $(shell find ${SRCDIRS} -name '*.py' $(foreach EXCLUDE,$(EXCLUDES),-not -path $(EXCLUDE)))
 
 # Check if playwright is installed (for browser tests)
 HAS_PLAYWRIGHT := $(shell poetry run command -v playwright 2> /dev/null)

gptme/cli.py

@@ -19,12 +19,8 @@
 from rich.console import Console
 
 from .commands import CMDFIX, action_descriptions, execute_cmd
-from .constants import (
-    HISTORY_FILE,
-    LOGSDIR,
-    MULTIPROMPT_SEPARATOR,
-    PROMPT_USER,
-)
+from .constants import MULTIPROMPT_SEPARATOR, PROMPT_USER
+from .dirs import get_logs_dir, get_readline_history_file
 from .llm import init_llm, reply
 from .logmanager import LogManager, _conversations
 from .message import Message
@@ -135,12 +131,11 @@ def main(
     stream: bool,
     verbose: bool,
     no_confirm: bool,
-    show_hidden: bool,
     interactive: bool,
+    show_hidden: bool,
     version: bool,
 ):
     """Main entrypoint for the CLI."""
-
     if version:
         # print version and exit
         print_builtin(f"gptme {importlib.metadata.version('gptme-python')}")
@@ -280,13 +275,14 @@ def loop(
 
 def get_name(name: str) -> Path:
     datestr = datetime.now().strftime("%Y-%m-%d")
+    logsdir = get_logs_dir()
 
     # returns a name for the new conversation
     if name == "random":
         # check if name exists, if so, generate another one
         for _ in range(3):
             name = generate_name()
-            logpath = LOGSDIR / f"{datestr}-{name}"
+            logpath = logsdir / f"{datestr}-{name}"
             if not logpath.exists():
                 break
         else:
@@ -296,7 +292,7 @@ def get_name(name: str) -> Path:
             # ask for name, or use random name
             name = input("Name for conversation (or empty for random words): ")
             name = f"{datestr}-{name}"
-            logpath = LOGSDIR / name
+            logpath = logsdir / name
 
             # check that name is unique/doesn't exist
             if not logpath.exists():
@@ -309,7 +305,7 @@ def get_name(name: str) -> Path:
             datetime.strptime(name[:10], "%Y-%m-%d")
         except ValueError:
             name = f"{datestr}-{name}"
-        logpath = LOGSDIR / name
+        logpath = logsdir / name
     return logpath
 
 
@@ -330,13 +326,14 @@ def _load_readline_history() -> None:
     readline.set_auto_history(True)
     # had some bugs where it grew to gigs, which should be fixed, but still good precaution
     readline.set_history_length(100)
+    history_file = get_readline_history_file()
     try:
-        readline.read_history_file(HISTORY_FILE)
+        readline.read_history_file(history_file)
     except FileNotFoundError:
         for line in history_examples:
             readline.add_history(line)
 
-    atexit.register(readline.write_history_file, HISTORY_FILE)
+    atexit.register(readline.write_history_file, history_file)
 
 
 def get_logfile(name: str, interactive=True) -> Path:
@@ -351,7 +348,7 @@ def is_test(name: str) -> bool:
 
     # filter out test conversations
     # TODO: save test convos to different folder instead
-    prev_conv_files = [f for f in prev_conv_files if not is_test(f.parent.name)]
+    # prev_conv_files = [f for f in prev_conv_files if not is_test(f.parent.name)]
 
     NEWLINE = "\n"
     prev_convs = [
@@ -370,7 +367,7 @@ def is_test(name: str) -> bool:
         if index == 0:
             logdir = get_name(name)
         else:
-            logdir = LOGSDIR / prev_conv_files[index - 1].parent
+            logdir = get_logs_dir() / prev_conv_files[index - 1].parent
     else:
         logdir = get_name(name)
 
@@ -474,6 +471,7 @@ def _parse_prompt(prompt: str) -> str:
                 "Failed to import browser tool, skipping URL expansion."
                 "You might have to install browser extras."
             )
+            continue
 
         try:
             content = read_url(url)

gptme/commands.py

@@ -1,4 +1,5 @@
 import logging
+import re
 import sys
 from collections.abc import Generator
 from pathlib import Path
@@ -18,7 +19,7 @@
 from .tools.context import gen_context_msg
 from .tools.summarize import summarize
 from .tools.useredit import edit_text_with_editor
-from .util import len_tokens
+from .util import ask_execute, len_tokens
 
 logger = logging.getLogger(__name__)
 
@@ -79,12 +80,13 @@ def handle_cmd(
     """Handles a command."""
     cmd = cmd.lstrip(CMDFIX)
     logger.debug(f"Executing command: {cmd}")
-    name, *args = cmd.split(" ")
+    name, *args = re.split(r"[\n\s]", cmd)
+    full_args = cmd.split(" ", 1)[1] if " " in cmd else ""
     match name:
         case "bash" | "sh" | "shell":
-            yield from execute_shell(" ".join(args), ask=not no_confirm)
+            yield from execute_shell(full_args, ask=not no_confirm)
         case "python" | "py":
-            yield from execute_python(" ".join(args), ask=not no_confirm)
+            yield from execute_python(full_args, ask=not no_confirm)
         case "log":
             log.undo(1, quiet=True)
             log.print(show_hidden="--hidden" in args)
@@ -94,7 +96,7 @@ def handle_cmd(
             # rename the conversation
             print("Renaming conversation (enter empty name to auto-generate)")
             new_name = args[0] if args else input("New name: ")
-            rename(log, new_name)
+            rename(log, new_name, ask=not no_confirm)
         case "fork":
             # fork the conversation
             new_name = args[0] if args else input("New name: ")
@@ -134,7 +136,7 @@ def handle_cmd(
                     for msg in execute_msg(msg, ask=True):
                         print_msg(msg, oneline=False)
         case "impersonate":
-            content = " ".join(args) if args else input("[impersonate] Assistant: ")
+            content = full_args if full_args else input("[impersonate] Assistant: ")
             msg = Message("assistant", content)
             yield msg
             yield from execute_msg(msg, ask=not no_confirm)
@@ -165,8 +167,7 @@ def edit(log: LogManager) -> Generator[Message, None, None]:  # pragma: no cover
             except KeyboardInterrupt:
                 yield Message("system", "Interrupted")
                 return
-    log.log = list(reversed(res))
-    log.write()
+    log.edit(list(reversed(res)))
     # now we need to redraw the log so the user isn't seeing stale messages in their buffer
     # log.print()
     print("Applied edited messages, write /log to see the result")
@@ -179,23 +180,24 @@ def save(log: LogManager, filename: str):
         print("No code block found")
         return
     if Path(filename).exists():
-        ans = input("File already exists, overwrite? [y/N] ")
-        if ans.lower() != "y":
+        confirm = ask_execute("File already exists, overwrite?", default=False)
+        if not confirm:
             return
     with open(filename, "w") as f:
         f.write(code)
     print(f"Saved code block to {filename}")
 
 
-def rename(log: LogManager, new_name: str):
+def rename(log: LogManager, new_name: str, ask: bool = True):
     if new_name in ["", "auto"]:
         new_name = llm.generate_name(log.prepare_messages())
         assert " " not in new_name
         print(f"Generated name: {new_name}")
-        confirm = input("Confirm? [y/N] ")
-        if confirm.lower() not in ["y", "yes"]:
-            print("Aborting")
-            return
+        if ask:
+            confirm = ask_execute("Confirm?")
+            if not confirm:
+                print("Aborting")
+                return
         log.rename(new_name, keep_date=True)
     else:
         log.rename(new_name, keep_date=False)

gptme/constants.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 # prefix for commands, e.g. /help
 CMDFIX = "/"
 
@@ -23,15 +21,3 @@
 PROMPT_ASSISTANT = (
     f"[bold {ROLE_COLOR['assistant']}]Assistant[/bold {ROLE_COLOR['assistant']}]"
 )
-
-# Config
-CONFIG_PATH = Path("~/.config/gptme").expanduser()
-HISTORY_FILE = CONFIG_PATH / "history"
-
-# Data
-DATADIR = Path("~/.local/share/gptme").expanduser()
-LOGSDIR = DATADIR / "logs"
-
-# create all paths
-for path in [CONFIG_PATH, DATADIR, LOGSDIR]:
-    path.mkdir(parents=True, exist_ok=True)

gptme/dirs.py

@@ -0,0 +1,43 @@
+import os
+from pathlib import Path
+
+from platformdirs import user_config_dir, user_data_dir
+
+
+def get_config_dir() -> Path:
+    return Path(user_config_dir("gptme"))
+
+
+def get_readline_history_file() -> Path:
+    # TODO: move to data dir
+    return get_config_dir() / "history"
+
+
+def get_data_dir() -> Path:
+    # used in testing, so must take precedence
+    if "XDG_DATA_HOME" in os.environ:
+        return Path(os.environ["XDG_DATA_HOME"]) / "gptme"
+
+    # just a workaround for me personally
+    old = Path("~/.local/share/gptme").expanduser()
+    if old.exists():
+        return old
+
+    return Path(user_data_dir("gptme"))
+
+
+def get_logs_dir() -> Path:
+    """Get the path for **conversation logs** (not to be confused with the logger file)"""
+    path = get_data_dir() / "logs"
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def _init_paths():
+    # create all paths
+    for path in [get_config_dir(), get_data_dir(), get_logs_dir()]:
+        path.mkdir(parents=True, exist_ok=True)
+
+
+# run once on init
+_init_paths()