diff options
author | flu0r1ne <flu0r1ne@flu0r1ne.net> | 2023-05-09 15:57:05 -0500 |
---|---|---|
committer | flu0r1ne <flu0r1ne@flu0r1ne.net> | 2023-05-09 15:57:05 -0500 |
commit | 4f684f7073149c48ea9dd962b64a8bf7a169a1cf (patch) | |
tree | 3d642519b77bf9a6e451af62e7a50d7ee99f5a1d /src/gpt_chat_cli/cmd.py | |
parent | 61a9d6426cdb0e32436b704561b7a3ff47b60b04 (diff) | |
download | gpt-chat-cli-4f684f7073149c48ea9dd962b64a8bf7a169a1cf.tar.xz gpt-chat-cli-4f684f7073149c48ea9dd962b64a8bf7a169a1cf.zip |
Rewrote the argparsing functionality to enable autocompletion via the "kislyuk/argcomplete" package.
This essentially means:
- `argparsing.py` does a minimal amount of work to initialize the arg parser.
Then it attempts dynamic completion.
- `argvalidation.py` processes the raw arguments parsed in
`argparsing.py`, validates them, issues warnings (if required), and
splits them into logical groupings
- Commands in `gcli.py` have been moved to `cmd.py`
- `main.py` provides an initial control path to call these functions in
succession
Diffstat (limited to 'src/gpt_chat_cli/cmd.py')
-rw-r--r-- | src/gpt_chat_cli/cmd.py | 326 |
1 files changed, 326 insertions, 0 deletions
diff --git a/src/gpt_chat_cli/cmd.py b/src/gpt_chat_cli/cmd.py new file mode 100644 index 0000000..d7aed6c --- /dev/null +++ b/src/gpt_chat_cli/cmd.py @@ -0,0 +1,326 @@ +#!/bin/env python3 + +import sys +import openai +import pickle +import os +import datetime + +from collections import defaultdict +from dataclasses import dataclass +from typing import Tuple, Optional + +from .openai_wrappers import ( + create_chat_completion, + list_models, + OpenAIChatResponse, + OpenAIChatResponseStream, + FinishReason, + Role, + ChatMessage +) + +from .argvalidation import ( + Arguments, + DisplayArguments, + CompletionArguments, + DebugArguments, + MessageSource +) + +from .version import VERSION +from .color import get_color_codes +from .chat_colorizer import ChatColorizer + +########################### +#### UTILS #### +########################### + +def resolve_initial_message(src: MessageSource, interactive=False) -> str: + msg = None + + if src.message: + msg = src.message + elif src.prompt_from_fd: + with os.fdopen(src.prompt_from_fd, "r") as f: + msg = f.read() + elif src.prompt_from_file: + with open(src.prompt_from_file, "r") as f: + msg = f.read() + elif not interactive: + msg = sys.stdin.read() + + return msg + +def get_system_message(system_message : Optional[str]): + + if not system_message: + + current_date_time = datetime.datetime.now() + + system_message = f'The current date is {current_date_time}. When emitting code or producing markdown, ensure to label fenced code blocks with the language in use.' + + return ChatMessage(Role.SYSTEM, system_message) + +def enable_emacs_editing(): + try: + import readline + except ImportError: + pass + +########################### +#### SAVE / REPLAY #### +########################### + +@dataclass +class CompletionContext: + message: str + completion_args: CompletionArguments + system_message: Optional[str] = None + +def create_singleton_chat_completion(ctx : CompletionContext): + + hist = [ + get_system_message(ctx.system_message), + ChatMessage(Role.USER, ctx.message) + ] + + completion = create_chat_completion(hist, ctx.completion_args) + + return completion + +def save_response_and_arguments(args : Arguments) -> None: + + message = resolve_initial_message(args.initial_message) + + ctx = CompletionContext( + message=message, + completion_args=args.completion_args, + system_message=args.system_message + ) + + completion = create_singleton_chat_completion( + message, + args.completion_args, + args.system_message, + ) + + completion = list(completion) + + filename = args.debug_args.save_response_to_file + + with open(filename, 'wb') as f: + pickle.dump((ctx, completion,), f) + +def load_response_and_arguments(args : Arguments) \ + -> Tuple[CompletionContext, OpenAIChatResponseStream]: + + filename = args.debug_args.load_response_from_file + + with open(filename, 'rb') as f: + ctx, completion = pickle.load(f) + + return (ctx, completion) + +######################### +#### PRETTY PRINTING #### +######################### + +@dataclass +class CumulativeResponse: + delta_content: str = "" + finish_reason: FinishReason = FinishReason.NONE + content: str = "" + + def take_delta(self : "CumulativeResponse"): + chunk = self.delta_content + self.delta_content = "" + return chunk + + def add_content(self : "CumulativeResponse", new_chunk : str): + self.content += new_chunk + self.delta_content += new_chunk + +def print_streamed_response( + display_args : DisplayArguments, + completion : OpenAIChatResponseStream, + n_completions : int, + return_responses : bool = False + ) -> None: + """ + Print the response in real time by printing the deltas as they occur. If multiple responses + are requested, print the first in real-time, accumulating the others in the background. One the + first response completes, move on to the second response printing the deltas in real time. Continue + on until all responses have been printed. + """ + + no_color = not display_args.color + + COLOR_CODE = get_color_codes(no_color = no_color) + adornments = display_args.adornments + + cumu_responses = defaultdict(CumulativeResponse) + display_idx = 0 + prompt_printed = False + + chat_colorizer = ChatColorizer(no_color = no_color) + + for update in completion: + + for choice in update.choices: + delta = choice.delta + + if delta.content: + cumu_responses[choice.index].add_content(delta.content) + + if choice.finish_reason is not FinishReason.NONE: + cumu_responses[choice.index].finish_reason = choice.finish_reason + + display_response = cumu_responses[display_idx] + + if not prompt_printed and adornments: + res_indicator = '' if n_completions == 1 else \ + f' {display_idx + 1}/{n_completions}' + PROMPT = f'[{COLOR_CODE.GREEN}{update.model}{COLOR_CODE.RESET}{COLOR_CODE.RED}{res_indicator}{COLOR_CODE.RESET}]' + prompt_printed = True + print(PROMPT, end=' ', flush=True) + + content = display_response.take_delta() + chat_colorizer.add_chunk( content ) + + chat_colorizer.print() + + if display_response.finish_reason is not FinishReason.NONE: + chat_colorizer.finish() + chat_colorizer.print() + chat_colorizer = ChatColorizer( no_color=no_color ) + + if display_idx < n_completions: + display_idx += 1 + prompt_printed = False + + if adornments: + print(end='\n\n', flush=True) + else: + print(end='\n', flush=True) + + if return_responses: + return [ cumu_responses[i].content for i in range(n_completions) ] + +######################### +#### 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}') + +def list_models(): + for model in list_models(): + print(model) + +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 ) ] + + PROMPT = f'[{COLOR_CODE.WHITE}#{COLOR_CODE.RESET}] ' + PROMPT = surround_ansi_escapes(PROMPT) + + def prompt_message() -> bool: + + # Control-D closes the input stream + try: + message = input( PROMPT ) + except (EOFError, KeyboardInterrupt): + print() + return False + + hist.append( ChatMessage( Role.USER, message ) ) + + return True + + print(f'GPT Chat CLI version {VERSION}') + print(f'Press Control-D to exit') + + initial_message = resolve_initial_message(args.initial_message, interactive=True) + + 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: + + completion = create_chat_completion(hist, completion_args) + + try: + response = print_streamed_response( + display_args, completion, 1, return_responses=True, + )[0] + + hist.append( ChatMessage(Role.ASSISTANT, response) ) + except KeyboardInterrupt: + print() + + if not prompt_message(): + return + +def singleton(args: Arguments): + completion_args = args.completion_args + + debug_args : DebugArguments = args.debug_args + message = args.initial_message + + if debug_args.save_response_to_file: + save_response_and_arguments(args) + return + elif debug_args.load_response_from_file: + ctx, completion = load_response_and_arguments(args) + + message = ctx.message + completion_args = ctx.completion_args + else: + # message is only None is a TTY is not attached + message = resolve_initial_message(args.initial_message) + + ctx = CompletionContext( + message=message, + completion_args=completion_args, + system_message=args.system_message + ) + + completion = create_singleton_chat_completion(ctx) + + print_streamed_response( + args.display_args, + completion, + completion_args.n_completions + ) |