Merge pull request 'Changes from https://github.com/juk0de/ChatMastermind.git' (#2) from juk into main

Reviewed-on: https://stage.fjerkroa.no/gitea/ok/ChatMastermind/pulls/2

README.md

@@ -28,15 +28,19 @@ 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 | -l] [-c CONFIG] [-m MAX_TOKENS] [-T TEMPERATURE] [-M MODEL] [-n NUMBER] [-t [TAGS [TAGS ...]]] [-e [EXTAGS [EXTAGS ...]]] [-o [OTAGS [OTAGS ...]]] [-a] [-w] [-W]
 ```
 
 ### Arguments
 
 - `-p`, `--print`: YAML file to print.
 - `-q`, `--question`: Question to ask.
-- `-D`, `--chat-dump`: Print chat as a Python structure.
+- `-D`, `--chat-dump`: Print chat history as a Python structure.
-- `-d`, `--chat`: Print chat as readable text.
+- `-d`, `--chat`: Print chat history as readable text.
+- `-a`, `--match-all-tags`: All given tags must match when selecting chat history entries.
+- `-w`, `--with-tags`: Print chat history with tags.
+- `-W`, `--with-tags`: Print chat history with filenames.
+- `-l`, `--list-tags`: List all tags and their frequency.
 - `-c`, `--config`: Config file name (defaults to `.config.yaml`).
 - `-m`, `--max-tokens`: Max tokens to use.
 - `-T`, `--temperature`: Temperature to use.
@@ -111,4 +115,4 @@ After adding this line, restart your shell or run `source <your-shell-config-fil
 
 ## License
 
-This project is licensed under the terms of the WTFPL License.
+This project is licensed under the terms of the WTFPL License.

chatmastermind/main.py

@@ -7,8 +7,8 @@ import sys
 import argcomplete
 import argparse
 import pathlib
-from .utils import terminal_width, process_tags, display_chat, display_source_code
+from .utils import terminal_width, process_tags, display_chat, display_source_code, display_tags_frequency
-from .storage import save_answers, create_chat, get_tags, read_file, dump_data
+from .storage import save_answers, create_chat, get_tags, get_tags_unique, read_file, dump_data
 from .api_client import ai, openai_api_key
 from itertools import zip_longest
 
@@ -55,10 +55,18 @@ def process_and_display_chat(args: argparse.Namespace,
                 question_parts.append(f"```\n{r.read().strip()}\n```")
 
     full_question = '\n\n'.join(question_parts)
-    chat = create_chat(full_question, tags, extags, config)
+    chat = create_chat(full_question, tags, extags, config,
+                       args.match_all_tags, args.with_tags,
+                       args.with_file)
     display_chat(chat, dump, args.only_source_code)
     return chat, full_question, tags
 
+def process_and_display_tags(args: argparse.Namespace,
+                             config: dict,
+                             dump: bool=False
+                             ) -> None:
+    display_tags_frequency(get_tags(config, None), dump)
+
 
 def handle_question(args: argparse.Namespace,
                     config: dict,
@@ -75,7 +83,7 @@ def handle_question(args: argparse.Namespace,
 def tags_completer(prefix, parsed_args, **kwargs):
     with open(parsed_args.config, 'r') as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
-    return get_tags(config, prefix)
+    return get_tags_unique(config, prefix)
 
 
 def create_parser() -> argparse.ArgumentParser:
@@ -85,8 +93,9 @@ def create_parser() -> argparse.ArgumentParser:
     group = parser.add_mutually_exclusive_group(required=True)
     group.add_argument('-p', '--print', help='File to print')
     group.add_argument('-q', '--question', nargs='*', help='Question to ask')
-    group.add_argument('-D', '--chat-dump', help="Print chat as Python structure", action='store_true')
+    group.add_argument('-D', '--chat-dump', help="Print chat history as Python structure", action='store_true')
-    group.add_argument('-d', '--chat', help="Print chat as readable text", action='store_true')
+    group.add_argument('-d', '--chat', help="Print chat history as readable text", action='store_true')
+    group.add_argument('-l', '--list-tags', help="List all tags and their frequency", action='store_true')
     parser.add_argument('-c', '--config', help='Config file name.', default=default_config)
     parser.add_argument('-m', '--max-tokens', help='Max tokens to use', type=int)
     parser.add_argument('-T', '--temperature', help='Temperature to use', type=float)
@@ -94,6 +103,9 @@ def create_parser() -> 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 history with tags.", action='store_true')
+    parser.add_argument('-W', '--with-file', help="Print chat history with filename.", action='store_true')
+    parser.add_argument('-a', '--match-all-tags', help="All given tags must match when selecting chat history entries.", 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')
@@ -130,6 +142,8 @@ def main() -> int:
         process_and_display_chat(args, config, dump=True)
     elif args.chat:
         process_and_display_chat(args, config)
+    elif args.list_tags:
+        process_and_display_tags(args, config)
 
     return 0
 

chatmastermind/storage.py

@@ -15,7 +15,8 @@ def read_file(fname: str, tags_only: bool = False) -> Dict[str, Any]:
         answer_idx = text.index("==== ANSWER ====")
         question = "\n".join(text[question_idx:answer_idx]).strip()
         answer = "\n".join(text[answer_idx + 1:]).strip()
-        return {"question": question, "answer": answer, "tags": tags}
+        return {"question": question, "answer": answer, "tags": tags,
+                "file": pathlib.Path(fname).name}
 
 
 def dump_data(data: Dict[str, Any]) -> str:
@@ -62,7 +63,10 @@ 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],
+                match_all_tags: bool = False,
+                with_tags: bool = False,
+                with_file: bool = False
                 ) -> List[Dict[str, str]]:
     chat: List[Dict[str, str]] = []
     append_message(chat, 'system', config['system'].strip())
@@ -75,12 +79,15 @@ def create_chat(question: Optional[str],
         else:
             continue
         data_tags = set(data.get('tags', []))
-        tags_match = \
+        tags_match: bool
-            not tags or data_tags.intersection(tags)
+        if match_all_tags:
+            tags_match = not tags or set(tags).issubset(data_tags)
+        else:
+            tags_match = not tags or bool(data_tags.intersection(tags))
         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, with_file)
     if question:
         append_message(chat, 'user', question)
     return chat
@@ -102,4 +109,7 @@ def get_tags(config: Dict[str, Any], prefix: Optional[str]) -> List[str]:
                     result.append(tag)
             else:
                 result.append(tag)
-    return list(set(result))
+    return result
+
+def get_tags_unique(config: Dict[str, Any], prefix: Optional[str]) -> List[str]:
+    return list(set(get_tags(config, prefix)))

chatmastermind/utils.py

@@ -34,10 +34,17 @@ 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,
+                    with_file: 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)
+    if with_file:
+        append_message(chat, 'file', message['file'])
 
 
 def display_source_code(content: str) -> None:
@@ -66,3 +73,10 @@ def display_chat(chat, dump=False, source_code=False) -> None:
             print(message['content'])
         else:
             print(f"{message['role'].upper()}: {message['content']}")
+
+def display_tags_frequency(tags: List[str], dump=False) -> None:
+    if dump:
+        pp(tags)
+        return
+    for tag in set(tags):
+        print(f"-{tag} : {tags.count(tag)}")