Add slash command and editing with an arbitrary editor

4 files changed, 140 insertions, 61 deletions

diff --git a/src/gpt_chat_cli/argparsing.py b/src/gpt_chat_cli/argparsing.py
index a96e81c..703a13b 100644
--- a/src/gpt_chat_cli/argparsing.py
+++ b/src/gpt_chat_cli/argparsing.py
@@ -124,14 +124,14 @@ def _construct_parser() \
         "-n",
         "--n-completions",
         type=int,
-        default=os.getenv('f{_GPT_CLI_ENV_PREFIX}N_COMPLETIONS', 1),
+        default=os.getenv(f'{_GPT_CLI_ENV_PREFIX}N_COMPLETIONS', 1),
         help="How many chat completion choices to generate for each input message.",
     )
 
     parser.add_argument(
         "--system-message",
         type=str,
-        default=os.getenv('f{_GPT_CLI_ENV_PREFIX}SYSTEM_MESSAGE'),
+        default=os.getenv(f'{_GPT_CLI_ENV_PREFIX}SYSTEM_MESSAGE'),
         help="Specify an alternative system message.",
     )
 
@@ -174,6 +174,13 @@ def _construct_parser() \
         help="Start an interactive session"
     )
 
+    parser.add_argument(
+        "--interactive-editor",
+        type=str,
+        default=os.getenv(f'{_GPT_CLI_ENV_PREFIX}INTERACTIVE_EDITOR'),
+        help="The editor which is launched by default using the /edit command in interactive mode"
+    )
+
     initial_prompt = parser.add_mutually_exclusive_group()
 
     initial_prompt.add_argument(
diff --git a/src/gpt_chat_cli/argvalidation.py b/src/gpt_chat_cli/argvalidation.py
index 16987a9..8e6ef52 100644
--- a/src/gpt_chat_cli/argvalidation.py
+++ b/src/gpt_chat_cli/argvalidation.py
@@ -53,6 +53,7 @@ class Arguments:
     openai_key: str
     system_message: Optional[str] = None
     debug_args: Optional[DebugArguments] = None
+    interactive_editor: Optional[str] = None
 
 def post_process_raw_args(raw_args : RawArguments) -> Arguments:
     _populate_defaults(raw_args)
@@ -104,7 +105,8 @@ def _restructure_arguments(raw_args : RawArguments) -> Arguments:
            version=args.version,
            list_models=args.list_models,
            interactive=args.interactive,
-           system_message=args.system_message
+           system_message=args.system_message,
+           interactive_editor=args.interactive_editor
     )
 
 def _die_validation_err(err : str):
diff --git a/src/gpt_chat_cli/cmd.py b/src/gpt_chat_cli/cmd.py
index 86ce0e6..899c705 100644
--- a/src/gpt_chat_cli/cmd.py
+++ b/src/gpt_chat_cli/cmd.py
@@ -10,6 +10,10 @@ from collections import defaultdict
 from dataclasses import dataclass
 from typing import Tuple, Optional
 
