diff options
| -rw-r--r-- | src/gpt_chat_cli/argparsing.py | 11 | ||||
| -rw-r--r-- | src/gpt_chat_cli/argvalidation.py | 4 | ||||
| -rw-r--r-- | src/gpt_chat_cli/cmd.py | 166 | ||||
| -rw-r--r-- | src/gpt_chat_cli/color.py | 20 |
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 |