From 031d94a1bbae291cc5f235c576fc20439c3b97b7 Mon Sep 17 00:00:00 2001
From: Emery Berger <emery.berger@gmail.com>
Date: Sat, 30 Sep 2023 14:30:36 -0400
Subject: [PATCH] Updated to handle when ``` is not at start of line. Fixes
 https://github.com/plasma-umass/ChatDBG/issues/12.

---
 src/chatdbg/chatdbg_utils.py | 48 +++++++++++++++++++++++++++++++++++-
 1 file changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/chatdbg/chatdbg_utils.py b/src/chatdbg/chatdbg_utils.py
index a180ad6..79fcc98 100644
--- a/src/chatdbg/chatdbg_utils.py
+++ b/src/chatdbg/chatdbg_utils.py
@@ -18,8 +18,54 @@ def get_model() -> str:
         
     return model
 
-    
+
 def word_wrap_except_code_blocks(text: str) -> str:
+    """
+    Wraps text except for code blocks.
+
+    Splits the text into paragraphs and wraps each paragraph,
+    except for paragraphs that are inside of code blocks denoted
+    by ` ``` `. Returns the updated text.
+
+    Args:
+        text: The text to wrap.
+
+    Returns:
+        The wrapped text.
+    """
+    # Split text into paragraphs
+    paragraphs = text.split("\n\n")
+    wrapped_paragraphs = []
+    # Check if currently in a code block.
+    in_code_block = False
+    # Loop through each paragraph and apply appropriate wrapping.
+    for paragraph in paragraphs:
+        # Check for the presence of triple quotes in the paragraph
+        if "```" in paragraph:
+            # Split paragraph by triple quotes
+            parts = paragraph.split("```")
+            for i, part in enumerate(parts):
+                # If we are inside a code block, do not wrap the text
+                if in_code_block:
+                    wrapped_paragraphs.append(part)
+                else:
+                    # Otherwise, apply text wrapping to the part
+                    wrapped_paragraphs.append(textwrap.fill(part))
+                # Toggle the in_code_block flag for each triple quote encountered
+                if i < len(parts) - 1:
+                    wrapped_paragraphs.append("```")
+                    in_code_block = not in_code_block
+        else:
+            # If the paragraph does not contain triple quotes and is not inside a code block, wrap the text
+            if not in_code_block:
+                wrapped_paragraphs.append(textwrap.fill(paragraph))
+            else:
+                wrapped_paragraphs.append(paragraph)
+    # Join all paragraphs into a single string
+    wrapped_text = "\n\n".join(wrapped_paragraphs)
+    return wrapped_text
+
+def word_wrap_except_code_blocks_previous(text: str) -> str:
     """Wraps text except for code blocks.
 
     Splits the text into paragraphs and wraps each paragraph,