From ad10af8f4f23d7a3d32c0ce2f1a7d5172329036b Mon Sep 17 00:00:00 2001
From: Oleksandr Kozachuk <ddeus.lp@mailnull.com>
Date: Sat, 8 Apr 2023 11:32:29 +0200
Subject: [PATCH] Improve source code extraction

---
 chatmastermind/utils.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/chatmastermind/utils.py b/chatmastermind/utils.py
index d0d05ae..5f4bcad 100644
--- a/chatmastermind/utils.py
+++ b/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
-                for line in message['content'].splitlines():
-                    if line.strip().startswith('```'):
-                        out += 1
-                    elif out == 1:
-                        print(f"{line}")
+                try:
+                    content_start = message['content'].index('```')
+                    content_end = message['content'].rindex('```')
+                    if content_start + 3 < content_end:
+                        print(message['content'][content_start+3:content_end].strip())
+                except ValueError:
+                    pass
             else:
                 print(message['content'])
         elif not source_code: