Improve handling of printing source code only output.

chatmastermind/utils.py

@@ -40,24 +40,28 @@ def message_to_chat(message: Dict[str, str],
     append_message(chat, 'assistant', message['answer'])
 
 
+def display_source_code(content: str) -> None:
+    code_block_count = 0
+    for line in content.splitlines():
+        if line.strip().startswith('```'):
+            code_block_count += 1
+        elif code_block_count == 1:
+            print(line)
+
+
 def display_chat(chat, dump=False, source_code=False) -> None:
     if dump:
         pp(chat)
         return
     for message in chat:
-        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()}:")
-            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:
+        text_too_long = len(message['content']) > terminal_width() - len(message['role']) - 2
+        if source_code:
+            display_source_code(message['content'])
+            continue
+        if message['role'] == 'user':
+            print('-' * terminal_width())
+        if text_too_long:
+            print(f"{message['role'].upper()}:")
+            print(message['content'])
+        else:
             print(f"{message['role'].upper()}: {message['content']}")