From f6cd64a7870dafa80cc81e52ad9a8f7fe6529154 Mon Sep 17 00:00:00 2001 From: flu0r1ne Date: Fri, 5 May 2023 02:35:50 -0500 Subject: Add command to list models --- src/gpt_chat_cli/argparsing.py | 12 ++++++++++-- src/gpt_chat_cli/gcli.py | 8 ++++++++ src/gpt_chat_cli/openai_wrappers.py | 13 +++++++++++++ 3 files changed, 31 insertions(+), 2 deletions(-) (limited to 'src/gpt_chat_cli') diff --git a/src/gpt_chat_cli/argparsing.py b/src/gpt_chat_cli/argparsing.py index 81a9b47..4597e84 100644 --- a/src/gpt_chat_cli/argparsing.py +++ b/src/gpt_chat_cli/argparsing.py @@ -65,6 +65,7 @@ class Arguments: completion_args: CompletionArguments display_args: DisplayArguments version: bool + list_models: bool debug_args: Optional[DebugArguments] = None def split_arguments(args: argparse.Namespace) -> Arguments: @@ -93,7 +94,8 @@ def split_arguments(args: argparse.Namespace) -> Arguments: completion_args=completion_args, display_args=display_args, debug_args=debug_args, - version=args.version + version=args.version, + list_models=args.list_models, ) def parse_args() -> Arguments: @@ -203,11 +205,17 @@ def parse_args() -> Arguments: parser.add_argument( "--version", - # type=bool, action="store_true", help="Print version and exit" ) + parser.add_argument( + "-l", + "--list-models", + action="store_true", + help="List models and exit" + ) + parser.add_argument( "message", type=str, diff --git a/src/gpt_chat_cli/gcli.py b/src/gpt_chat_cli/gcli.py index f0b754e..a9d36e4 100644 --- a/src/gpt_chat_cli/gcli.py +++ b/src/gpt_chat_cli/gcli.py @@ -10,6 +10,7 @@ from typing import Tuple from .openai_wrappers import ( create_chat_completion, + list_models, OpenAIChatResponse, OpenAIChatResponseStream, FinishReason, @@ -138,6 +139,13 @@ def main(): print(f'version {VERSION}') sys.exit(0) + if args.list_models: + + for model in list_models(): + print(model) + + sys.exit(0) + completion_args = args.completion_args COLOR_CODE = get_color_codes(no_color = not args.display_args.color) diff --git a/src/gpt_chat_cli/openai_wrappers.py b/src/gpt_chat_cli/openai_wrappers.py index 784a9ce..413ec24 100644 --- a/src/gpt_chat_cli/openai_wrappers.py +++ b/src/gpt_chat_cli/openai_wrappers.py @@ -67,3 +67,16 @@ def create_chat_completion(*args, **kwargs) \ OpenAIChatResponse.from_json(update) \ for update in openai.ChatCompletion.create(*args, **kwargs) ) + +def list_models() -> List[str]: + + model_data = openai.Model.list() + + models = [] + + for model in model_data["data"]: + models.append(model["id"]) + + models.sort() + + return models -- cgit v1.2.3