use click instead of argparser

askchat/__init__.py

@@ -4,4 +4,6 @@
 __email__ = '[email protected]'
 __version__ = '0.3.4'
 
-from .askchat import ask
+import asyncio
+from .ask import ask
+from .askchat import askchat

askchat/ask.py

@@ -0,0 +1,25 @@
+import click
+import asyncio
+from chattool import Chat
+
+async def show_resp(chat):
+    msg = ''
+    async for char in chat.async_stream_responses(textonly=True):
+        print(char, end='', flush=True)
+        msg += char
+        await asyncio.sleep(0.01)
+    return msg
+
+@click.command()
+@click.argument('message', nargs=-1, required=True)
+def ask(message):
+    """Send a message to ChatGPT and display the response."""
+    message = ' '.join(message).strip()
+    if not message:
+        click.echo("Cannot send an empty message")
+        return
+    chat = Chat(message)
+    asyncio.run(show_resp(chat))
+
+if __name__ == '__main__':
+    ask()

askchat/askchat.py

@@ -1,190 +1,164 @@
 """Main module."""
 
-from argparse import ArgumentParser
+import click, asyncio, askchat
+from pathlib import Path
 from pprint import pprint
 from dotenv import load_dotenv, set_key
-import asyncio, os, uuid, shutil
+import asyncio, os, uuid, shutil, json
 from chattool import Chat, debug_log, load_envs
-import askchat
+from pathlib import Path
+from .ask import show_resp
 
+# Version and Config Path
 VERSION = askchat.__version__
-CONFIG_PATH = os.path.expanduser("~/.askchat")
-CONFIG_FILE = os.path.expanduser("~/.askchat/.env")
-LAST_CHAT_FILE = os.path.expanduser("~/.askchat/_last_chat.json")
+CONFIG_PATH = Path.home() / ".askchat"
+CONFIG_FILE = CONFIG_PATH / ".env"
+LAST_CHAT_FILE = CONFIG_PATH / "_last_chat.json"
 
-# print the response in a typewriter way
-async def show_resp(chat, delay=0.01):
-    msg = ''
-    async for char in chat.async_stream_responses(textonly=True):
-        print(char, end='', flush=True)
-        msg += char
-        await asyncio.sleep(delay)
-    return msg
+def setup():
+    """Application setup: Ensure that necessary folders and files exist."""
+    os.makedirs(CONFIG_PATH, exist_ok=True)
+    if not os.path.exists(CONFIG_FILE):
+        with open(CONFIG_FILE, 'w') as cf:
+            cf.write("# Initial configuration\n")
+    # if not os.path.exists(LAST_CHAT_FILE):
+    #     with open(LAST_CHAT_FILE, 'w') as lcf:
+    #         lcf.write('{"index": 0, "chat_log": []}')
 
-def ask():
-    """Interact with ChatGPT in terminal via chattool"""
-    # parse arguments
-    parser = ArgumentParser()
-    parser.add_argument('message', help='User message', default='', nargs='*')
-    args = parser.parse_args()
-    msg = args.message
-    if isinstance(msg, list):
-        msg = ' '.join(msg)
-    assert len(msg.strip()), 'Please specify message'
-    # call
-    chat = Chat(msg)
-    asyncio.run(show_resp(chat))
+def print_callback(ctx, param, value):
+    print("无参数，回调")
+    if value is None:
+        return "_last_chat"  # default value
+    return value
 
