Improve source code extraction

chatmastermind/utils.py

@@ -51,12 +51,13 @@ def display_chat(chat, dump=False, source_code=False) -> None:
             if not source_code:
                 print(f"{message['role'].upper()}:")
             if source_code:
-                out = 0
+                try:
-                for line in message['content'].splitlines():
+                    content_start = message['content'].index('```')
-                    if line.strip().startswith('```'):
+                    content_end = message['content'].rindex('```')
-                        out += 1
+                    if content_start + 3 < content_end:
-                    elif out == 1:
+                        print(message['content'][content_start+3:content_end].strip())
-                        print(f"{line}")
+                except ValueError:
+                    pass
             else:
                 print(message['content'])
         elif not source_code: