diff --git a/chatmastermind/main.py b/chatmastermind/main.py index 2e0cce1..cc634fc 100755 --- a/chatmastermind/main.py +++ b/chatmastermind/main.py @@ -8,7 +8,7 @@ import argcomplete import argparse import pathlib from .utils import terminal_width, print_tag_args, print_chat_hist, display_source_code, print_tags_frequency, ConfigType -from .storage import save_answers, create_chat_hist, get_tags, get_tags_unique, read_file, read_config, dump_data +from .storage import save_answers, create_chat_hist, get_tags, get_tags_unique, read_file, read_config, write_config, dump_data from .api_client import ai, openai_api_key, print_models from itertools import zip_longest @@ -63,12 +63,20 @@ def tag_cmd(args: argparse.Namespace, config: ConfigType) -> None: print_tags_frequency(get_tags(config, None)) -def model_cmd(args: argparse.Namespace, config: ConfigType) -> None: +def config_cmd(args: argparse.Namespace, config: ConfigType) -> None: """ - Handler for the 'model' command. + Handler for the 'config' command. """ - if args.list: + if type(config['openai']) is not dict: + raise RuntimeError('Configuration openai is not a dict.') + + if args.list_models: print_models() + elif args.show_model: + print(config['openai']['model']) + elif args.model: + config['openai']['model'] = args.model + write_config(args.config, config) def ask_cmd(args: argparse.Namespace, config: ConfigType) -> None: @@ -188,12 +196,16 @@ def create_parser() -> argparse.ArgumentParser: tag_cmd_parser.add_argument('-l', '--list', help="List all tags and their frequency", action='store_true') - # 'model' command parser - model_cmd_parser = cmdparser.add_parser('model', - help="Manage models.") - model_cmd_parser.set_defaults(func=model_cmd) - model_cmd_parser.add_argument('-l', '--list', help="List all available models", - action='store_true') + # 'config' command parser + config_cmd_parser = cmdparser.add_parser('config', + help="Manage configuration") + config_cmd_parser.set_defaults(func=config_cmd) + config_group = config_cmd_parser.add_mutually_exclusive_group(required=True) + config_group.add_argument('-L', '--list-models', help="List all available models", + action='store_true') + config_group.add_argument('-m', '--show-model', help="Show current model", + action='store_true') + config_group.add_argument('-M', '--model', help="Set model in the config file") # 'print' command parser print_cmd_parser = cmdparser.add_parser('print',