Adjust typing and error handling for incorrect diagrams

src/interface/web/app/components/mermaid/mermaid.tsx

@@ -1,5 +1,6 @@
 import React, { useEffect, useState, useRef } from "react";
 import mermaid from "mermaid";
+import { Info } from "@phosphor-icons/react";
 
 interface MermaidProps {
     chart: string;
@@ -18,8 +19,23 @@ const Mermaid: React.FC<MermaidProps> = ({ chart }) => {
         mermaid.parseError = (error) => {
             console.error("Mermaid errors:", error);
             // Extract error message from error object
-            const errorMessage = JSON.stringify(error);
+            // Parse error message safely
-            setMermaidError(errorMessage);
+            let errorMessage;
+            try {
+                errorMessage = typeof error === "string" ? JSON.parse(error) : error;
+            } catch (e) {
+                errorMessage = error?.toString() || "Unknown error";
+            }
+
+            console.log("Mermaid error message:", errorMessage);
+
+            if (errorMessage.str !== "element is null") {
+                setMermaidError(
+                    "Something went wrong while rendering the diagram. Please try again later or downvote the message if the issue persists.",
+                );
+            } else {
+                setMermaidError(null);
+            }
         };
 
         mermaid.contentLoaded();
@@ -37,30 +53,32 @@ const Mermaid: React.FC<MermaidProps> = ({ chart }) => {
                 .then(() => {
                     setMermaidError(null);
                 })
-                .catch((error) => {
+                .catch((error) => {});
-                    console.error("Mermaid error:", error);
-                    // Extract error message from error object
-                    const errorMessage = error?.str || error?.message || JSON.stringify(error);
-                    console.log("Mermaid error message:", errorMessage);
-                });
         }
     }, [chart]);
 
     return (
         <div>
-            <div
+            {mermaidError ? (
-                id={mermaidId}
+                <div className="flex items-center gap-2 bg-red-100 border border-red-500 rounded-md p-3 mt-3 text-red-900 text-sm">
-                ref={elementRef}
+                    <Info className="w-12 h-12" />
-                className="mermaid"
+                    <span>Error rendering diagram: {mermaidError}</span>
-                style={{
+                </div>
-                    width: "auto",
+            ) : (
-                    height: "auto",
+                <div
-                    boxSizing: "border-box",
+                    id={mermaidId}
-                    overflow: "auto",
+                    ref={elementRef}
-                }}
+                    className="mermaid"
-            >
+                    style={{
-                {chart}
+                        width: "auto",
-            </div>
+                        height: "auto",
+                        boxSizing: "border-box",
+                        overflow: "auto",
+                    }}
+                >
+                    {chart}
+                </div>
+            )}
         </div>
     );
 };

src/khoj/processor/conversation/prompts.py

@@ -428,27 +428,27 @@ pie title Pets adopted by volunteers
     "Rats" : 60
 
 flowchart TB
-    c1-->a2
+    subgraph "Group 1"
-    subgraph one
+        a1["Start Node"] --> a2["End Node"]
-    a1-->a2
     end
-    subgraph two
+    subgraph "Group 2"
-    b1-->b2["this is b2"]
+        b1["Process 1"] --> b2["Process 2"]
     end
-    subgraph three
+    subgraph "Group 3"
-    c1["this is c1"]-->c2["this is c2"]
+        c1["Input"] --> c2["Output"]
     end
-    one --> two
+    a["Group 1"] --> b["Group 2"]
-    three --> two
+    c["Group 3"] --> d["Group 2"]
-    two --> c2
 
 ----Process----
 Create your diagram with great composition and intuitiveness from the provided context and user prompt below.
 - You may use subgraphs to group elements together. Each subgraph must have a title.
-- **You must wrap ALL entity and node labels in double quotes**. For example, "Entity Name".
+- **You must wrap ALL entity and node labels in double quotes**, example: "My Node Label"
+- **All nodes MUST use the id["label"] format**. For example: node1["My Node Label"]
 - Custom style are not permitted. Default styles only.
 - JUST provide the diagram, no additional text or context. Say nothing else in your response except the diagram.
 - Keep diagrams simple - maximum 15 nodes
+- Every node inside a subgraph MUST use square bracket notation: id["label"]
 
 output: {query}
 

src/khoj/routers/helpers.py

@@ -1000,7 +1000,7 @@ async def generate_mermaidjs_diagram_from_description(
     user: KhojUser = None,
     agent: Agent = None,
     tracer: dict = {},
-) -> Dict[str, Any]:
+) -> str:
     personality_context = (
         prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
     )