Merge branch 'main' of fstage:ok/ChatMastermind

chatmastermind/main.py

@@ -6,16 +6,27 @@ import yaml
 import sys
 import argcomplete
 import argparse
-from .utils import terminal_width, pp, process_tags, display_chat
-from .storage import save_answers, create_chat, get_tags
+import pathlib
+from .utils import terminal_width, process_tags, display_chat, display_source_code
+from .storage import save_answers, create_chat, get_tags, read_file, dump_data
 from .api_client import ai, openai_api_key
 from itertools import zip_longest
 
 
 def run_print_command(args: argparse.Namespace, config: dict) -> None:
-    with open(args.print, 'r') as f:
-        data = yaml.load(f, Loader=yaml.FullLoader)
-    pp(data)
+    fname = pathlib.Path(args.print)
+    if fname.suffix == '.yaml':
+        with open(args.print, 'r') as f:
+            data = yaml.load(f, Loader=yaml.FullLoader)
+    elif fname.suffix == '.txt':
+        data = read_file(fname)
+    else:
+        print(f"Unknown file type: {args.print}")
+        sys.exit(1)
+    if args.only_source_code:
+        display_source_code(data['answer'])
+    else:
+        print(dump_data(data).strip())
 
 
 def process_and_display_chat(args: argparse.Namespace,
@@ -72,7 +83,7 @@ def create_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(
         description="ChatMastermind is a Python application that automates conversation with AI")
     group = parser.add_mutually_exclusive_group(required=True)
-    group.add_argument('-p', '--print', help='YAML file to print')
+    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', help="Print chat as readable text", action='store_true')
@@ -80,7 +91,7 @@ def create_parser() -> argparse.ArgumentParser:
     parser.add_argument('-m', '--max-tokens', help='Max tokens to use', type=int)
     parser.add_argument('-T', '--temperature', help='Temperature to use', type=float)
     parser.add_argument('-M', '--model', help='Model to use')
-    parser.add_argument('-n', '--number', help='Number of answers to produce', type=int, default=3)
+    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')
     tags_arg = parser.add_argument('-t', '--tags', nargs='*', help='List of tag names', metavar='TAGS')