#!/usr/bin/env python # -*- coding: utf-8 -*- # vim: set fileencoding=utf-8 : import sys import os import argcomplete import argparse from pathlib import Path from typing import Any from .configuration import Config, default_config_file from .message import Message from .commands.question import question_cmd from .commands.tags import tags_cmd from .commands.config import config_cmd from .commands.hist import hist_cmd from .commands.print import print_cmd from .commands.translation import translation_cmd from .commands.glossary import glossary_cmd from .chat import msg_location 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_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_file) # subcommand-parser cmdparser = parser.add_subparsers(dest='command', title='commands', description='supported commands', required=True) # a parent parser for all commands that support tag selection tag_parser = argparse.ArgumentParser(add_help=False) tag_arg = tag_parser.add_argument('-t', '--or-tags', nargs='*', help='List of tags (one must match)', metavar='OTAGS') tag_arg.completer = tags_completer # type: ignore atag_arg = tag_parser.add_argument('-k', '--and-tags', nargs='*', help='List of tags (all must match)', metavar='ATAGS') atag_arg.completer = tags_completer # type: ignore etag_arg = tag_parser.add_argument('-x', '--exclude-tags', nargs='*', help='List of tags to exclude', metavar='XTAGS') etag_arg.completer = tags_completer # type: ignore otag_arg = tag_parser.add_argument('-o', '--output-tags', nargs='+', help='List of output tags (default: use input tags)', metavar='OUTAGS') otag_arg.completer = tags_completer # type: ignore # a parent parser for all commands that support AI configuration ai_parser = argparse.ArgumentParser(add_help=False) ai_parser.add_argument('-A', '--AI', help='AI ID to use', metavar='AI_ID') ai_parser.add_argument('-M', '--model', help='Model to use', metavar='MODEL') ai_parser.add_argument('-n', '--num-answers', help='Number of answers to request', type=int, default=1) ai_parser.add_argument('-m', '--max-tokens', help='Max. nr. of tokens', type=int) ai_parser.add_argument('-T', '--temperature', help='Temperature value', type=float) # 'question' command parser question_cmd_parser = cmdparser.add_parser('question', parents=[tag_parser, ai_parser], help="ask, create and process questions.", aliases=['q']) question_cmd_parser.set_defaults(func=question_cmd) question_group = question_cmd_parser.add_mutually_exclusive_group(required=True) question_group.add_argument('-a', '--ask', nargs='+', help='Ask a question', metavar='QUESTION') question_group.add_argument('-c', '--create', nargs='+', help='Create a question', metavar='QUESTION') question_group.add_argument('-r', '--repeat', nargs='*', help='Repeat a question', metavar='MESSAGE') question_group.add_argument('-p', '--process', nargs='*', help='Process existing questions', metavar='MESSAGE') question_cmd_parser.add_argument('-l', '--location', choices=[x.value for x in msg_location if x not in [msg_location.MEM, msg_location.DISK]], default='db', help='Use given location when building the chat history (default: \'db\')') question_cmd_parser.add_argument('-g', '--glob', help='Filter message files using the given glob pattern') question_cmd_parser.add_argument('-O', '--overwrite', help='Overwrite existing messages when repeating them', action='store_true') question_cmd_parser.add_argument('-s', '--source-text', nargs='+', help='Add content of a file to the query', metavar='FILE') question_cmd_parser.add_argument('-S', '--source-code', nargs='+', help='Add source code file content to the chat history', metavar='FILE') # 'hist' command parser hist_cmd_parser = cmdparser.add_parser('hist', parents=[tag_parser], help="Print and manage chat history.", aliases=['h']) hist_cmd_parser.set_defaults(func=hist_cmd) hist_group = hist_cmd_parser.add_mutually_exclusive_group(required=True) hist_group.add_argument('-p', '--print', help='Print the DB chat history', action='store_true') hist_group.add_argument('-c', '--convert', help='Convert all message files to the given format [txt|yaml]', metavar='FORMAT') hist_cmd_parser.add_argument('-w', '--with-metadata', help="Print chat history with metadata (tags, filename, AI, etc.).", action='store_true') hist_cmd_parser.add_argument('-S', '--source-code-only', help='Only print embedded source code', action='store_true') hist_cmd_parser.add_argument('-A', '--answer', help='Print only answers with given substring', metavar='SUBSTRING') hist_cmd_parser.add_argument('-Q', '--question', help='Print only questions with given substring', metavar='SUBSTRING') hist_cmd_parser.add_argument('-d', '--tight', help='Print without message separators', action='store_true') hist_cmd_parser.add_argument('-P', '--no-paging', help='Print without paging', action='store_true') hist_cmd_parser.add_argument('-l', '--location', choices=[x.value for x in msg_location if x not in [msg_location.MEM, msg_location.DISK]], default='db', help='Use given location when building the chat history (default: \'db\')') hist_cmd_parser.add_argument('-g', '--glob', help='Filter message files using the given glob pattern') # 'tags' command parser tags_cmd_parser = cmdparser.add_parser('tags', help="Manage tags.", aliases=['T']) tags_cmd_parser.set_defaults(func=tags_cmd) tags_group = tags_cmd_parser.add_mutually_exclusive_group(required=True) tags_group.add_argument('-l', '--list', help="List all tags and their frequency", action='store_true') tags_cmd_parser.add_argument('-p', '--prefix', help="Filter tags by prefix", metavar='PREFIX') tags_cmd_parser.add_argument('-c', '--contain', help="Filter tags by contained substring", metavar='SUBSTRING') # 'config' command parser config_cmd_parser = cmdparser.add_parser('config', help="Manage configuration", aliases=['c']) config_cmd_parser.set_defaults(func=config_cmd) config_cmd_parser.add_argument('-A', '--AI', help='AI ID to use') 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', '--print-model', help="Print the currently configured model", action='store_true') config_group.add_argument('-c', '--create', help="Create config with default settings in the given file", metavar='FILE') # 'print' command parser print_cmd_parser = cmdparser.add_parser('print', help="Print message files.", aliases=['p']) print_cmd_parser.set_defaults(func=print_cmd) print_group = print_cmd_parser.add_mutually_exclusive_group(required=True) print_group.add_argument('-f', '--file', help='Print given message file', metavar='FILE') print_group.add_argument('-l', '--latest', help='Print latest message', action='store_true') print_cmd_modes = print_cmd_parser.add_mutually_exclusive_group() print_cmd_modes.add_argument('-q', '--question', help='Only print the question', action='store_true') print_cmd_modes.add_argument('-a', '--answer', help='Only print the answer', action='store_true') print_cmd_modes.add_argument('-S', '--only-source-code', help='Only print embedded source code', action='store_true') # 'translation' command parser translation_cmd_parser = cmdparser.add_parser('translation', parents=[ai_parser, tag_parser], help="Ask, create and repeat translations.", aliases=['t']) translation_cmd_parser.set_defaults(func=translation_cmd) translation_group = translation_cmd_parser.add_mutually_exclusive_group(required=True) translation_group.add_argument('-a', '--ask', nargs='+', help='Ask to translate the given text', metavar='TEXT') translation_group.add_argument('-c', '--create', nargs='+', help='Create a translation', metavar='TEXT') translation_group.add_argument('-r', '--repeat', nargs='*', help='Repeat a translation', metavar='MESSAGE') translation_cmd_parser.add_argument('-S', '--source-lang', help="Source language", metavar="LANGUAGE", required=True) translation_cmd_parser.add_argument('-T', '--target-lang', help="Target language", metavar="LANGUAGE", required=True) translation_cmd_parser.add_argument('-G', '--glossaries', nargs='+', help="List of glossary names", metavar="GLOSSARY") translation_cmd_parser.add_argument('-d', '--input-document', help="Document to translate", metavar="FILE") translation_cmd_parser.add_argument('-D', '--output-document', help="Path for the translated document", metavar="FILE") # 'glossary' command parser glossary_cmd_parser = cmdparser.add_parser('glossary', parents=[ai_parser], help="Manage glossaries.", aliases=['g']) glossary_cmd_parser.set_defaults(func=glossary_cmd) glossary_group = glossary_cmd_parser.add_mutually_exclusive_group(required=True) glossary_group.add_argument('-c', '--create', help='Create a glossary', action='store_true') glossary_cmd_parser.add_argument('-n', '--name', help="Glossary name (not ID)", metavar="NAME") glossary_cmd_parser.add_argument('-S', '--source-lang', help="Source language", metavar="LANGUAGE") glossary_cmd_parser.add_argument('-T', '--target-lang', help="Target language", metavar="LANGUAGE") glossary_cmd_parser.add_argument('-f', '--file', help='File path of the goven glossary', metavar='GLOSSARY_FILE') glossary_cmd_parser.add_argument('-D', '--description', help="Glossary description", metavar="DESCRIPTION") glossary_group.add_argument('-l', '--list', help='List existing glossaries', action='store_true') glossary_cmd_parser.add_argument('-E', '--entries', help="Print entries when listing glossaries", action='store_true') argcomplete.autocomplete(parser) return parser def create_directories(config: Config) -> None: # noqa: 11 """ Create the directories in the given configuration if they don't exist. """ def make_dir(path: Path) -> None: try: os.makedirs(path.absolute()) except Exception as e: print(f"Creating directory '{path.absolute()}' failed with: {e}") sys.exit(1) # Cache cache_path = Path(config.cache) if not cache_path.exists(): answer = input(f"Cache directory '{cache_path}' does not exist. Create it? [y/n]") if answer.lower() in ['y', 'yes']: make_dir(cache_path.absolute()) else: print("Can't continue without a valid cache directory!") sys.exit(1) # DB db_path = Path(config.db) if not db_path.exists(): answer = input(f"DB directory '{db_path}' does not exist. Create it? [y/n]") if answer.lower() in ['y', 'yes']: make_dir(db_path.absolute()) else: print("Can't continue without a valid DB directory!") sys.exit(1) # Glossaries if config.glossaries: glossaries_path = Path(config.glossaries) if not glossaries_path.exists(): answer = input(f"Glossaries directory '{glossaries_path}' does not exist. Create it? [y/n]") if answer.lower() in ['y', 'yes']: make_dir(glossaries_path.absolute()) else: print("Can't continue without a valid glossaries directory. Create it or remove it from the configuration.") sys.exit(1) def main() -> int: parser = create_parser() args = parser.parse_args() command = parser.parse_args() if command.func == config_cmd: command.func(command) else: config = Config.from_file(args.config) create_directories(config) command.func(command, config) return 0 if __name__ == '__main__': sys.exit(main())