Update app.py

app.py

@@ -453,39 +453,82 @@ def remove_text_from_html(report_html: str, text_to_remove: str) -> str:
 
 def fix_visual_after_section(report_html: str, section_title: str, extra_instructions: str) -> str:
     """
-    Given a report HTML (as a string), the target section name (from inside [[ ]]) and any extra instructions,
-    this function finds the first header (preferably an <h1>) that contains the section name (ignoring ones such as "Table of Contents"),
-    then finds the first <iframe> after that header, sends its srcdoc to the LLM for improvement, and reinjects the improved code.
+    Given a report HTML, a target section name (extracted from within [[...]]), and extra instructions,
+    this function finds the first header (h1–h4) that contains the section name (ignoring common TOC headers),
+    then finds the first <iframe> following that header.
+    It sends the current iframe 'srcdoc' (which should contain a mermaid diagram)
+    to the LLM with extra instructions and expects back a corrected mermaid code (starting with a mermaid diagram keyword).
+    It then rebuilds and reassigns the iframe's srcdoc with the improved code.
     """
     soup = BeautifulSoup(report_html, "html.parser")
     header = None
-    # Look for header tags (h1–h4) that contain the target, excluding common TOC headings
+    # Look for header tags (h1–h4) that contain the section title (ignoring "table of contents")
     for tag in soup.find_all(["h1", "h2", "h3", "h4"]):
-        text = tag.get_text(strip=True)
-        if section_title.lower() in text.lower() and "table of contents" not in text.lower():
+        header_text = tag.get_text(strip=True)
+        if section_title.lower() in header_text.lower() and "table of contents" not in header_text.lower():
             header = tag
             break
-    if header:
-        # Look for the first <iframe> after that header
-        iframe = header.find_next("iframe")
-        if iframe and iframe.has_attr("srcdoc"):
-            current_code = iframe["srcdoc"]
-            prompt = (
-                f"Improve the following visual code by simplifying its formatting (e.g., adjust parenthesis, remove extra styling) "
-                f"based on these extra instructions: {extra_instructions}\n\n"
-                f"Current code:\n{current_code}\n\n"
-                "Return only the improved code (no extra commentary)."
-            )
-            improved_code = llm_call(prompt=prompt, model="o3-mini", temperature=0, max_tokens_param=1500).strip()
-            # Check that the improvement is valid and not an error message.
-            if improved_code and not improved_code.lower().startswith("error: empty response"):
-                iframe["srcdoc"] = improved_code
-            else:
-                logging.error("fix_visual_after_section: LLM returned an empty or error response.")
-        else:
-            logging.error(f"fix_visual_after_section: No iframe found after section '{section_title}'.")
-    else:
+    if not header:
         logging.error(f"fix_visual_after_section: Section '{section_title}' not found.")
+        return report_html
+
+    # Find the first <iframe> after the header
+    iframe = header.find_next("iframe")
+    if not iframe or not iframe.has_attr("srcdoc"):
+        logging.error(f"fix_visual_after_section: No <iframe> found after section '{section_title}'.")
+        return report_html
+
+    current_code = iframe["srcdoc"]
+    logging.info(f"fix_visual_after_section: Current iframe srcdoc length: {len(current_code)}")
+    
+    # Build prompt: ask the LLM to simplify/adjust the mermaid code.
+    prompt = f"""Improve the following mermaid visualization code.
+Extra instructions: {extra_instructions}
+Current mermaid code:
+{current_code}
+
+Return only the improved mermaid code (the content that should go inside the <div class="mermaid">...</div>) with no extra wrapping HTML or commentary.
+Ensure that the output starts with one of the following keywords: mindmap, flowchart, sequence, gantt, or pie."""
+    
+    improved_code = llm_call(prompt=prompt, model="o3-mini", temperature=0, max_tokens_param=1500)
+    improved_code = improved_code.strip().strip("```").strip()
+    logging.info(f"fix_visual_after_section: Improved code received from LLM:\n{improved_code}")
+    
+    # Verify that the result starts with a valid mermaid diagram keyword.
+    valid_keywords = ["mindmap", "flowchart", "sequence", "gantt", "pie"]
+    if not any(improved_code.lower().startswith(keyword) for keyword in valid_keywords):
+        logging.error("fix_visual_after_section: Improved code does not start with a valid mermaid keyword.")
+        return report_html
+
+    # Rebuild the srcdoc using the improved mermaid code.
+    new_srcdoc = f"""<!DOCTYPE html>
+<html>
+    <head>
+        <script src="https://cdn.jsdelivr.net/npm/mermaid@10.4.0/dist/mermaid.min.js"></script>
+        <script>
+            mermaid.initialize({{ startOnLoad: true, securityLevel: "loose", theme: "default" }});
+        </script>
+        <style>
+            .mermaid {{
+                background-color: #f7f2f2;
+                color: black;
+                padding: 10px;
+                border-radius: 5px;
+                border: 1px solid #ccc;
+                max-width: 1200px;
+                margin: 0 auto;
+                min-height: 400px;
+            }}
+        </style>
+    </head>
+    <body>
+        <div class="mermaid">
+{improved_code}
+        </div>
+    </body>
+</html>"""
+    iframe["srcdoc"] = new_srcdoc
+    logging.info("fix_visual_after_section: Successfully updated the iframe srcdoc with improved mermaid code.")
     return str(soup)
 
 def snippet_in_tag(tag: Tag, snippet: str) -> bool: