diff options
Diffstat (limited to 'src/gpt_chat_cli/cmd.py')
| -rw-r--r-- | src/gpt_chat_cli/cmd.py | 166 |
1 files changed, 108 insertions, 58 deletions
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 |