From f6cd64a7870dafa80cc81e52ad9a8f7fe6529154 Mon Sep 17 00:00:00 2001
From: flu0r1ne <flu0r1ne@flu0r1ne.net>
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