From a8e627f1f83cd91dfd433415998ae0f192f72f93 Mon Sep 17 00:00:00 2001 From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com> Date: Wed, 3 Dec 2025 16:14:53 +0900 Subject: [PATCH] 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) --- app/api/chat/route.ts | 92 ++++++++++++++++++++++++++++------ app/api/chat/xml_guide.md | 30 +++++++++-- components/chat-panel.tsx | 26 +++++++--- lib/utils.ts | 103 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 225 insertions(+), 26 deletions(-) diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index a4aeaf9..ecb3962 100644 --- 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 + + + + + + + +\`\`\` + +CRITICAL RULES: +1. Always include the two root cells: and +2. ALL mxCell elements must be DIRECT children of - 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="" for grouped elements + +Shape (vertex) example: +\`\`\`xml + + + +\`\`\` + +Connector (edge) example: +\`\`\`xml + + + +\`\`\` + +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 tag (including the tag itself) in the XML string. - For example: - - - - - - - - - - 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 tags. + +VALIDATION RULES (XML will be rejected if violated): +1. All mxCell elements must be DIRECT children of - 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: < > & " +6. Always start with: + +Example with swimlanes and edges (note: all mxCells are siblings): + + + + + + + + + + + + + + + + + + + + +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") }) diff --git a/app/api/chat/xml_guide.md b/app/api/chat/xml_guide.md index 3490378..a15744d 100644 --- 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 ``. 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 ``. Edges are NOT nested inside swimlanes or steps.** ```xml - - - + + + + + + + + + + + + + + + + + + + + + + + + ``` ### Tables diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index 8016eef..3170cec 100644 --- 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 }>; diff --git a/lib/utils.ts b/lib/utils.ts index a74eecf..8fe61c3 100644 --- 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 < for <, > for >, & for &, " 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(); + 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 , 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)