-def main():
-    """Interact with ChatGPT in terminal via chattool"""
-    # parse arguments
-    parser = ArgumentParser()
-    ## arguments for chat message
-    parser.add_argument('message', help='User message', default='', nargs='*')
-    parser.add_argument('-m', '--model', default=None, help='Model name')
-    parser.add_argument('-b', '--base-url', default=None, help='base url of the api(without suffix `/v1`)')
-    parser.add_argument('-a', "--api-key", default=None, help="OpenAI API key")
-    ## Chat with history
-    parser.add_argument('-c', action='store_true', help='Continue the last conversation')
-    parser.add_argument('-r', action='store_true', help='Regenerate the last conversation')
-    parser.add_argument('-s', "--save", default=None, help="Save the conversation to a file")
-    parser.add_argument("-l", "--load", default=None, help="Load the conversation from a file")
-    parser.add_argument("-p", "--print", default=None, nargs='*', help="Print the conversation from " +\
-                        "a file or the last conversation if no file is specified")
-    parser.add_argument("-d", "--delete", default=None, help="Delete the conversation from a file")
-    parser.add_argument("--list", action="store_true", help="List all the conversation files")
-    ## other options
-    parser.add_argument('--debug', action='store_true', help='Print debug log')
-    parser.add_argument('--valid-models', action='store_true', help='Print valid models that contain "gpt" in their names')
-    parser.add_argument('--all-valid-models', action='store_true', help='Print all valid models')
-    parser.add_argument('--generate-config', action="store_true", help="Generate a configuration file by environment table")
-    parser.add_argument('-v', '--version', action='version', version=VERSION)
-    args = parser.parse_args()
-    os.makedirs(CONFIG_PATH, exist_ok=True)
+def generate_config():
+    pass
 
-    # generate config file
-    if args.generate_config:
-        api_key, model = os.getenv("OPENAI_API_KEY"), os.getenv("OPENAI_API_MODEL")
-        base_url, api_base = os.getenv("OPENAI_API_BASE_URL"), os.getenv("OPENAI_API_BASE")
-        # move the old config file to a temporary file
-        if os.path.exists(CONFIG_FILE):
-            # move the old config file to a temporary file
-            os.makedirs("/tmp", exist_ok=True)
-            tmp_file = os.path.join("/tmp", str(uuid.uuid4())[:8] + ".askchat.env")
-            shutil.move(CONFIG_FILE, tmp_file)
-            print(f"Moved old config file to {tmp_file}")
-        # save the config file
-        with open(CONFIG_FILE, "w") as f:
-            # description for the config file
-            f.write("#!/bin/bash\n" +\
-                    "# Description: Env file for askchat.\n" +\
-                    "# Current version: " + VERSION + "\n\n" +\
-                    "# The base url of the API (without suffix /v1)\n" +\
-                    "OPENAI_API_BASE_URL=\n\n" +\
-                    "# The base url of the API (with suffix /v1)\n" +\
-                    "OPENAI_API_BASE=\n\n" +\
-                    "# Your API key\n" +\
-                    "OPENAI_API_KEY=\n\n" +\
-                    "# The model name\n" +\
-                    "# You can use `askchat --all-valid-models` to see the valid models\n" +\
-                    "OPENAI_API_MODEL=\n\n")
-        # write the environment table
-        if api_key: set_key(CONFIG_FILE, "OPENAI_API_KEY", api_key)
-        if base_url: set_key(CONFIG_FILE, "OPENAI_API_BASE_URL", base_url)
-        if api_base: set_key(CONFIG_FILE, "OPENAI_API_BASE", api_base)
-        if model: set_key(CONFIG_FILE, "OPENAI_API_MODEL", model)
-        print("Created config file at", CONFIG_FILE)
-        return
+def load_config():
+    pass
+
+@click.group()
+def cli():
+    """A CLI for interacting with ChatGPT with advanced options."""
+    pass
 
