Updated cleaning system

README.md

@@ -1,8 +1,8 @@
 ---
 title: CodeFlow
 emoji: 📊
-colorFrom: pink
-colorTo: purple
+colorFrom: indigo
+colorTo: cyan
 sdk: gradio
 sdk_version: 6.16.0
 python_version: '3.13'

app.py

@@ -35,6 +35,24 @@ llm = Llama(
 app = gr.Server(title="Code-to-Flowchart Generator")
 
 # ----- Functions ----- #
+
+# This is a cleaning function to resolve common syntax errors.
+def quote_labels(text: str) -> str:
+    # Mermaid node labels can't hold raw code char, so quote-wrap each label body.
+    END = r'(?=\s*(?:--|==|-\.|<|\||;|$))'
+
+    def esc(body: str) -> str:
+        return (body.replace('"', "'")
+                    .replace('[', '&#91;').replace(']', '&#93;')
+                    .replace('{', '&#123;').replace('}', '&#125;'))
+
+    out = []
+    for line in text.split('\n'):
+        line = re.sub(r'(?<=\w)\[(.*?)\]' + END, lambda m: '["' + esc(m.group(1)) + '"]', line)
+        line = re.sub(r'(?<=\w)\{(.*?)\}' + END, lambda m: '{"' + esc(m.group(1)) + '"}', line)
+        out.append(line)
+    return '\n'.join(out)
+
 @app.api(name="generate_flowchart")
 def generate_flowchart(src_code: str) -> str:
     # check if src_code is empty
@@ -113,11 +131,9 @@ def generate_flowchart(src_code: str) -> str:
     # remove the thinking tags from the response
     cleaned = re.sub(r'<thinking>.*?</thinking>', '', content, flags=re.DOTALL)
 
-    # Remove raw quotes on both parentheses and non-parentheses
-    cleaned = cleaned.replace('"', "'")
-    cleaned = re.sub(r'\[([^\[\]]+)\]', r'["\g<1>"]', cleaned)
-    cleaned = re.sub(r'\{([^{}]+)\}', r'{"\g<1>"}', cleaned)
-    
+    # Quote-wrap each node label and escape any leaked code characters
+    cleaned = quote_labels(cleaned)
+
     return cleaned.strip() # and remove excess whitespace
 
 # ----- Custom Frontend ----- #