diff --git a/README.md b/README.md index 57067d2..e42d1df 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ pip install . ## Usage ```bash -cmm [-h] [-p PRINT | -q QUESTION | -D | -d] [-c CONFIG] [-m MAX_TOKENS] [-T TEMPERATURE] [-M MODEL] [-n NUMBER] [-t [TAGS [TAGS ...]]] [-e [EXTAGS [EXTAGS ...]]] [-o [OTAGS [OTAGS ...]]] +cmm [-h] [-p PRINT | -q QUESTION | -D | -d] [-c CONFIG] [-m MAX_TOKENS] [-T TEMPERATURE] [-M MODEL] [-n NUMBER] [-t [TAGS [TAGS ...]]] [-e [EXTAGS [EXTAGS ...]]] [-o [OTAGS [OTAGS ...]]] [-w] ``` ### Arguments @@ -37,6 +37,7 @@ cmm [-h] [-p PRINT | -q QUESTION | -D | -d] [-c CONFIG] [-m MAX_TOKENS] [-T TEMP - `-q`, `--question`: Question to ask. - `-D`, `--chat-dump`: Print chat as a Python structure. - `-d`, `--chat`: Print chat as readable text. +- `-w`, `--with-tags`: Print chat with tags. - `-c`, `--config`: Config file name (defaults to `.config.yaml`). - `-m`, `--max-tokens`: Max tokens to use. - `-T`, `--temperature`: Temperature to use. @@ -111,4 +112,4 @@ After adding this line, restart your shell or run `source argparse.ArgumentParser: parser.add_argument('-n', '--number', help='Number of answers to produce', type=int, default=1) parser.add_argument('-s', '--source', nargs='*', help='Source add content of a file to the query') parser.add_argument('-S', '--only-source-code', help='Print only source code', action='store_true') + parser.add_argument('-w', '--with-tags', help="Print chat with tags.", action='store_true') tags_arg = parser.add_argument('-t', '--tags', nargs='*', help='List of tag names', metavar='TAGS') tags_arg.completer = tags_completer # type: ignore extags_arg = parser.add_argument('-e', '--extags', nargs='*', help='List of tag names to exclude', metavar='EXTAGS') diff --git a/chatmastermind/storage.py b/chatmastermind/storage.py index 4215e5a..9fa3a84 100644 --- a/chatmastermind/storage.py +++ b/chatmastermind/storage.py @@ -62,7 +62,8 @@ def save_answers(question: str, def create_chat(question: Optional[str], tags: Optional[List[str]], extags: Optional[List[str]], - config: Dict[str, Any] + config: Dict[str, Any], + with_tags: bool = False ) -> List[Dict[str, str]]: chat: List[Dict[str, str]] = [] append_message(chat, 'system', config['system'].strip()) @@ -80,7 +81,7 @@ def create_chat(question: Optional[str], extags_do_not_match = \ not extags or not data_tags.intersection(extags) if tags_match and extags_do_not_match: - message_to_chat(data, chat) + message_to_chat(data, chat, with_tags) if question: append_message(chat, 'user', question) return chat diff --git a/chatmastermind/utils.py b/chatmastermind/utils.py index 4db8f0c..c4270c7 100644 --- a/chatmastermind/utils.py +++ b/chatmastermind/utils.py @@ -34,10 +34,14 @@ def append_message(chat: List[Dict[str, str]], def message_to_chat(message: Dict[str, str], - chat: List[Dict[str, str]] + chat: List[Dict[str, str]], + with_tags: bool = False ) -> None: append_message(chat, 'user', message['question']) append_message(chat, 'assistant', message['answer']) + if with_tags: + tags = ", ".join(message['tags']) + append_message(chat, 'tags', tags) def display_source_code(content: str) -> None: