feat: many changes, now supports local inference via llama_cpp.server, some refactoring, improved README

Author: ErikBjare

.gitignore

@@ -3,3 +3,4 @@
 *.log
 __pycache__
 .coverage
+*logs*

README.md

@@ -8,14 +8,41 @@ Just me playing with large language models, langchain, etc.
 
 ## gptme
 
-An interactive CLI to let you chat with ChatGPT, with extra tools like:
+An interactive CLI to let you interact with LLMs in a Chat-style interface.
 
- - Execute shell/Python code on the local machine.
-   - Command output (stdout & stderr + error code) will be feeded back to the agent, making it able to self-correct errors etc.
- - Handle long context sizes through summarization.
+With **features** like:
+
+ - Supports OpenAI and **any model that runs in llama**
+   - Thanks to llama-cpp-server!
+ - Tools
+   - Access to the local machine
+   - Execute shell/Python code on the local machine.
+     - Command output (stdout & stderr + error code) will be feeded back to the agent, making it able to self-correct errors etc.
+ - Can handle long context sizes through summarization.
    - (not very well developed)
 
 
+### Usage
+
+Install deps:
+
+```sh
+poetry install
+```
+
+To use locally, you need to start llama-cpp-server:
+
+```sh
+poetry run python -m llama_cpp.server --model ~/ML/Manticore-13B.ggmlv3.q4_1.bin
+```
+
+Then you can interact with it using:
+```sh
+gptme --llm llama
+
+```
+
+
 ## TODO
 
 Ideas for things to try:

gpt_playground/shell.py

@@ -3,23 +3,33 @@
 """
 
 from copy import copy
-import openai
+import subprocess
+import sys
 
 import click
+import openai
+
+
+# TODO: use this generation in gptme
+if sys.platform == "linux":
+    system_info = """
+    OS: Arch Linux
+    """
+elif sys.platform == "darwin":
+    system_info = f"""
+    $ uname -a
+    {subprocess.call(["uname", "-a"])}
+    Darwin erb-m2.localdomain 21.6.0 Darwin Kernel Version 21.6.0: Sat Jun 18 17:07:28 PDT 2022; root:xnu-8020.140.41~1/RELEASE_ARM64_T8110 arm64 arm Darwin
+    $ sw_vers
+    {subprocess.call(["sw_vers"])}
+    ProductName:	macOS
+    ProductVersion:	12.5
+    BuildVersion:	21G72
+    """
+else:
+    system_info = "Unknown/unsupported OS (Windows?)"
 
-# TODO: Generate automatically
-system_info_arch = """
-OS: Arch Linux
-"""
 
-system_info_macos = """
-$ uname -a
-Darwin erb-m2.localdomain 21.6.0 Darwin Kernel Version 21.6.0: Sat Jun 18 17:07:28 PDT 2022; root:xnu-8020.140.41~1/RELEASE_ARM64_T8110 arm64 arm Darwin
-$ sw_vers
-ProductName:	macOS
-ProductVersion:	12.5
-BuildVersion:	21G72
-"""
 
 initial_messages = [
     {
@@ -29,7 +39,7 @@
      },
     {
         "role": "system",
-        "content": """First, let's check the environment:\n"""  + system_info_macos
+        "content": """First, let's check the environment:\n"""  + system_info
     }
 ]
 

gptme/cli.py

@@ -1,48 +1,67 @@
 """
-This is a long-living agent that is designed to be a companion to the user.
+GPTMe
+=====
+
+This is a long-living AI language model called GPTMe, it is designed to be a helpful companion.
 
 It should be able to help the user in various ways, such as:
 
- - Acting as an executive assistant
- - Answering questions
- - Helping strategize
- - Giving advice
  - Writing code
+ - Using the shell
+ - Assisting with technical tasks
  - Writing prose (such as email, code docs, etc.)
- - Providing companionship
+ - Acting as an executive assistant
 
 The agent should be able to learn from the user and adapt to their needs.
-The agent should try to always output information using markdown formatting, preferably using GitHub Flavored Markdown.
+The agent should always output information using GitHub Flavored Markdown.
+THe agent should always output code and commands in markdown code blocks with the appropriate language tag.
 
 Since the agent is long-living, it should be able to remember things that the user has told it,
 to do so, it needs to be able to store and query past conversations in a database.
 """
