fix: improved main loop logic, squashed bugs

Author: ErikBjare

CONTRIBUTING.md

@@ -7,3 +7,5 @@ Tests are currently covering:
  - tools like shell and Python
  - integration tests that make LLM calls, run the generated code, and checks the output
    - this could be used as a LLM eval harness
+
+There are also some integration tests in `./tests/test-integration.sh` for an alternative way to test.

gptme/cli.py

@@ -171,6 +171,9 @@ def handle_cmd(
             yield msg
             yield from execute_msg(msg, ask=not no_confirm)
         case _:
+            # first undo the '.help' command itself
+            logmanager.undo(1, quiet=True)
+
             print("Available commands:")
             for cmd, desc in action_descriptions.items():
                 print(f"  {cmd}: {desc}")
@@ -282,6 +285,13 @@ def main(
     else:
         promptmsgs = [Message("system", prompt_system)]
 
+    # we need to run this before checking stdin, since the interactive doesn't work with the switch back to interactive mode
+    logfile = get_logfile(
+        name, interactive=(not prompts and interactive) and sys.stdin.isatty()
+    )
+    print(f"Using logdir {logfile.parent}")
+    log = LogManager.load(logfile, initial_msgs=promptmsgs, show_hidden=show_hidden)
+
     # if stdin is not a tty, we're getting piped input
     if not sys.stdin.isatty():
         # fetch prompt from stdin
@@ -293,10 +303,6 @@ def main(
             sys.stdin.close()
             sys.stdin = open("/dev/tty")
 
-    logfile = get_logfile(name, interactive=not prompts and interactive)
-    print(f"Using logdir {logfile.parent}")
-    log = LogManager.load(logfile, initial_msgs=promptmsgs, show_hidden=show_hidden)
-
     # print log
     log.print()
     print("--- ^^^ past messages ^^^ ---")
@@ -318,15 +324,35 @@ def main(
         # if prompts given on cli, insert next prompt into log
         if prompts:
             prompt = prompts.pop(0)
-            log.append(Message("user", prompt))
+            msg = Message("user", prompt)
+            log.append(msg)
+            # if prompt is a user-command, execute it
+            if execute_cmd(msg, log):
+                continue
+        # if prompts exhausted and non-interactive, exit
+        elif not interactive:
+            logger.info("Non-interactive and exhausted prompts, exiting")
+            exit(0)
 
+        # ask for input if no prompt, generate reply, and run tools
         for msg in loop(log, no_confirm, model, llm, stream=stream):
             log.append(msg)
+            # run any user-commands, if msg is from user
+            if msg.role == "user" and execute_cmd(msg, log):
+                break
 
-        # if non-interactive and prompts have been exhausted, exit
-        if not interactive and not prompts:
-            logger.info("Non-interactive and exhausted prompts, exiting")
-            exit(0)
+
+def execute_cmd(msg, log):
+    """Executes any user-command, returns True if command was executed."""
+    assert msg.role == "user"
+
+    # if message starts with ., treat as command
+    # when command has been run,
+    if msg.content[:1] in [".", "$", "/"]:
+        for resp in handle_cmd(msg.content, log, no_confirm=True):
+            log.append(resp)
+        return True
+    return False
 
 
 def loop(
@@ -342,18 +368,6 @@ def loop(
     if log[-1].role == "assistant":
         yield from execute_msg(log[-1], ask=not no_confirm)
 
-    # execute user command
-    if log[-1].role == "user":
-        inquiry = log[-1].content
-        # if message starts with ., treat as command
-        # when command has been run,
-        if inquiry.startswith(".") or inquiry.startswith("$"):
-            yield from handle_cmd(inquiry, log, no_confirm=no_confirm)
-            # we need to re-assign `log` here since it may be replaced by `handle_cmd`
-            # FIXME: this is pretty bad hack to get things working, needs to be refactored
-            if inquiry != ".continue":
-                return
-
     # If last message was a response, ask for input.
     # If last message was from the user (such as from crash/edited log),
     # then skip asking for input and generate response

gptme/prompts.py

@@ -1,15 +1,13 @@
+import logging
 import os
 import shutil
 import subprocess
 from datetime import date
+from typing import Generator
 
 from .config import get_config
 from .message import Message
 
-from typing import Generator
-
-import logging
-
 USER = os.environ.get("USER", None)
 
 
@@ -97,7 +95,8 @@ def initial_prompt(short: bool = False) -> Generator[Message, None, None]:
         yield Message(
             "system",
             """
-You are gptme, a AI assistant powered by large language models that helps the user and can run code and commands on their behalf.
+You are gptme, an AI assistant CLI tool powered powered by large language models that helps the user.
+You can run code and execute terminal commands on their local machine.
 The assistant shows the user to write code, interact with the system, and access the internet. The user will then choose to execute the suggested commands.
 All code should be copy-pasteable or saved, and runnable as-is. Do not use placeholders like `$REPO` unless they have been set.
 When the output of a command is of interest, end the code block so that the user can execute it before continuing.

gptme/tools/useredit.py

@@ -35,7 +35,6 @@ def edit_text_with_editor(initial_text: str, ext=None) -> str:
     # Check that the user actually edited the file.
     if edited_text == initial_text:
         logger.info("No changes made, exiting.")
-        print("No changes made, exiting.")
 
     # Delete the temporary file.
     os.remove(temp_filename)

tests/test-integration.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+
+set -e
+set -x
+
+# set this to indicate tests are run (non-interactive)
+export PYTEST_CURRENT_TEST=1
+
+# test stdin and cli-provided prompt
+echo "The project mascot is a flying pig" | gptme "What is the project mascot?"
+
+# test python command
+gptme ".python print('hello world')"
+
+# test shell command
+gptme ".shell echo 'hello world'"