Add possibility to print just the source code

chatmastermind/main.py

@@ -24,6 +24,8 @@ def process_and_display_chat(args: argparse.Namespace,
     tags = args.tags or []
     extags = args.extags or []
     otags = args.output_tags or []
+
+    if not args.only_source_code:
         process_tags(tags, extags, otags)
 
     question_parts = []
@@ -45,7 +47,7 @@ def process_and_display_chat(args: argparse.Namespace,
     question = '\n\n'.join(question_parts)
 
     chat = create_chat(question, tags, extags, config)
-    display_chat(chat, dump)
+    display_chat(chat, dump, args.only_source_code)
     return chat, question, tags
 
 
@@ -82,6 +84,7 @@ def create_parser() -> argparse.ArgumentParser:
     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('-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')
     tags_arg.completer = tags_completer  # type: ignore
     extags_arg = parser.add_argument('-e', '--extags', nargs='*', help='List of tag names to exclude', metavar='EXTAGS')

chatmastermind/utils.py

@@ -40,15 +40,24 @@ def message_to_chat(message: Dict[str, str],
     append_message(chat, 'assistant', message['answer'])
 
 
-def display_chat(chat, dump=False) -> None:
+def display_chat(chat, dump=False, source_code=False) -> None:
     if dump:
         pp(chat)
         return
     for message in chat:
-        if message['role'] == 'user':
+        if message['role'] == 'user' and not source_code:
             print('-' * (terminal_width()))
         if len(message['content']) > terminal_width() - len(message['role']) - 2:
+            if not source_code:
                 print(f"{message['role'].upper()}:")
-            print(message['content'])
+            if source_code:
+                out = 0
+                for line in message['content'].splitlines():
+                    if line.strip().startswith('```'):
+                        out += 1
+                    elif out == 1:
+                        print(f"{line}")
             else:
+                print(message['content'])
+        elif not source_code:
             print(f"{message['role'].upper()}: {message['content']}")

tests/test_main.py

@@ -95,6 +95,7 @@ class TestHandleQuestion(unittest.TestCase):
             output_tags=None,
             question=[self.question],
             source=None,
+            only_source_code=False,
             number=3
         )
         self.config = {