From 810332d1409c54bb71211cbc8a23879c4a3652e4 Mon Sep 17 00:00:00 2001 From: juk0de Date: Fri, 8 Sep 2023 13:31:01 +0200 Subject: [PATCH] cmm: splitted commands into separate modules (and more cleanup) --- chatmastermind/commands/config.py | 11 +++ chatmastermind/commands/hist.py | 23 +++++ chatmastermind/commands/print.py | 19 ++++ chatmastermind/commands/question.py | 57 ++++++++++++ chatmastermind/commands/tags.py | 17 ++++ chatmastermind/main.py | 131 +++++----------------------- setup.py | 2 +- tests/test_ai_factory.py | 48 ++++++++++ 8 files changed, 196 insertions(+), 112 deletions(-) create mode 100644 chatmastermind/commands/config.py create mode 100644 chatmastermind/commands/hist.py create mode 100644 chatmastermind/commands/print.py create mode 100644 chatmastermind/commands/question.py create mode 100644 chatmastermind/commands/tags.py create mode 100644 tests/test_ai_factory.py diff --git a/chatmastermind/commands/config.py b/chatmastermind/commands/config.py new file mode 100644 index 0000000..262164c --- /dev/null +++ b/chatmastermind/commands/config.py @@ -0,0 +1,11 @@ +import argparse +from pathlib import Path +from ..configuration import Config + + +def config_cmd(args: argparse.Namespace) -> None: + """ + Handler for the 'config' command. + """ + if args.create: + Config.create_default(Path(args.create)) diff --git a/chatmastermind/commands/hist.py b/chatmastermind/commands/hist.py new file mode 100644 index 0000000..88ed3be --- /dev/null +++ b/chatmastermind/commands/hist.py @@ -0,0 +1,23 @@ +import argparse +from pathlib import Path +from ..configuration import Config +from ..chat import ChatDB +from ..message import MessageFilter + + +def hist_cmd(args: argparse.Namespace, config: Config) -> None: + """ + Handler for the 'hist' command. + """ + + mfilter = MessageFilter(tags_or=args.or_tags, + tags_and=args.and_tags, + tags_not=args.exclude_tags, + question_contains=args.question, + answer_contains=args.answer) + chat = ChatDB.from_dir(Path('.'), + Path(config.db), + mfilter=mfilter) + chat.print(args.source_code_only, + args.with_tags, + args.with_files) diff --git a/chatmastermind/commands/print.py b/chatmastermind/commands/print.py new file mode 100644 index 0000000..51e76f8 --- /dev/null +++ b/chatmastermind/commands/print.py @@ -0,0 +1,19 @@ +import sys +import argparse +from pathlib import Path +from ..configuration import Config +from ..message import Message, MessageError + + +def print_cmd(args: argparse.Namespace, config: Config) -> None: + """ + Handler for the 'print' command. + """ + fname = Path(args.file) + try: + message = Message.from_file(fname) + if message: + print(message.to_str(source_code_only=args.source_code_only)) + except MessageError: + print(f"File is not a valid message: {args.file}") + sys.exit(1) diff --git a/chatmastermind/commands/question.py b/chatmastermind/commands/question.py new file mode 100644 index 0000000..1d926f6 --- /dev/null +++ b/chatmastermind/commands/question.py @@ -0,0 +1,57 @@ +import argparse +from pathlib import Path +from ..configuration import Config +from ..chat import ChatDB +from ..message import Message, Question +from ..ai_factory import create_ai +from ..ai import AI, AIResponse + + +def create_message(chat: ChatDB, args: argparse.Namespace) -> Message: + """ + Creates (and writes) a new message from the given arguments. + """ + # FIXME: add sources to the question + message = Message(question=Question(args.question), + tags=args.ouput_tags, # FIXME + ai=args.ai, + model=args.model) + chat.add_to_cache([message]) + return message + + +def question_cmd(args: argparse.Namespace, config: Config) -> None: + """ + Handler for the 'question' command. + """ + chat = ChatDB.from_dir(cache_path=Path('.'), + db_path=Path(config.db)) + # if it's a new question, create and store it immediately + if args.ask or args.create: + message = create_message(chat, args) + if args.create: + return + + # create the correct AI instance + ai: AI = create_ai(args, config) + if args.ask: + response: AIResponse = ai.request(message, + chat, + args.num_answers, # FIXME + args.otags) # FIXME + assert response + # TODO: + # * add answer to the message above (and create + # more messages for any additional answers) + pass + elif args.repeat: + lmessage = chat.latest_message() + assert lmessage + # TODO: repeat either the last question or the + # one(s) given in 'args.repeat' (overwrite + # existing ones if 'args.overwrite' is True) + pass + elif args.process: + # TODO: process either all questions without an + # answer or the one(s) given in 'args.process' + pass diff --git a/chatmastermind/commands/tags.py b/chatmastermind/commands/tags.py new file mode 100644 index 0000000..2906a5b --- /dev/null +++ b/chatmastermind/commands/tags.py @@ -0,0 +1,17 @@ +import argparse +from pathlib import Path +from ..configuration import Config +from ..chat import ChatDB + + +def tags_cmd(args: argparse.Namespace, config: Config) -> None: + """ + Handler for the 'tags' command. + """ + chat = ChatDB.from_dir(cache_path=Path('.'), + db_path=Path(config.db)) + if args.list: + tags_freq = chat.tags_frequency(args.prefix, args.contain) + for tag, freq in tags_freq.items(): + print(f"- {tag}: {freq}") + # TODO: add renaming diff --git a/chatmastermind/main.py b/chatmastermind/main.py index 857bb5a..88121b4 100755 --- a/chatmastermind/main.py +++ b/chatmastermind/main.py @@ -6,12 +6,14 @@ import sys import argcomplete import argparse from pathlib import Path -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 typing import Any +from .configuration import Config, default_config_path +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 def tags_completer(prefix: str, parsed_args: Any, **kwargs: Any) -> list[str]: @@ -19,101 +21,6 @@ def tags_completer(prefix: str, parsed_args: Any, **kwargs: Any) -> list[str]: return list(Message.tags_from_dir(Path(config.db), prefix=prefix)) -def tags_cmd(args: argparse.Namespace, config: Config) -> None: - """ - Handler for the 'tags' command. - """ - chat = ChatDB.from_dir(cache_path=Path('.'), - db_path=Path(config.db)) - if args.list: - tags_freq = chat.tags_frequency(args.prefix, args.contain) - for tag, freq in tags_freq.items(): - print(f"- {tag}: {freq}") - # TODO: add renaming - - -def config_cmd(args: argparse.Namespace) -> None: - """ - Handler for the 'config' command. - """ - if args.create: - Config.create_default(Path(args.create)) - - -def question_cmd(args: argparse.Namespace, config: Config) -> None: - """ - Handler for the 'question' command. - """ - chat = ChatDB.from_dir(cache_path=Path('.'), - db_path=Path(config.db)) - # if it's a new question, create and store it immediately - if args.ask or args.create: - # FIXME: add sources to the question - message = Message(question=Question(args.question), - tags=args.ouput_tags, # FIXME - ai=args.ai, - model=args.model) - chat.add_to_cache([message]) - if args.create: - return - - # create the correct AI instance - ai: AI = create_ai(args, config) - if args.ask: - response: AIResponse = ai.request(message, - chat, - args.num_answers, # FIXME - args.otags) # FIXME - assert response - # TODO: - # * add answer to the message above (and create - # more messages for any additional answers) - pass - elif args.repeat: - lmessage = chat.latest_message() - assert lmessage - # TODO: repeat either the last question or the - # one(s) given in 'args.repeat' (overwrite - # existing ones if 'args.overwrite' is True) - pass - elif args.process: - # TODO: process either all questions without an - # answer or the one(s) given in 'args.process' - pass - - -def hist_cmd(args: argparse.Namespace, config: Config) -> None: - """ - Handler for the 'hist' command. - """ - - mfilter = MessageFilter(tags_or=args.or_tags, - tags_and=args.and_tags, - tags_not=args.exclude_tags, - question_contains=args.question, - answer_contains=args.answer) - chat = ChatDB.from_dir(Path('.'), - Path(config.db), - mfilter=mfilter) - chat.print(args.source_code_only, - args.with_tags, - args.with_files) - - -def print_cmd(args: argparse.Namespace, config: Config) -> None: - """ - Handler for the 'print' command. - """ - fname = Path(args.file) - try: - message = Message.from_file(fname) - if message: - print(message.to_str(source_code_only=args.source_code_only)) - except MessageError: - print(f"File is not a valid message: {args.file}") - sys.exit(1) - - def create_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="ChatMastermind is a Python application that automates conversation with AI") @@ -128,20 +35,28 @@ def create_parser() -> argparse.ArgumentParser: # 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 tag names (one must match)', metavar='OTAGS') + 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 tag names (all must match)', metavar='ATAGS') + 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 tag names to exclude', metavar='XTAGS') + 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 tag names, default is input', metavar='OUTTAGS') + help='List of output tags (default: use input tags)', metavar='OUTTAGS') 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') + ai_parser.add_argument('-M', '--model', help='Model to use') + 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], + 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) @@ -152,12 +67,6 @@ def create_parser() -> argparse.ArgumentParser: question_group.add_argument('-p', '--process', nargs='*', help='Process existing questions') question_cmd_parser.add_argument('-O', '--overwrite', help='Overwrite existing messages when repeating them', action='store_true') - question_cmd_parser.add_argument('-m', '--max-tokens', help='Max tokens to use', type=int) - question_cmd_parser.add_argument('-T', '--temperature', help='Temperature to use', type=float) - question_cmd_parser.add_argument('-A', '--AI', help='AI to use') - question_cmd_parser.add_argument('-M', '--model', help='Model to use') - question_cmd_parser.add_argument('-n', '--num-answers', help='Number of answers to produce', type=int, - default=1) question_cmd_parser.add_argument('-s', '--source', nargs='+', help='Source add content of a file to the query') question_cmd_parser.add_argument('-S', '--source-code-only', help='Add pure source code to the chat history', action='store_true') diff --git a/setup.py b/setup.py index 8484629..a311605 100644 --- a/setup.py +++ b/setup.py @@ -12,7 +12,7 @@ setup( long_description=long_description, long_description_content_type="text/markdown", url="https://github.com/ok2/ChatMastermind", - packages=find_packages() + ["chatmastermind.ais"], + packages=find_packages() + ["chatmastermind.ais", "chatmastermind.commands"], classifiers=[ "Development Status :: 3 - Alpha", "Environment :: Console", diff --git a/tests/test_ai_factory.py b/tests/test_ai_factory.py new file mode 100644 index 0000000..d63970e --- /dev/null +++ b/tests/test_ai_factory.py @@ -0,0 +1,48 @@ +import argparse +import unittest +from unittest.mock import MagicMock +from chatmastermind.ai_factory import create_ai +from chatmastermind.configuration import Config +from chatmastermind.ai import AIError +from chatmastermind.ais.openai import OpenAI + + +class TestCreateAI(unittest.TestCase): + def setUp(self) -> None: + self.args = MagicMock(spec=argparse.Namespace) + self.args.ai = 'default' + self.args.model = None + self.args.max_tokens = None + self.args.temperature = None + + def test_create_ai_from_args(self) -> None: + # Create an AI with the default configuration + config = Config() + self.args.ai = 'default' + ai = create_ai(self.args, config) + self.assertIsInstance(ai, OpenAI) + + def test_create_ai_from_default(self) -> None: + self.args.ai = None + # Create an AI with the default configuration + config = Config() + ai = create_ai(self.args, config) + self.assertIsInstance(ai, OpenAI) + + def test_create_empty_ai_error(self) -> None: + self.args.ai = None + # Create Config with empty AIs + config = Config() + config.ais = {} + # Call create_ai function and assert that it raises AIError + with self.assertRaises(AIError): + create_ai(self.args, config) + + def test_create_unsupported_ai_error(self) -> None: + # Mock argparse.Namespace with ai='invalid_ai' + self.args.ai = 'invalid_ai' + # Create default Config + config = Config() + # Call create_ai function and assert that it raises AIError + with self.assertRaises(AIError): + create_ai(self.args, config)