+import subprocess
+import tempfile
+import os
+
 from .openai_wrappers import (
     create_chat_completion,
     OpenAIChatResponse,
@@ -30,8 +34,13 @@ from .argvalidation import (
 )
 
 from .version import VERSION
-from .color import get_color_codes
+from .color import (
+    get_color_codes,
+    surround_ansi_escapes
+)
+
 from .chat_colorizer import ChatColorizer
+from .prompt import Prompter
 
 ###########################
 ####      UTILS        ####
@@ -63,11 +72,56 @@ def get_system_message(system_message : Optional[str]):
 
     return ChatMessage(Role.SYSTEM, system_message)
 
-def enable_emacs_editing():
-    try:
-        import readline
-    except ImportError:
-        pass
+def _resolve_system_editor() -> Optional[str]:
+    '''
+        Attempt to resolve the system if one is not explicitly specified.
+        This is either the editor stored in the EDITOR environmental variable
+        or one of editor, vim, emacs, vi, or nano.
+    '''
+
+    fallback_editors = [
+        "editor", # debian default
+        "vim",
+        "emacs",
+        "vi",
+        "nano",
+    ]
+
+    editor = os.getenv("EDITOR")
+
+    if editor:
+        return editor
+
+    paths = os.getenv("PATH")
+
+    if not paths:
+        return None
+
+    for editor in fallback_editors:
+        for path in paths.split(os.pathsep):
+            if os.path.exists(os.path.join(path, editor)):
+                return editor
+
+    return None
+
+def _launch_interactive_editor(editor : Optional[str] = None) -> str:
+
+    with tempfile.NamedTemporaryFile(suffix="gpt.msg") as tmp_file:
+
+        editor = editor or _resolve_system_editor()
+
+        try:
+            subprocess.call([editor, tmp_file.name])
+        except FileNotFoundError:
+            print(f'error: the specified editor \"{editor}\" does not exist')
+            sys.exit(1)
+
+        # Read the resulting file into a string
+        with open(tmp_file.name, "r") as edited_file:
+            edited_content = edited_file.read()
+
+        return edited_content
+
 
 ###########################
 ####   SAVE / REPLAY   ####
@@ -212,26 +266,6 @@ def print_streamed_response(
 ####    COMMANDS     ####
 #########################
 
-def surround_ansi_escapes(prompt, start = "\x01", end = "\x02"):
-        '''
-        Fixes issue on Linux with the readline module
-        See: https://github.com/python/cpython/issues/61539
-        '''
-        escaped = False
-        result = ""
-
-        for c in prompt:
-                if c == "\x1b" and not escaped:
-                        result += start + c
-                        escaped = True
-                elif c.isalpha() and escaped:
-                        result += c + end
-                        escaped = False
-                else:
-                        result += c
-
-        return result
-
 def version():
     print(f'version {VERSION}')
 
@@ -241,30 +275,22 @@ def list_models():
 
 def interactive(args : Arguments):
 
-    enable_emacs_editing()
-
-    COLOR_CODE = get_color_codes(no_color = not args.display_args.color)
-
     completion_args = args.completion_args
     display_args = args.display_args
 
-    hist = [ get_system_message( args.system_message ) ]
+    system_message = get_system_message(args.system_message)
+    interactive_editor = args.interactive_editor
 
-    PROMPT = f'[{COLOR_CODE.WHITE}#{COLOR_CODE.RESET}] '
-    PROMPT = surround_ansi_escapes(PROMPT)
+    hist = [ system_message ]
 
-    def prompt_message() -> bool:
-
-        # Control-D closes the input stream
-        try:
-            message = input( PROMPT )
-        except (EOFError, KeyboardInterrupt):
-            print()
-            return False
+    no_color = not display_args.color
 
-        hist.append( ChatMessage( Role.USER, message ) )
+    prompter = Prompter(no_color = no_color)
 
-        return True
+    CLEAR = prompter.add_command('clear', 'clears the context (excluding the system message)')
+    EDIT = prompter.add_command('edit', 'launches a terminal editor')
+    EXIT = prompter.add_command('exit', 'exit the terminal')
+    HELP = prompter.add_command('help', 'print all interactive commands')
 
     print(f'GPT Chat CLI version {VERSION}')
     print(f'Press Control-D to exit')
@@ -273,26 +299,50 @@ def interactive(args : Arguments):
 
     if initial_message:
         print( PROMPT, initial_message, sep='', flush=True )
-        hist.append( ChatMessage( Role.USER, initial_message ) )
-    else:
-        if not prompt_message():
-            return
 
-    while True:
+    with prompter as prompt:
+        while True:
+            try:
+                if initial_message:
+                    cmd, args = None, initial_message
+                    initial_message = None
+                else:
+                    cmd, args = prompt.input()
+
+                if cmd == CLEAR:
+                    hist = [ system_message ]
+                    continue
+                elif cmd == EDIT:
+                    message = _launch_interactive_editor(
+                        editor=interactive_editor
+                    )
+                    print( prompt.prompt, end='')
+                    print( message, end='' )
+                elif cmd == EXIT:
+                    return
+                elif cmd == HELP:
+                    prompt.print_help()
+                    continue
+                else:
+                    message = args
 
-        completion = create_chat_completion(hist, completion_args)
+                if message == '':
+                    continue
 
-        try:
-            response = print_streamed_response(
-                display_args, completion, 1, return_responses=True,
-            )[0]
+                hist.append( ChatMessage( Role.USER, message ) )
+
+                completion = create_chat_completion(hist, completion_args)
 
-            hist.append( ChatMessage(Role.ASSISTANT, response) )
-        except KeyboardInterrupt:
-            print()
+                response = print_streamed_response(
+                    display_args, completion, 1, return_responses=True,
+                )[0]
 
-        if not prompt_message():
-            return
+                hist.append( ChatMessage(Role.ASSISTANT, response) )
+            except KeyboardInterrupt: # Skip to next prompt
+                continue
+            except EOFError: # Exit on Control-D
+                print()
+                sys.exit(1)
 
 def singleton(args: Arguments):
     completion_args = args.completion_args
diff --git a/src/gpt_chat_cli/color.py b/src/gpt_chat_cli/color.py
index ce1b182..9de82f1 100644
--- a/src/gpt_chat_cli/color.py
+++ b/src/gpt_chat_cli/color.py
@@ -90,3 +90,23 @@ def get_color_codes(no_color=False) -> ColorCode:
         return NoColorColorCode
     else:
         return VT100ColorCode
+
+def surround_ansi_escapes(prompt, start = "\x01", end = "\x02"):
+        '''
+        Fixes issue on Linux with the readline module
+        See: https://github.com/python/cpython/issues/61539
+        '''
+        escaped = False
+        result = ""
+
+        for c in prompt:
+                if c == "\x1b" and not escaped:
+                        result += start + c
+                        escaped = True
+                elif c.isalpha() and escaped:
+                        result += c + end
+                        escaped = False
+                else:
+                        result += c
+
+        return result