From 0851b32b67b9fa9646ee02a6ccea1e2bb6e48071 Mon Sep 17 00:00:00 2001 From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com> Date: Sun, 14 Dec 2025 14:04:44 +0900 Subject: [PATCH] refactor: simplify LLM XML format to output bare mxCells only (#254) * refactor: simplify LLM XML format to output bare mxCells only - Update wrapWithMxFile() to always add root cells (id=0, id=1) automatically - LLM now generates only mxCell elements starting from id=2 (no wrapper tags) - Update system prompts and tool descriptions with new format instructions - Update cached responses to remove root cells and wrapper tags - Update truncation detection to check for complete mxCell endings - Update documentation in xml_guide.md * fix: address PR review issues for XML format refactor - Fix critical bug: inconsistent truncation check using old pattern - Fix stale error message referencing tag - Add isMxCellXmlComplete() helper for consistent truncation detection - Improve regex patterns to handle any attribute order in root cells - Update wrapWithMxFile JSDoc to document root cell removal behavior * fix: handle non-self-closing root cells in wrapWithMxFile regex --- app/api/chat/route.ts | 54 ++++++++++----------- app/api/chat/xml_guide.md | 21 ++++---- components/chat-message-display.tsx | 12 +++-- components/chat-panel.tsx | 29 ++++++----- lib/cached-responses.ts | 49 ++++--------------- lib/system-prompts.ts | 75 ++++++++++++++--------------- lib/utils.ts | 47 +++++++++++++++--- 7 files changed, 143 insertions(+), 144 deletions(-) diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index 19e5560..05b9de1 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -367,36 +367,32 @@ ${userInputText} tools: { // Client-side tool that will be executed on the client display_diagram: { - description: `Display a diagram on draw.io. Pass the XML content inside tags. + description: `Display a diagram on draw.io. Pass ONLY the mxCell elements - wrapper tags and root cells are added automatically. 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: +1. Generate ONLY mxCell elements - NO wrapper tags (, , ) +2. Do NOT include root cells (id="0" or id="1") - they are added automatically +3. All mxCell elements must be siblings - never nested +4. Every mxCell needs a unique id (start from "2") +5. Every mxCell needs a valid parent attribute (use "1" for top-level) +6. Escape special chars in values: < > & " -Example with swimlanes and edges (note: all mxCells are siblings): - - - - - - - - - - - - - - - - - - - +Example (generate ONLY this - no wrapper tags): + + + + + + + + + + + + + + + Notes: - For AWS diagrams, use **AWS 2025 icons**. @@ -444,9 +440,9 @@ IMPORTANT: Keep edits concise: WHEN TO USE: Only call this tool after display_diagram was truncated (you'll see an error message about truncation). CRITICAL INSTRUCTIONS: -1. Do NOT start with , , or - they already exist in the partial +1. Do NOT include any wrapper tags - just continue the mxCell elements 2. Continue from EXACTLY where your previous output stopped -3. End with the closing tag to complete the diagram +3. Complete the remaining mxCell elements 4. If still truncated, call append_diagram again with the next fragment Example: If previous output ended with '...' and complete the remaining elements.`, diff --git a/app/api/chat/xml_guide.md b/app/api/chat/xml_guide.md index a15744d..e76edf6 100644 --- a/app/api/chat/xml_guide.md +++ b/app/api/chat/xml_guide.md @@ -81,16 +81,15 @@ Contains the actual diagram data. ## Root Cell Container: `` -Contains all the cells in the diagram. +Contains all the cells in the diagram. **Note:** When generating diagrams, you only need to provide the mxCell elements - the root container and root cells (id="0", id="1") are added automatically. -**Example:** +**Internal structure (auto-generated):** ```xml - - - - + + + ``` @@ -203,15 +202,15 @@ Draw.io files contain two special cells that are always present: 1. **Root Cell** (id = "0"): The parent of all cells 2. **Default Parent Cell** (id = "1", parent = "0"): The default layer and parent for most cells -## Tips for Manually Creating Draw.io XML +## Tips for Creating Draw.io XML -1. Start with the basic structure (`mxfile`, `diagram`, `mxGraphModel`, `root`) -2. Always include the two special cells (id = "0" and id = "1") +1. **Generate ONLY mxCell elements** - wrapper tags and root cells (id="0", id="1") are added automatically +2. Start IDs from "2" (id="0" and id="1" are reserved for root cells) 3. Assign unique and sequential IDs to all cells -4. Define parent relationships correctly +4. Define parent relationships correctly (use parent="1" for top-level shapes) 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.** +7. **CRITICAL: All mxCell elements must be siblings. NEVER nest mxCell inside another mxCell.** ## Common Patterns diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx index 81c80e7..b3fc76c 100644 --- a/components/chat-message-display.tsx +++ b/components/chat-message-display.tsx @@ -29,7 +29,12 @@ import { ReasoningTrigger, } from "@/components/ai-elements/reasoning" import { ScrollArea } from "@/components/ui/scroll-area" -import { convertToLegalXml, replaceNodes, validateAndFixXml } from "@/lib/utils" +import { + convertToLegalXml, + isMxCellXmlComplete, + replaceNodes, + validateAndFixXml, +} from "@/lib/utils" import ExamplePanel from "./chat-example-panel" import { CodeBlock } from "./code-block" @@ -457,8 +462,7 @@ export function ChatMessageDisplay({ const isTruncated = (toolName === "display_diagram" || toolName === "append_diagram") && - (!input?.xml || - !input.xml.includes("")) + !isMxCellXmlComplete(input?.xml) return isTruncated ? ( Truncated @@ -507,7 +511,7 @@ export function ChatMessageDisplay({ const isTruncated = (toolName === "display_diagram" || toolName === "append_diagram") && - (!input?.xml || !input.xml.includes("")) + !isMxCellXmlComplete(input?.xml) return (
indicates incomplete output) - const isTruncated = - !xml.includes("") && !xml.trim().endsWith("/>") + // Check if XML is truncated (incomplete mxCell indicates truncated output) + const isTruncated = !isMxCellXmlComplete(xml) if (isTruncated) { // Store the partial XML for continuation via append_diagram @@ -244,9 +243,9 @@ ${partialEnding} \`\`\` NEXT STEP: Call append_diagram with the continuation XML. -- Do NOT start with , , or (they already exist) +- Do NOT include wrapper tags or root cells (id="0", id="1") - Start from EXACTLY where you stopped -- End with the closing tag to complete the diagram`, +- Complete all remaining mxCell elements`, }) return } @@ -376,17 +375,21 @@ Please retry with an adjusted search pattern or use display_diagram if retries a const { xml } = toolCall.input as { xml: string } // Detect if LLM incorrectly started fresh instead of continuing + // LLM should only output bare mxCells now, so wrapper tags indicate error + const trimmed = xml.trim() const isFreshStart = - xml.trim().startsWith(", , or . + errorText: `ERROR: You started fresh with wrapper tags. Do NOT include wrapper tags or root cells (id="0", id="1"). Continue from EXACTLY where the partial ended: \`\`\` @@ -401,8 +404,8 @@ Start your continuation with the NEXT character after where it stopped.`, // Append to accumulated XML partialXmlRef.current += xml - // Check if XML is now complete - const isComplete = partialXmlRef.current.includes("") + // Check if XML is now complete (last mxCell is complete) + const isComplete = isMxCellXmlComplete(partialXmlRef.current) if (isComplete) { // Wrap and display the complete diagram @@ -439,7 +442,7 @@ Please use display_diagram with corrected XML.`, tool: "append_diagram", toolCallId: toolCall.toolCallId, state: "output-error", - errorText: `XML still incomplete (missing ). Call append_diagram again to continue. + errorText: `XML still incomplete (mxCell not closed). Call append_diagram again to continue. Current ending: \`\`\` diff --git a/lib/cached-responses.ts b/lib/cached-responses.ts index 2b9a740..8b8375f 100644 --- a/lib/cached-responses.ts +++ b/lib/cached-responses.ts @@ -9,12 +9,7 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [ promptText: "Give me a **animated connector** diagram of transformer's architecture", hasImage: false, - xml: ` - - - - - + xml: ` @@ -254,18 +249,12 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [ - -`, + `, }, { promptText: "Replicate this in aws style", hasImage: true, - xml: ` - - - - - + xml: ` @@ -324,18 +313,12 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [ - -`, + `, }, { promptText: "Replicate this flowchart.", hasImage: true, - xml: ` - - - - - + xml: ` @@ -391,16 +374,12 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [ - -`, + `, }, { promptText: "Summarize this paper as a diagram", hasImage: true, - xml: ` - - - @@ -751,18 +730,12 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [ value="Foundational technique for modern LLM reasoning - inspired many follow-up works including Self-Consistency, Tree-of-Thought, etc." vertex="1"> - - `, + `, }, { promptText: "Draw a cat for me", hasImage: false, - xml: ` - - - - - + xml: ` @@ -902,9 +875,7 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [ - - -`, + `, }, ] diff --git a/lib/system-prompts.ts b/lib/system-prompts.ts index d878acb..08875e8 100644 --- a/lib/system-prompts.ts +++ b/lib/system-prompts.ts @@ -104,22 +104,21 @@ When using edit_diagram tool: ## Draw.io XML Structure Reference -Basic structure: +**IMPORTANT:** You only generate the mxCell elements. The wrapper structure and root cells (id="0", id="1") are added automatically. + +Example - generate ONLY this: \`\`\`xml - - - - - - + + + \`\`\` -Note: All other mxCell elements go as siblings after id="1". 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 +1. Generate ONLY mxCell elements - NO wrapper tags (, , ) +2. Do NOT include root cells (id="0" or id="1") - they are added automatically +3. ALL mxCell elements must be siblings - NEVER nest mxCell inside another mxCell +4. Use unique sequential IDs starting from "2" +5. Set parent="1" for top-level shapes, or parent="" for grouped elements Shape (vertex) example: \`\`\`xml @@ -151,34 +150,30 @@ const EXTENDED_ADDITIONS = ` ### display_diagram Details **VALIDATION RULES** (XML will be rejected if violated): -1. All mxCell elements must be DIRECT children of - never nested inside other mxCell elements -2. Every mxCell needs a unique id attribute -3. Every mxCell (except id="0") needs a valid parent attribute referencing an existing cell -4. Edge source/target attributes must reference existing cell IDs -5. Escape special characters in values: < for <, > for >, & for &, " for " -6. Always start with the two root cells: +1. Generate ONLY mxCell elements - wrapper tags and root cells are added automatically +2. All mxCell elements must be siblings - never nested inside other mxCell elements +3. Every mxCell needs a unique id attribute (start from "2") +4. Every mxCell needs a valid parent attribute (use "1" for top-level, or container-id for grouped) +5. Edge source/target attributes must reference existing cell IDs +6. Escape special characters in values: < for <, > for >, & for &, " for " -**Example with swimlanes and edges** (note: all mxCells are siblings under ): +**Example with swimlanes and edges** (generate ONLY this - no wrapper tags): \`\`\`xml - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + \`\`\` ### append_diagram Details @@ -186,9 +181,9 @@ const EXTENDED_ADDITIONS = ` **WHEN TO USE:** Only call this tool when display_diagram output was truncated (you'll see an error message about truncation). **CRITICAL RULES:** -1. Do NOT start with , , or - they already exist in the partial +1. Do NOT include any wrapper tags - just continue the mxCell elements 2. Continue from EXACTLY where your previous output stopped -3. End with the closing tag to complete the diagram +3. Complete the remaining mxCell elements 4. If still truncated, call append_diagram again with the next fragment **Example:** If previous output ended with \`...\` and complete the remaining elements. diff --git a/lib/utils.ts b/lib/utils.ts index c878006..65c3276 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -29,6 +29,22 @@ const STRUCTURAL_ATTRS = [ /** Valid XML entity names */ const VALID_ENTITIES = new Set(["lt", "gt", "amp", "quot", "apos"]) +// ============================================================================ +// mxCell XML Helpers +// ============================================================================ + +/** + * Check if mxCell XML output is complete (not truncated). + * Complete XML ends with a self-closing tag (/>) or closing mxCell tag. + * @param xml - The XML string to check (can be undefined/null) + * @returns true if XML appears complete, false if truncated or empty + */ +export function isMxCellXmlComplete(xml: string | undefined | null): boolean { + const trimmed = xml?.trim() || "" + if (!trimmed) return false + return trimmed.endsWith("/>") || trimmed.endsWith("") +} + // ============================================================================ // XML Parsing Helpers // ============================================================================ @@ -197,13 +213,17 @@ export function convertToLegalXml(xmlString: string): string { /** * Wrap XML content with the full mxfile structure required by draw.io. - * Handles cases where XML is just , , or already has . - * @param xml - The XML string (may be partial or complete) - * @returns Full mxfile-wrapped XML string + * Always adds root cells (id="0" and id="1") automatically. + * If input already contains root cells, they are removed to avoid duplication. + * LLM should only generate mxCell elements starting from id="2". + * @param xml - The XML string (bare mxCells, , , or full ) + * @returns Full mxfile-wrapped XML string with root cells included */ export function wrapWithMxFile(xml: string): string { - if (!xml) { - return `` + const ROOT_CELLS = '' + + if (!xml || !xml.trim()) { + return `${ROOT_CELLS}` } // Already has full structure @@ -216,9 +236,20 @@ export function wrapWithMxFile(xml: string): string { return `${xml}` } - // Just content - extract inner content and wrap fully - const rootContent = xml.replace(/<\/?root>/g, "").trim() - return `${rootContent}` + // Has wrapper - extract inner content + let content = xml + if (xml.includes("")) { + content = xml.replace(/<\/?root>/g, "").trim() + } + + // Remove any existing root cells from content (LLM shouldn't include them, but handle it gracefully) + // Use flexible patterns that match both self-closing (/>) and non-self-closing (>) formats + content = content + .replace(/]*\bid=["']0["'][^>]*(?:\/>|><\/mxCell>)/g, "") + .replace(/]*\bid=["']1["'][^>]*(?:\/>|><\/mxCell>)/g, "") + .trim() + + return `${ROOT_CELLS}${content}` } /**