From eed180e7a614f5c20a3fb7b91fcf05fc878bafdf Mon Sep 17 00:00:00 2001
From: Oleksandr Kozachuk <ddeus.lp@mailnull.com>
Date: Fri, 7 Apr 2023 23:42:26 +0200
Subject: [PATCH] Add possibility to print just the source code

---
 chatmastermind/main.py  |  7 +++++--
 chatmastermind/utils.py | 19 ++++++++++++++-----
 tests/test_main.py      |  1 +
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/chatmastermind/main.py b/chatmastermind/main.py
index 27e00f1..1c937d4 100755
--- a/chatmastermind/main.py
+++ b/chatmastermind/main.py
@@ -24,7 +24,9 @@ def process_and_display_chat(args: argparse.Namespace,
     tags = args.tags or []
     extags = args.extags or []
     otags = args.output_tags or []
-    process_tags(tags, extags, otags)
+
+    if not args.only_source_code:
+        process_tags(tags, extags, otags)
 
     question_parts = []
     question_list = args.question if args.question is not None else []
@@ -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')
diff --git a/chatmastermind/utils.py b/chatmastermind/utils.py
index 2c07ce3..d0d05ae 100644
--- a/chatmastermind/utils.py
+++ b/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:
-            print(f"{message['role'].upper()}:")
-            print(message['content'])
-        else:
+            if not source_code:
+                print(f"{message['role'].upper()}:")
+            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']}")
diff --git a/tests/test_main.py b/tests/test_main.py
index d3c7755..5b5d178 100644
--- a/tests/test_main.py
+++ b/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 = {