diff --git a/chatmastermind/main.py b/chatmastermind/main.py index b10b97b..7f34b54 100755 --- a/chatmastermind/main.py +++ b/chatmastermind/main.py @@ -6,61 +6,19 @@ import sys import argcomplete import argparse from pathlib import Path -from .utils import terminal_width, print_tag_args, print_chat_hist, ChatType -from .storage import save_answers, create_chat_hist -from .api_client import ai, openai_api_key, print_models -from .configuration import Config +from .configuration import Config, default_config_path from .chat import ChatDB from .message import Message, MessageFilter, MessageError, Question from .ai_factory import create_ai from .ai import AI, AIResponse -from itertools import zip_longest from typing import Any -default_config = '.config.yaml' - def tags_completer(prefix: str, parsed_args: Any, **kwargs: Any) -> list[str]: config = Config.from_file(parsed_args.config) return list(Message.tags_from_dir(Path(config.db), prefix=prefix)) -def create_question_with_hist(args: argparse.Namespace, - config: Config, - ) -> tuple[ChatType, str, list[str]]: - """ - Creates the "AI request", including the question and chat history as determined - by the specified tags. - """ - tags = args.or_tags or [] - xtags = args.exclude_tags or [] - otags = args.output_tags or [] - - if not args.source_code_only: - print_tag_args(tags, xtags, otags) - - question_parts = [] - question_list = args.question if args.question is not None else [] - source_list = args.source if args.source is not None else [] - - for question, source in zip_longest(question_list, source_list, fillvalue=None): - if question is not None and source is not None: - with open(source) as r: - question_parts.append(f"{question}\n\n```\n{r.read().strip()}\n```") - elif question is not None: - question_parts.append(question) - elif source is not None: - with open(source) as r: - question_parts.append(f"```\n{r.read().strip()}\n```") - - full_question = '\n\n'.join(question_parts) - chat = create_chat_hist(full_question, tags, xtags, config, - match_all_tags=True if args.and_tags else False, # FIXME - with_tags=False, - with_file=False) - return chat, full_question, tags - - def tags_cmd(args: argparse.Namespace, config: Config) -> None: """ Handler for the 'tags' command. @@ -74,17 +32,12 @@ def tags_cmd(args: argparse.Namespace, config: Config) -> None: # TODO: add renaming -def config_cmd(args: argparse.Namespace, config: Config) -> None: +def config_cmd(args: argparse.Namespace) -> None: """ Handler for the 'config' command. """ - if args.list_models: - print_models() - elif args.print_model: - print(config.openai.model) - elif args.model: - config.openai.model = args.model - config.to_file(args.config) + if args.create: + Config.create_default(Path(args.create)) def question_cmd(args: argparse.Namespace, config: Config) -> None: @@ -128,25 +81,6 @@ def question_cmd(args: argparse.Namespace, config: Config) -> None: pass -def ask_cmd(args: argparse.Namespace, config: Config) -> None: - """ - Handler for the 'ask' command. - """ - if args.max_tokens: - config.openai.max_tokens = args.max_tokens - if args.temperature: - config.openai.temperature = args.temperature - if args.model: - config.openai.model = args.model - chat, question, tags = create_question_with_hist(args, config) - print_chat_hist(chat, False, args.source_code_only) - otags = args.output_tags or [] - answers, usage = ai(chat, config, args.num_answers) - save_answers(question, answers, tags, otags, config) - print("-" * terminal_width()) - print(f"Usage: {usage}") - - def hist_cmd(args: argparse.Namespace, config: Config) -> None: """ Handler for the 'hist' command. @@ -182,7 +116,7 @@ def print_cmd(args: argparse.Namespace, config: Config) -> None: def create_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="ChatMastermind is a Python application that automates conversation with AI") - parser.add_argument('-C', '--config', help='Config file name.', default=default_config) + parser.add_argument('-C', '--config', help='Config file name.', default=default_config_path) # subcommand-parser cmdparser = parser.add_subparsers(dest='command', @@ -227,22 +161,6 @@ def create_parser() -> argparse.ArgumentParser: question_cmd_parser.add_argument('-S', '--source-code-only', help='Add pure source code to the chat history', action='store_true') - # 'ask' command parser - ask_cmd_parser = cmdparser.add_parser('ask', parents=[tag_parser], - help="Ask a question.", - aliases=['a']) - ask_cmd_parser.set_defaults(func=ask_cmd) - ask_cmd_parser.add_argument('-q', '--question', nargs='+', help='Question to ask', - required=True) - ask_cmd_parser.add_argument('-m', '--max-tokens', help='Max tokens to use', type=int) - ask_cmd_parser.add_argument('-T', '--temperature', help='Temperature to use', type=float) - ask_cmd_parser.add_argument('-M', '--model', help='Model to use') - ask_cmd_parser.add_argument('-n', '--num-answers', help='Number of answers to produce', type=int, - default=1) - ask_cmd_parser.add_argument('-s', '--source', nargs='+', help='Source add content of a file to the query') - ask_cmd_parser.add_argument('-S', '--source-code-only', help='Add pure source code to the chat history', - action='store_true') - # 'hist' command parser hist_cmd_parser = cmdparser.add_parser('hist', parents=[tag_parser], help="Print chat history.", @@ -278,7 +196,7 @@ def create_parser() -> argparse.ArgumentParser: action='store_true') config_group.add_argument('-m', '--print-model', help="Print the currently configured model", action='store_true') - config_group.add_argument('-M', '--model', help="Set model in the config file") + config_group.add_argument('-c', '--create', help="Create config with default settings in the given file") # 'print' command parser print_cmd_parser = cmdparser.add_parser('print', @@ -297,11 +215,12 @@ def main() -> int: parser = create_parser() args = parser.parse_args() command = parser.parse_args() - config = Config.from_file(args.config) - openai_api_key(config.openai.api_key) - - command.func(command, config) + if command.func == config_cmd: + command.func(command) + else: + config = Config.from_file(args.config) + command.func(command, config) return 0