From b26a5a37921bd357ea1605b250d3426429d64a5b Mon Sep 17 00:00:00 2001 From: juk0de Date: Sat, 2 Sep 2023 08:21:49 +0200 Subject: [PATCH] cmm: the 'hist' command now uses the new 'ChatDB' --- chatmastermind/main.py | 60 +++++++++++++++++++++++------------------- tests/test_main.py | 15 ++++++----- 2 files changed, 42 insertions(+), 33 deletions(-) diff --git a/chatmastermind/main.py b/chatmastermind/main.py index 3f31aee..08c5e3e 100755 --- a/chatmastermind/main.py +++ b/chatmastermind/main.py @@ -12,6 +12,7 @@ from .storage import save_answers, create_chat_hist, get_tags_unique, read_file, from .api_client import ai, openai_api_key, print_models from .configuration import Config from .chat import ChatDB +from .message import MessageFilter from itertools import zip_longest from typing import Any @@ -32,11 +33,11 @@ def create_question_with_hist(args: argparse.Namespace, by the specified tags. """ tags = args.tags or [] - extags = args.extags or [] + etags = args.etags or [] otags = args.output_tags or [] - if not args.only_source_code: - print_tag_args(tags, extags, otags) + if not args.source_code_only: + print_tag_args(tags, etags, otags) question_parts = [] question_list = args.question if args.question is not None else [] @@ -53,8 +54,10 @@ def create_question_with_hist(args: argparse.Namespace, question_parts.append(f"```\n{r.read().strip()}\n```") full_question = '\n\n'.join(question_parts) - chat = create_chat_hist(full_question, tags, extags, config, - args.match_all_tags, False, False) + chat = create_chat_hist(full_question, tags, etags, config, + match_all_tags=True if args.atags else False, # FIXME + with_tags=False, + with_file=False) return chat, full_question, tags @@ -95,7 +98,7 @@ def ask_cmd(args: argparse.Namespace, config: Config) -> None: if args.model: config.openai.model = args.model chat, question, tags = create_question_with_hist(args, config) - print_chat_hist(chat, False, args.only_source_code) + print_chat_hist(chat, False, args.source_code_only) otags = args.output_tags or [] answers, usage = ai(chat, config, args.number) save_answers(question, answers, tags, otags, config) @@ -107,14 +110,18 @@ def hist_cmd(args: argparse.Namespace, config: Config) -> None: """ Handler for the 'hist' command. """ - tags = args.tags or [] - extags = args.extags or [] - chat = create_chat_hist(None, tags, extags, config, - args.match_all_tags, - args.with_tags, - args.with_files) - print_chat_hist(chat, args.dump, args.only_source_code) + mfilter = MessageFilter(tags_or=args.tags, + tags_and=args.atags, + tags_not=args.etags, + 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: @@ -130,7 +137,7 @@ def print_cmd(args: argparse.Namespace, config: Config) -> None: else: print(f"Unknown file type: {args.file}") sys.exit(1) - if args.only_source_code: + if args.source_code_only: display_source_code(data['answer']) else: print(dump_data(data).strip()) @@ -150,18 +157,17 @@ 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', '--tags', nargs='+', - help='List of tag names', metavar='TAGS') + help='List of tag names (one must match)', metavar='TAGS') tag_arg.completer = tags_completer # type: ignore - extag_arg = tag_parser.add_argument('-e', '--extags', nargs='+', - help='List of tag names to exclude', metavar='EXTAGS') - extag_arg.completer = tags_completer # type: ignore + atag_arg = tag_parser.add_argument('-a', '--atags', nargs='+', + help='List of tag names (all must match)', metavar='TAGS') + atag_arg.completer = tags_completer # type: ignore + etag_arg = tag_parser.add_argument('-e', '--etags', nargs='+', + help='List of tag names to exclude', metavar='ETAGS') + 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='OTAGS') otag_arg.completer = tags_completer # type: ignore - tag_parser.add_argument('-a', '--match-all-tags', - help="All given tags must match when selecting chat history entries", - action='store_true') - # enable autocompletion for tags # 'ask' command parser ask_cmd_parser = cmdparser.add_parser('ask', parents=[tag_parser], @@ -176,7 +182,7 @@ def create_parser() -> argparse.ArgumentParser: ask_cmd_parser.add_argument('-n', '--number', 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', '--only-source-code', help='Add pure source code to the chat history', + ask_cmd_parser.add_argument('-S', '--source-code-only', help='Add pure source code to the chat history', action='store_true') # 'hist' command parser @@ -184,14 +190,14 @@ def create_parser() -> argparse.ArgumentParser: help="Print chat history.", aliases=['h']) hist_cmd_parser.set_defaults(func=hist_cmd) - hist_cmd_parser.add_argument('-d', '--dump', help="Print chat history as Python structure", - action='store_true') hist_cmd_parser.add_argument('-w', '--with-tags', help="Print chat history with tags.", action='store_true') hist_cmd_parser.add_argument('-W', '--with-files', help="Print chat history with filenames.", action='store_true') - hist_cmd_parser.add_argument('-S', '--only-source-code', help='Print only source code', + hist_cmd_parser.add_argument('-S', '--source-code-only', help='Print only source code', action='store_true') + hist_cmd_parser.add_argument('-A', '--answer', help='Search for answer substring') + hist_cmd_parser.add_argument('-Q', '--question', help='Search for question substring') # 'tags' command parser tags_cmd_parser = cmdparser.add_parser('tags', @@ -222,7 +228,7 @@ def create_parser() -> argparse.ArgumentParser: aliases=['p']) print_cmd_parser.set_defaults(func=print_cmd) print_cmd_parser.add_argument('-f', '--file', help='File to print', required=True) - print_cmd_parser.add_argument('-S', '--only-source-code', help='Print only source code', + print_cmd_parser.add_argument('-S', '--source-code-only', help='Print only source code', action='store_true') argcomplete.autocomplete(parser) diff --git a/tests/test_main.py b/tests/test_main.py index 23c3d00..bb9aa2a 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -115,11 +115,12 @@ class TestHandleQuestion(CmmTestCase): self.question = "test question" self.args = argparse.Namespace( tags=['tag1'], - extags=['extag1'], + atags=None, + etags=['etag1'], output_tags=None, question=[self.question], source=None, - only_source_code=False, + source_code_only=False, number=3, max_tokens=None, temperature=None, @@ -143,16 +144,18 @@ class TestHandleQuestion(CmmTestCase): with patch("chatmastermind.storage.open", open_mock): ask_cmd(self.args, self.config) mock_print_tag_args.assert_called_once_with(self.args.tags, - self.args.extags, + self.args.etags, []) mock_create_chat_hist.assert_called_once_with(self.question, self.args.tags, - self.args.extags, + self.args.etags, self.config, - False, False, False) + match_all_tags=False, + with_tags=False, + with_file=False) mock_print_chat_hist.assert_called_once_with('test_chat', False, - self.args.only_source_code) + self.args.source_code_only) mock_ai.assert_called_with("test_chat", self.config, self.args.number)