+# The above docstring is the first message that the agent will see.
 
 from typing import Literal, Generator
 from datetime import datetime
 import logging
 import os
 import sys
 import shutil
+import readline  # noqa: F401
 import itertools
 from pathlib import Path
 
-from termcolor import colored
+from termcolor import colored  # type: ignore
 import openai
 import click
 
-import typing
 
 from .constants import role_color
-from .tools import _execute_linecmd, _execute_codeblock, _execute_save, _execute_shell, _execute_python
+from .tools import (
+    _execute_linecmd,
+    _execute_codeblock,
+    _execute_save,
+    _execute_shell,
+    _execute_python,
+)
 from .util import msgs2dicts
 from .message import Message
 from .logmanager import LogManager
+from .prompts import initial_prompt
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
 
+LLMChoice = Literal["openai", "llama"]
+
+readline.add_history("What is love?")
+readline.add_history("Have you heard about an open-source app called ActivityWatch?")
+readline.add_history(
+    "Explain the 'Attention is All You Need' paper in the style of Andrej Karpathy."
+)
+
+
 def get_logfile(logdir: str) -> str:
     logdir = logdir + "/"
     if not os.path.exists(logdir):
@@ -106,28 +125,52 @@ def handle_cmd(cmd: str, logmanager: LogManager) -> Generator[Message, None, Non
             sys.exit(0)
         case _:
             print("Available commands:")
-            for cmd in typing.get_args(Actions):
-                desc = action_descriptions.get(cmd, default="missing description")
+            for cmd, desc in action_descriptions.items():
                 print(f"  {cmd}: {desc}")
 
 
 @click.group()
 def cli():
     pass
 
+
 script_path = Path(os.path.realpath(__file__))
 
+
 @cli.command()
-@click.argument("command" , default=None, required=False)
+@click.argument("command", default=None, required=False)
+@click.option(
+    "--logs",
+    default=script_path.parent.parent / "logs",
+    help="Folder where conversation logs are stored",
+)
+@click.option("--llm", default="openai", help="LLM to use")
 @click.option(
-    "--logs", default=script_path.parent.parent / "logs", help="Folder where conversation logs are stored"
+    "--stream",
+    is_flag=True,
+    default=True,
+    help="Wether to use streaming (only supported for openai atm)",
 )
-def main(command: str | None, logs: str):
-    """Main interactivity loop."""
+@click.option(
+    "--prompt",
+    default="short",
+    help="Can be 'short', 'full', or a custom prompt",
+)
+def main(command: str | None, logs: str, llm: LLMChoice, stream: bool, prompt: str):
+    """
+    GPTMe, a CLI interface for LLMs.
+    """
     openai.api_key = os.environ["OPENAI_API_KEY"]
+    openai.api_base = "http://localhost:8000/v1"
 
+    if prompt in ["full", "short"]:
+        promptmsgs = initial_prompt(short=prompt == "short")
+    else:
+        promptmsgs = [Message("system", prompt)]
+
+    print(f"Using logdir {logs}")
     logfile = get_logfile(logs)
-    logmanager = LogManager.load(logfile)
+    logmanager = LogManager.load(logfile, initial_msgs=promptmsgs)
     logmanager.print()
     print("--- ^^^ past messages ^^^ ---")
 
@@ -143,53 +186,117 @@ def main(command: str | None, logs: str):
     while True:
         # if non-interactive command given on cli, exit
         if command_triggered:
+            print("Command triggered, exiting")
             break
 
         # If last message was a response, ask for input.
-        # If last message was from the user (such as from crash/edited log), 
+        # If last message was from the user (such as from crash/edited log),
         # then skip asking for input and generate response
-        if log[-1].role in ["system", "assistant"]:
-            prompt = colored("User", role_color["user"]) + ": "
+        last_msg = log[-1] if log else None
+        if not last_msg or (
+            (last_msg.role in ["system", "assistant"])
+            or (log[-1].role == "user" and log[-1].content.startswith("."))
+        ):
+            inquiry = prompt_user(command)
             if command:
-                print(prompt + command)
-                inquiry = command
                 command = None
                 command_triggered = True
-            else:
-                inquiry = input(prompt)
-    
+
             if not inquiry:
+                print("Continue 1 (rare!)")
                 continue
             logmanager.append(Message("user", inquiry))
 
         assert log[-1].role == "user"
         inquiry = log[-1].content
         # if message starts with ., treat as command
-        # when command has been run, 
+        # when command has been run,
         if inquiry.startswith("."):
             for msg in handle_cmd(inquiry, logmanager):
                 logmanager.append(msg)
+            if command:
+                command_triggered = True
+                print("Continue 2")
             continue
 
         # if large context, try to reduce/summarize
         # print response
-        msg_response = reply(logmanager.prepare_messages())
+        try:
+            msg_response = reply(logmanager.prepare_messages(), stream)
+
+            # log response and run tools
+            if msg_response:
+                for msg in itertools.chain([msg_response], execute_msg(msg_response)):
+                    logmanager.append(msg)
+        except KeyboardInterrupt:
+            print("Interrupted")
+
+
+def prompt_user(value=None) -> str:
+    return prompt_input(colored("User", role_color["user"]) + ": ", value)
+
+
+def prompt_input(prompt: str, value=None) -> str:
+    if value:
+        print(prompt + value)
+    else:
+        value = input(prompt)
+    return value
+
 
-        # log response and run tools
-        for msg in itertools.chain([msg_response], execute_msg(msg_response)):
-            logmanager.append(msg)
+def reply(messages: list[Message], stream: bool = False) -> Message:
+    if stream:
+        return reply_stream(messages)
+    else:
+        prefix = colored("Assistant", "green", attrs=["bold"])
+        print(f"{prefix}: Thinking...", end="\r")
+        response = _chat_complete(messages)
+        print(" " * shutil.get_terminal_size().columns, end="\r")
+        return Message("assistant", response)
 
 
-def reply(messages: list[Message]) -> Message:
-    # print in-progress indicator
-    print(colored("Assistant", "green", attrs=["bold"]) + ": Thinking...", end="\r")
-    response = openai.ChatCompletion.create(
+def _chat_complete(messages: list[Message]) -> str:
+    response = openai.ChatCompletion.create(  # type: ignore
         model="gpt-3.5-turbo",
         messages=msgs2dicts(messages),
         temperature=0,
     )
-    print(" " * shutil.get_terminal_size().columns, end="\r")
-    return Message("assistant", response.choices[0].message.content)
+    return response.choices[0].message.content
+
+
+def reply_stream(messages: list[Message]) -> Message:
+    prefix = colored("Assistant", "green", attrs=["bold"])
+    print(f"{prefix}: Thinking...", end="\r")
+    response = openai.ChatCompletion.create(  # type: ignore
+        model="gpt-3.5-turbo",
+        messages=msgs2dicts(messages),
+        temperature=0,
+        stream=True,
+        max_tokens=1000,
+    )
+
+    def deltas_to_str(deltas: list[dict]):
+        return "".join([d.get("content", "") for d in deltas])
+
+    def print_clear():
+        print(" " * shutil.get_terminal_size().columns, end="\r")
+
+    deltas: list[dict] = []
+    print_clear()
+    print(f"{prefix}: ", end="")
+    stop_reason = None
+    for chunk in response:
+        delta = chunk["choices"][0]["delta"]
+        deltas.append(delta)
+        stop_reason = chunk["choices"][0]["finish_reason"]
+        print(deltas_to_str([delta]), end="")
+        # need to flush stdout to get the print to show up
+        sys.stdout.flush()
+    print_clear()
+    verbose = True
+    if verbose:
+        print(f" - Stop reason: {stop_reason}")
+    return Message("assistant", deltas_to_str(deltas))
 
 
 if __name__ == "__main__":