feat: add XML structure guide to system prompt for smaller models (#51)

- Add essential draw.io XML structure rules to system prompt - Include critical rules about mxCell nesting (all must be direct children of root) - Add shape/vertex and connector/edge examples with proper structure - Improve tool description for display_diagram with validation rules - Update xml_guide.md with better swimlane examples showing flat structure - Add client-side XML validation to catch nested mxCell errors early Helps address issues #40 (local Ollama models not working) and #39 (mxCell nesting errors)
2026-01-02 14:22:28 +08:00 · 2025-12-03 16:14:53 +09:00
parent c458947553
commit a8e627f1f8
4 changed files with 225 additions and 26 deletions
--- a/app/api/chat/route.ts
+++ b/app/api/chat/route.ts
@@ -52,9 +52,8 @@ export async function POST(req: Request) {

    const systemMessage = `
 You are an expert diagram creation assistant specializing in draw.io XML generation.
-Your primary function is crafting clear, well-organized visual diagrams through precise XML specifications.
+Your primary function is chat with user and crafting clear, well-organized visual diagrams through precise XML specifications.
 You can see the image that user uploaded.
-Note that when you need to generate diagram about aws architecture, use **AWS 2025 icons**.

 You utilize the following tools:
 ---Tool1---
@@ -95,6 +94,9 @@ Layout constraints:
 - Avoid spreading elements too far apart horizontally - users should see the complete diagram without a page break line

 Note that:
+- Use proper tool calls to generate or edit diagrams;
+  - never return raw XML in text responses,
+  - never use display_diagram to generate messages that you want to send user directly. e.g. to generate a "hello" text box when you want to greet user.
 - Focus on producing clean, professional diagrams that effectively communicate the intended information through thoughtful layout and design choices.
 - When artistic drawings are requested, creatively compose them using standard diagram shapes and connectors while maintaining visual clarity.
 - Return XML only via tool calls, never in text responses.
@@ -110,6 +112,44 @@ When using edit_diagram tool:
  * You may retry edit_diagram up to 3 times with adjusted search patterns
  * After 3 failed attempts, you MUST fall back to using display_diagram to regenerate the entire diagram
  * The error message will indicate how many retries remain
+
+## Draw.io XML Structure Reference
+
+Basic structure:
+\`\`\`xml
+<mxGraphModel>
+  <root>
+    <mxCell id="0"/>
+    <mxCell id="1" parent="0"/>
+    <!-- All other cells go here as siblings -->
+  </root>
+</mxGraphModel>
+\`\`\`
+
+CRITICAL RULES:
+1. Always include the two root cells: <mxCell id="0"/> and <mxCell id="1" parent="0"/>
+2. ALL mxCell elements must be DIRECT children of <root> - NEVER nest mxCell inside another mxCell
+3. Use unique sequential IDs for all cells (start from "2" for user content)
+4. Set parent="1" for top-level shapes, or parent="<container-id>" for grouped elements
+
+Shape (vertex) example:
+\`\`\`xml
+<mxCell id="2" value="Label" style="rounded=1;whiteSpace=wrap;html=1;" vertex="1" parent="1">
+  <mxGeometry x="100" y="100" width="120" height="60" as="geometry"/>
+</mxCell>
+\`\`\`
+
+Connector (edge) example:
+\`\`\`xml
+<mxCell id="3" style="endArrow=classic;html=1;" edge="1" parent="1" source="2" target="4">
+  <mxGeometry relative="1" as="geometry"/>
+</mxCell>
+\`\`\`
+
+Common styles:
+- Shapes: rounded=1 (rounded corners), fillColor=#hex, strokeColor=#hex
+- Edges: endArrow=classic/block/open/none, startArrow=none/classic, curved=1, edgeStyle=orthogonalEdgeStyle
+- Text: fontSize=14, fontStyle=1 (bold), align=center/left/right
 `;

    const lastMessage = messages[messages.length - 1];
@@ -228,19 +268,41 @@ ${lastMessageText}
      tools: {
        // Client-side tool that will be executed on the client
        display_diagram: {
-          description: `Display a diagram on draw.io. You only need to pass the nodes inside the <root> tag (including the <root> tag itself) in the XML string.
-          For example:
-          <root>
-            <mxCell id="0"/>
-            <mxCell id="1" parent="0"/>
-            <mxGeometry x="20" y="20" width="100" height="100" as="geometry"/>
-            <mxCell id="2" value="Hello, World!" style="shape=rectangle" parent="1">
-              <mxGeometry x="20" y="20" width="100" height="100" as="geometry"/>
-            </mxCell>
-          </root>
-          - Note that when you need to generate diagram about aws architecture, use **AWS 2025 icons**.
-          - If you are asked to generate animated connectors, make sure to include "flowAnimation=1" in the style of the connector elements.
-          `,
+          description: `Display a diagram on draw.io. Pass the XML content inside <root> tags.
+
+VALIDATION RULES (XML will be rejected if violated):
+1. All mxCell elements must be DIRECT children of <root> - never nested
+2. Every mxCell needs a unique id
+3. Every mxCell (except id="0") needs a valid parent attribute
+4. Edge source/target must reference existing cell IDs
+5. Escape special chars in values: &lt; &gt; &amp; &quot;
+6. Always start with: <mxCell id="0"/><mxCell id="1" parent="0"/>
+
+Example with swimlanes and edges (note: all mxCells are siblings):
+<root>
+  <mxCell id="0"/>
+  <mxCell id="1" parent="0"/>
+  <mxCell id="lane1" value="Frontend" style="swimlane;" vertex="1" parent="1">
+    <mxGeometry x="40" y="40" width="200" height="200" as="geometry"/>
+  </mxCell>
+  <mxCell id="step1" value="Step 1" style="rounded=1;" vertex="1" parent="lane1">
+    <mxGeometry x="20" y="60" width="160" height="40" as="geometry"/>
+  </mxCell>
+  <mxCell id="lane2" value="Backend" style="swimlane;" vertex="1" parent="1">
+    <mxGeometry x="280" y="40" width="200" height="200" as="geometry"/>
+  </mxCell>
+  <mxCell id="step2" value="Step 2" style="rounded=1;" vertex="1" parent="lane2">
+    <mxGeometry x="20" y="60" width="160" height="40" as="geometry"/>
+  </mxCell>
+  <mxCell id="edge1" style="edgeStyle=orthogonalEdgeStyle;endArrow=classic;" edge="1" parent="1" source="step1" target="step2">
+    <mxGeometry relative="1" as="geometry"/>
+  </mxCell>
+</root>
+
+Notes:
+- For AWS diagrams, use **AWS 2025 icons**.
+- For animated connectors, add "flowAnimation=1" to edge style.
+`,
          inputSchema: z.object({
            xml: z.string().describe("XML string to be displayed on draw.io")
          })
--- a/app/api/chat/xml_guide.md
+++ b/app/api/chat/xml_guide.md
@@ -211,6 +211,7 @@ Draw.io files contain two special cells that are always present:
 4.  Define parent relationships correctly
 5.  Use `mxGeometry` elements to position shapes
 6.  For connectors, specify `source` and `target` attributes
+7.  **CRITICAL: All mxCell elements must be DIRECT children of `<root>`. NEVER nest mxCell inside another mxCell.**

 ## Common Patterns

@@ -234,12 +235,33 @@ To group elements, create a parent cell and set other cells' `parent` attribute

 ### Swimlanes

-Swimlanes use the `swimlane` shape style:
+Swimlanes use the `swimlane` shape style. **IMPORTANT: All mxCell elements (swimlanes, steps, and edges) must be siblings under `<root>`. Edges are NOT nested inside swimlanes or steps.**

 ```xml
-<mxCell id="20" value="Swimlane 1" style="swimlane;fontStyle=0;childLayout=stackLayout;horizontal=1;startSize=30;horizontalStack=0;resizeParent=1;resizeParentMax=0;resizeLast=0;collapsible=1;marginBottom=0;whiteSpace=wrap;html=1;" vertex="1" parent="1">
-  <mxGeometry x="200" y="200" width="140" height="120" as="geometry" />
-</mxCell>
+<root>
+  <mxCell id="0"/>
+  <mxCell id="1" parent="0"/>
+  <!-- Swimlane 1 -->
+  <mxCell id="lane1" value="Frontend" style="swimlane;startSize=30;" vertex="1" parent="1">
+    <mxGeometry x="40" y="40" width="200" height="300" as="geometry"/>
+  </mxCell>
+  <!-- Swimlane 2 -->
+  <mxCell id="lane2" value="Backend" style="swimlane;startSize=30;" vertex="1" parent="1">
+    <mxGeometry x="280" y="40" width="200" height="300" as="geometry"/>
+  </mxCell>
+  <!-- Step inside lane1 (parent="lane1") -->
+  <mxCell id="step1" value="Send Request" style="rounded=1;" vertex="1" parent="lane1">
+    <mxGeometry x="20" y="60" width="160" height="40" as="geometry"/>
+  </mxCell>
+  <!-- Step inside lane2 (parent="lane2") -->
+  <mxCell id="step2" value="Process" style="rounded=1;" vertex="1" parent="lane2">
+    <mxGeometry x="20" y="60" width="160" height="40" as="geometry"/>
+  </mxCell>
+  <!-- Edge connecting step1 to step2 (sibling element, NOT nested inside steps) -->
+  <mxCell id="edge1" style="edgeStyle=orthogonalEdgeStyle;endArrow=classic;" edge="1" parent="1" source="step1" target="step2">
+    <mxGeometry relative="1" as="geometry"/>
+  </mxCell>
+</root>
 ```

 ### Tables
--- a/components/chat-panel.tsx
+++ b/components/chat-panel.tsx
@@ -18,7 +18,7 @@ import { DefaultChatTransport } from "ai";
 import { ChatInput } from "@/components/chat-input";
 import { ChatMessageDisplay } from "./chat-message-display";
 import { useDiagram } from "@/contexts/diagram-context";
-import { replaceNodes, formatXML } from "@/lib/utils";
+import { replaceNodes, formatXML, validateMxCellStructure } from "@/lib/utils";
 import { ButtonWithTooltip } from "@/components/button-with-tooltip";

 interface ChatPanelProps {
@@ -73,12 +73,24 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr
            }),
            async onToolCall({ toolCall }) {
                if (toolCall.toolName === "display_diagram") {
-                    // Diagram is handled streamingly in the ChatMessageDisplay component
-                    addToolResult({
-                        tool: "display_diagram",
-                        toolCallId: toolCall.toolCallId,
-                        output: "Successfully displayed the diagram.",
-                    });
+                    const { xml } = toolCall.input as { xml: string };
+
+                    // Validate XML structure before confirming success
+                    const validationError = validateMxCellStructure(xml);
+
+                    if (validationError) {
+                        addToolResult({
+                            tool: "display_diagram",
+                            toolCallId: toolCall.toolCallId,
+                            output: validationError,
+                        });
+                    } else {
+                        addToolResult({
+                            tool: "display_diagram",
+                            toolCallId: toolCall.toolCallId,
+                            output: "Successfully displayed the diagram.",
+                        });
+                    }
                } else if (toolCall.toolName === "edit_diagram") {
                    const { edits } = toolCall.input as {
                        edits: Array<{ search: string; replace: string }>;
--- a/lib/utils.ts
+++ b/lib/utils.ts
@@ -306,6 +306,109 @@ export function replaceXMLParts(
  return result;
 }

+/**
+ * Validates draw.io XML structure for common issues
+ * @param xml - The XML string to validate
+ * @returns null if valid, error message string if invalid
+ */
+export function validateMxCellStructure(xml: string): string | null {
+  const parser = new DOMParser();
+  const doc = parser.parseFromString(xml, "text/xml");
+
+  // Check for XML parsing errors (includes unescaped special characters)
+  const parseError = doc.querySelector('parsererror');
+  if (parseError) {
+    return `Invalid XML: The XML contains syntax errors (likely unescaped special characters like <, >, & in attribute values). Please escape special characters: use &lt; for <, &gt; for >, &amp; for &, &quot; for ". Regenerate the diagram with properly escaped values.`;
+  }
+
+  // Get all mxCell elements once for all validations
+  const allCells = doc.querySelectorAll('mxCell');
+
+  // Single pass: collect IDs, check for duplicates, nesting, orphans, and invalid parents
+  const cellIds = new Set<string>();
+  const duplicateIds: string[] = [];
+  const nestedCells: string[] = [];
+  const orphanCells: string[] = [];
+  const invalidParents: { id: string; parent: string }[] = [];
+  const edgesToValidate: { id: string; source: string | null; target: string | null }[] = [];
+
+  allCells.forEach(cell => {
+    const id = cell.getAttribute('id');
+    const parent = cell.getAttribute('parent');
+    const isEdge = cell.getAttribute('edge') === '1';
+
+    // Check for duplicate IDs
+    if (id) {
+      if (cellIds.has(id)) {
+        duplicateIds.push(id);
+      } else {
+        cellIds.add(id);
+      }
+    }
+
+    // Check for nested mxCell (parent element is also mxCell)
+    if (cell.parentElement?.tagName === 'mxCell') {
+      nestedCells.push(id || 'unknown');
+    }
+
+    // Check parent attribute (skip root cell id="0")
+    if (id !== '0') {
+      if (!parent) {
+        if (id) orphanCells.push(id);
+      } else {
+        // Store for later validation (after all IDs collected)
+        invalidParents.push({ id: id || 'unknown', parent });
+      }
+    }
+
+    // Collect edges for connection validation
+    if (isEdge) {
+      edgesToValidate.push({
+        id: id || 'unknown',
+        source: cell.getAttribute('source'),
+        target: cell.getAttribute('target')
+      });
+    }
+  });
+
+  // Return errors in priority order
+  if (nestedCells.length > 0) {
+    return `Invalid XML: Found nested mxCell elements (IDs: ${nestedCells.slice(0, 3).join(', ')}). All mxCell elements must be direct children of <root>, never nested inside other mxCell elements. Please regenerate the diagram with correct structure.`;
+  }
+
+  if (duplicateIds.length > 0) {
+    return `Invalid XML: Found duplicate cell IDs (${duplicateIds.slice(0, 3).join(', ')}). Each mxCell must have a unique ID. Please regenerate the diagram with unique IDs for all elements.`;
+  }
+
+  if (orphanCells.length > 0) {
+    return `Invalid XML: Found cells without parent attribute (IDs: ${orphanCells.slice(0, 3).join(', ')}). All mxCell elements (except id="0") must have a parent attribute. Please regenerate the diagram with proper parent references.`;
+  }
+
+  // Validate parent references (now that all IDs are collected)
+  const badParents = invalidParents.filter(p => !cellIds.has(p.parent));
+  if (badParents.length > 0) {
+    const details = badParents.slice(0, 3).map(p => `${p.id} (parent: ${p.parent})`).join(', ');
+    return `Invalid XML: Found cells with invalid parent references (${details}). Parent IDs must reference existing cells. Please regenerate the diagram with valid parent references.`;
+  }
+
+  // Validate edge connections
+  const invalidConnections: string[] = [];
+  edgesToValidate.forEach(edge => {
+    if (edge.source && !cellIds.has(edge.source)) {
+      invalidConnections.push(`${edge.id} (source: ${edge.source})`);
+    }
+    if (edge.target && !cellIds.has(edge.target)) {
+      invalidConnections.push(`${edge.id} (target: ${edge.target})`);
+    }
+  });
+
+  if (invalidConnections.length > 0) {
+    return `Invalid XML: Found edges with invalid source/target references (${invalidConnections.slice(0, 3).join(', ')}). Edge source and target must reference existing cell IDs. Please regenerate the diagram with valid edge connections.`;
+  }
+
+  return null;
+}
+
 export function extractDiagramXML(xml_svg_string: string): string {
  try {
    // 1. Parse the SVG string (using built-in DOMParser in a browser-like environment)