diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index b83e2b6..dcebca5 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -75,7 +75,16 @@ ${lastMessage.content} tools: { // Client-side tool that will be executed on the client display_diagram: { - description: "Display a diagram on draw.io", + 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: + + + + + + + + `, parameters: z.object({ xml: z.string().describe("XML string to be displayed on draw.io") }) diff --git a/app/page.tsx b/app/page.tsx index bc81fba..a5f398d 100644 --- a/app/page.tsx +++ b/app/page.tsx @@ -2,7 +2,6 @@ import { DrawIoEmbed, DrawIoEmbedRef } from "react-drawio"; import { useRef, useState } from "react"; -import { Button } from "@/components/ui/button"; import { extractDiagramXML } from "./extract_xml"; import ChatPanel from "@/components/chat-panel"; @@ -28,6 +27,7 @@ export default function Home() { const loadDiagram = (chart: string) => { if (drawioRef.current) { + console.log("xml before load", chart); drawioRef.current.load({ xml: chart, }); @@ -65,6 +65,7 @@ export default function Home() {
loadDiagram(xml)} onFetchChart={() => { return new Promise((resolve) => { diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx index ee3a766..6cc38d8 100644 --- a/components/chat-message-display.tsx +++ b/components/chat-message-display.tsx @@ -6,9 +6,10 @@ import Image from "next/image"; import { ScrollArea } from "@/components/ui/scroll-area"; import ExamplePanel from "./chat-example-panel"; import { Message } from "ai"; -import { convertToLegalXml } from "@/lib/utils"; +import { convertToLegalXml, replaceNodes } from "@/lib/utils"; interface ChatMessageDisplayProps { + chartXML: string; messages: Message[]; error?: Error | null; setInput: (input: string) => void; @@ -17,6 +18,7 @@ interface ChatMessageDisplayProps { } export function ChatMessageDisplay({ + chartXML, messages, error, setInput, @@ -24,7 +26,7 @@ export function ChatMessageDisplay({ onDisplayChart, }: ChatMessageDisplayProps) { const messagesEndRef = useRef(null); - const stepCounterRef = useRef(0); + const previousXML = useRef(""); useEffect(() => { if (messagesEndRef.current) { messagesEndRef.current.scrollIntoView({ behavior: "smooth" }); @@ -42,19 +44,26 @@ export function ChatMessageDisplay({ const currentXml = toolInvocation.args?.xml || ""; // Increment the step counter - stepCounterRef.current += 1; // Determine whether to show details based on a simple threshold - if (stepCounterRef.current % 20 === 0) { - const convertedXml = - convertToLegalXml(currentXml); + const convertedXml = convertToLegalXml(currentXml); + if (convertedXml !== previousXML.current) { + previousXML.current = convertedXml; // if "/root" in convertedXml - if (convertedXml.includes("/root")) { - onDisplayChart(convertedXml); - console.log("converted xml", convertedXml); - } + const replacedXML = replaceNodes( + chartXML, + convertedXml + ); + console.log("currentXml", currentXml); + console.log("converted xml", convertedXml); + console.log("replaced xml", replacedXML); + onDisplayChart(replacedXML); + // if convertedXml changed } + // if "/root" in convertedXml + + // if convertedXml changed } return (
void; onFetchChart: () => Promise; diagramHistory?: { svg: string; xml: string }[]; @@ -21,9 +22,9 @@ interface ChatPanelProps { } export default function ChatPanel({ + chartXML, onDisplayChart, onFetchChart, - mergeXML, diagramHistory = [], onAddToHistory = () => {}, }: ChatPanelProps) { @@ -49,7 +50,8 @@ export default function ChatPanel({ async onToolCall({ toolCall }) { if (toolCall.toolName === "display_diagram") { const { xml } = toolCall.args as { xml: string }; - onDisplayChart(xml); + // do nothing because we will handle this streamingly in the ChatMessageDisplay component + // onDisplayChart(xml); return "Successfully displayed the flowchart."; } }, @@ -105,6 +107,7 @@ export default function ChatPanel({ tag does not have an mxGeometry child (e.g. ), + * it removes that tag from the output. * @param xmlString The potentially incomplete XML string - * @returns A legal XML string with properly closed tags + * @returns A legal XML string with properly closed tags and removed incomplete mxCell elements. */ export function convertToLegalXml(xmlString: string): string { - const stack: string[] = []; - let result = ''; - let tagStart = -1; + // This regex will match either self-closing or a block element + // ... . Unfinished ones are left out because they don't match. + const regex = /]*(?:\/>|>([\s\S]*?)<\/mxCell>)/g; + let match: RegExpExecArray | null; + let result = "\n"; - for (let i = 0; i < xmlString.length; i++) { - const char = xmlString[i]; - result += char; - - if (char === '<' && tagStart === -1) { - // Start of a new tag - tagStart = i; - } else if (char === '>' && tagStart !== -1) { - // End of a tag - const tagContent = xmlString.substring(tagStart + 1, i); - - if (tagContent.startsWith('/')) { - // Closing tag - const tagName = tagContent.substring(1).trim().split(/\s+/)[0]; - if (stack.length && stack[stack.length - 1] === tagName) { - stack.pop(); - } - } else if (!tagContent.endsWith('/') && !tagContent.startsWith('?') && !tagContent.startsWith('!')) { - // Opening tag (not self-closing, processing instruction, or comment) - const tagName = tagContent.trim().split(/\s+/)[0]; - stack.push(tagName); - } - - tagStart = -1; - } - } - - // If we have an incomplete tag at the end, don't include it in the result - if (tagStart !== -1) { - result = result.substring(0, tagStart); - } - - // Close all remaining open tags - for (let j = stack.length - 1; j >= 0; j--) { - result += ``; + while ((match = regex.exec(xmlString)) !== null) { + // match[0] contains the entire matched mxCell block + // Indent each line of the matched block for readability. + const formatted = match[0].split('\n').map(line => " " + line.trim()).join('\n'); + result += formatted + "\n"; } + result += ""; return result; } +/** + * Replace nodes in a Draw.io XML diagram + * @param currentXML - The original Draw.io XML string + * @param nodes - The XML string containing new nodes to replace in the diagram + * @returns The updated XML string with replaced nodes + */ +export function replaceNodes(currentXML: string, nodes: string): string { + // Check for valid inputs + if (!currentXML || !nodes) { + throw new Error("Both currentXML and nodes must be provided"); + } + + try { + // Parse the XML strings to create DOM objects + const parser = new DOMParser(); + const currentDoc = parser.parseFromString(currentXML, "text/xml"); + + // Handle nodes input - if it doesn't contain , wrap it + let nodesString = nodes; + if (!nodes.includes("")) { + nodesString = `${nodes}`; + } + + const nodesDoc = parser.parseFromString(nodesString, "text/xml"); + + // Find the root element in the current document + let currentRoot = currentDoc.querySelector("mxGraphModel > root"); + if (!currentRoot) { + // If no root element is found, create the proper structure + const mxGraphModel = currentDoc.querySelector("mxGraphModel") || + currentDoc.createElement("mxGraphModel"); + + if (!currentDoc.contains(mxGraphModel)) { + currentDoc.appendChild(mxGraphModel); + } + + currentRoot = currentDoc.createElement("root"); + mxGraphModel.appendChild(currentRoot); + } + + // Find the root element in the nodes document + const nodesRoot = nodesDoc.querySelector("root"); + if (!nodesRoot) { + throw new Error("Invalid nodes: Could not find or create element"); + } + + // Clear all existing child elements from the current root + while (currentRoot.firstChild) { + currentRoot.removeChild(currentRoot.firstChild); + } + + // Ensure the base cells exist + const hasCell0 = Array.from(nodesRoot.childNodes).some( + node => node.nodeName === "mxCell" && + (node as Element).getAttribute("id") === "0" + ); + + const hasCell1 = Array.from(nodesRoot.childNodes).some( + node => node.nodeName === "mxCell" && + (node as Element).getAttribute("id") === "1" + ); + + // Copy all child nodes from the nodes root to the current root + Array.from(nodesRoot.childNodes).forEach(node => { + const importedNode = currentDoc.importNode(node, true); + currentRoot.appendChild(importedNode); + }); + + // Add default cells if they don't exist + if (!hasCell0) { + const cell0 = currentDoc.createElement("mxCell"); + cell0.setAttribute("id", "0"); + currentRoot.insertBefore(cell0, currentRoot.firstChild); + } + + if (!hasCell1) { + const cell1 = currentDoc.createElement("mxCell"); + cell1.setAttribute("id", "1"); + cell1.setAttribute("parent", "0"); + + // Insert after cell0 if possible + const cell0 = currentRoot.querySelector('mxCell[id="0"]'); + if (cell0 && cell0.nextSibling) { + currentRoot.insertBefore(cell1, cell0.nextSibling); + } else { + currentRoot.appendChild(cell1); + } + } + + // Convert the modified DOM back to a string + const serializer = new XMLSerializer(); + return serializer.serializeToString(currentDoc); + } catch (error) { + throw new Error(`Error replacing nodes: ${error}`); + } +} \ No newline at end of file