-    # set values
-    ## 1. read from config file
-    if os.path.exists(CONFIG_FILE):
-        load_dotenv(CONFIG_FILE, override=True)
+@cli.command()
+@click.argument('message', nargs=-1)
+@click.option('-m', '--model', default=None, help='Model name')
+@click.option('-b', '--base-url', default=None, help='Base URL of the API (without suffix `/v1`)')
+@click.option('--api-base', default=None, help='Base URL of the API (with suffix `/v1`)')
+@click.option('-a', '--api-key', default=None, help='OpenAI API key')
+# Chat with history
+@click.option('-c', is_flag=True, help='Continue the last conversation')
+@click.option('-r', '--regenerate', is_flag=True, help='Regenerate the last conversation')
+@click.option('-s', '--save', default=None, help='Save the conversation to a file')
+@click.option('-l', '--load', default=None, help='Load the conversation from a file')
+@click.option('-p', '--print', is_flag=True, help='Print the last conversation or a specific conversation')
+@click.option('-d', '--delete', default=None, help='Delete the conversation from a file')
+@click.option('--list', is_flag=True, help='List all the conversation files')
+# Other options
+@click.option('--generate-config', is_flag=True, help='Generate a configuration file by environment table')
+@click.option('--debug', is_flag=True, help='Print debug log')
+@click.option('--valid-models', is_flag=True, help='Print valid models that contain "gpt" in their names')
+@click.option('--all-valid-models', is_flag=True, help='Print all valid models')
+@click.option('-v', '--version', is_flag=True, help='Print the version')
+def askchat( message, model, base_url, api_base, api_key
+           , c, regenerate, save, load, print, delete, list
+           , generate_config, debug, valid_models, all_valid_models, version):
+    """Interact with ChatGPT in terminal via chattool"""
+    setup()
+    message_text = ' '.join(message).strip()
+    # generate config file
+    if generate_config:
+        return generate_config()
+    # set values for the environment variables
+    ## 1. read from config file `~/.askchat/.env`
+    # load_config()
     ## 2. read from command line
-    if args.api_key:
-        os.environ['OPENAI_API_KEY'] = args.api_key
-    if args.base_url:
-        os.environ['OPENAI_API_BASE_URL'] = args.base_url
-    if args.model:
-        os.environ['OPENAI_API_MODEL'] = args.model
-    ## 3. read from environment variables
+    if api_key:
+        os.environ['OPENAI_API_KEY'] = api_key
+    if base_url:
+        os.environ['OPENAI_API_BASE_URL'] = base_url
+    if api_base:
+        os.environ['OPENAI_API_BASE'] = api_base
+    if model:
+        os.environ['OPENAI_API_MODEL'] = model
+    # update environment variables of chattool
     load_envs()
-
     # show debug log
-    if args.debug:
-        debug_log()
-        return
-
+    if debug:
+        return debug_log()
     # show valid models
-    if args.valid_models:
-        print('Valid models that contain "gpt" in their names:')
-        pprint(Chat().get_valid_models())
+    if valid_models:
+        click.echo('Valid models that contain "gpt" in their names:')
+        click.echo(pprint(Chat().get_valid_models()))
         return
-    if args.all_valid_models:
-        print('All valid models:')
-        pprint(Chat().get_valid_models(gpt_only=False))
+    if all_valid_models:
+        click.echo('All valid models:')
+        click.echo(pprint(Chat().get_valid_models(gpt_only=False)))
         return
-
-    # deal with chat history
-    call_history = False
-    ## load chat
-    if args.load is not None:
-        new_file = os.path.join(CONFIG_PATH, args.load) + ".json"
-        shutil.copyfile(new_file, LAST_CHAT_FILE)
-        print("Loaded conversation from", new_file)
-        call_history = True
-    ## save chat
-    if args.save is not None:
-        new_file = os.path.join(CONFIG_PATH, args.save) + ".json"
-        shutil.copyfile(LAST_CHAT_FILE, new_file)
-        print("Saved conversation to", new_file)
-        call_history = True
-    ## delete chat
-    if args.delete is not None:
-        new_file = os.path.join(CONFIG_PATH, args.delete) + ".json"
-        if os.path.exists(new_file):
-            os.remove(new_file)
-            print("Deleted conversation at", new_file)
-        else:
-            print("No such file", new_file)
-        call_history = True
-    ## list chat
-    if args.list:
-        print("All conversation files:")
-        for file in os.listdir(CONFIG_PATH):
-            if not file.startswith("_") and file.endswith(".json"):
-                print(" -", file[:-5])
-        call_history = True
-    ## print chat
-    if args.print is not None:
-        names = args.print
-        assert len(names) <= 1, "Only one file can be specified"
-        new_file = os.path.join(CONFIG_PATH, names[0]) + ".json" if len(names) else LAST_CHAT_FILE
-        assert os.path.exists(new_file), "No conversation file found, check the chat list with `--list` option"
-        chat = Chat.load(new_file)
-        chat.print_log()
-        call_history = True
-    if call_history: return
-    # Initial message
-    msg = args.message
-    if isinstance(msg, list):
-        msg = ' '.join(msg).strip()
-    chat = Chat(msg)
-    if os.path.exists(LAST_CHAT_FILE):
-        if args.c:
-            chat = Chat.load(LAST_CHAT_FILE)
-            chat.user(msg)
-        elif args.r:
-            # pop out the last two messages
+    # Handle chat history operations
+    if load:
+        try:
+            shutil.copyfile(CONFIG_PATH / f"{load}.json", LAST_CHAT_FILE)
+            click.echo(f"Loaded conversation from {CONFIG_PATH}/{load}.json")
+        except FileNotFoundError:
+            click.echo(f"The specified conversation {load} does not exist." +\
+                       "Please check the chat list with `--list` option.")
+        return
+    if save:
+        try:
+            shutil.copyfile(LAST_CHAT_FILE, CONFIG_PATH / f"{save}.json")
+            click.echo(f"Saved conversation to {CONFIG_PATH}/{save}.json")
+        except FileNotFoundError:
+            click.echo("No last conversation to save.")
+        return
+    if delete:
+        try:
+            os.remove(CONFIG_PATH / f"{delete}.json")
+            click.echo(f"Deleted conversation at {CONFIG_PATH}/{delete}.json")
+        except FileNotFoundError:
+            click.echo(f"The specified conversation {CONFIG_PATH}/{delete}.json does not exist.")
+        return
+    if list:
+        click.echo("All conversation files:")
+        for file in CONFIG_PATH.glob("*.json"):
+            if not file.name.startswith("_"):
+                click.echo(f" - {file.stem}")
+        return
+    if print:
+        fname = message_text if message_text else '_last_chat'
+        fname = f"{CONFIG_PATH}/{fname}.json"
+        try:
+            Chat().load(fname).print_log()
+        except FileNotFoundError:
+            click.echo(f"The specified conversation {fname} does not exist.")
+        return
+    # Handle version option
+    if version:
+        click.echo(f"askchat version: {VERSION}")
+        return
+    # Main chat
+    chat = Chat()
+    if c or regenerate: # Load last chat if -c or -r is used
+        try:
             chat = Chat.load(LAST_CHAT_FILE)
-            assert len(chat) > 1, "You should have at least two messages in the conversation"
-            chat.pop()
-            if len(msg) != 0: # not empty message
-                chat.pop()
-                chat.user(msg)
-            # if msg is empty, regenerate the last message
-    assert len(chat) > 0 and len(chat.last_message) > 0, "Please specify message!"
-    # call the function
-    newmsg = asyncio.run(show_resp(chat))
-    chat.assistant(newmsg)
-    chat.save(LAST_CHAT_FILE, mode='w')
+        except FileNotFoundError:
+            click.echo("No last conversation found. Starting a new conversation.")
+            return
+    if regenerate:
+        if len(chat) < 2:
+            click.echo("You should have at least two messages in the conversation")
+            return
+        chat.pop()
+    else:
+        if not message_text:
+            click.echo("Please specify message!")
+            return
+        chat.user(message_text)
+    # Simulate chat response
+    chat.assistant(asyncio.run(show_resp(chat)))
+    chat.save(LAST_CHAT_FILE, mode='w')
+
+if __name__ == '__main__':
+    cli()

setup.cfg

@@ -9,4 +9,5 @@ addopts = --ignore=setup.py
 [options.entry_points]
 console_scripts =
     ask = askchat:ask
-    askchat = askchat.askchat:main
+    askchat = askchat:askchat
+

setup.py

@@ -9,7 +9,7 @@
 with open('README.md') as readme_file:
     readme = readme_file.read()
 
-requirements = ['chattool>=3.1.1', "python-dotenv>=0.17.0"]
+requirements = ['chattool>=3.1.1', "python-dotenv>=0.17.0", 'Click>=7.0']
 
 test_requirements = ['pytest>=3']