From 43e5993f477a73c67fe0f34d6b42f6adaf1ffb36 Mon Sep 17 00:00:00 2001 From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com> Date: Wed, 10 Dec 2025 18:04:37 +0900 Subject: [PATCH] fix: improve LLM diagram context awareness and image preview (#202) - Add replaceHistoricalToolInputs to replace XML in tool calls with placeholders - Send both previousXml and current xml so LLM can understand user's manual edits - Update system message to mark current XML as authoritative source of truth - Fix React StrictMode issue with blob URL cleanup in FilePreviewList - Add unoptimized prop to Image components for blob URLs --- app/api/chat/route.ts | 40 ++++++++++++++++++++++++++++---- components/chat-panel.tsx | 30 ++++++++++++++++++++++++ components/file-preview-list.tsx | 4 ++++ 3 files changed, 70 insertions(+), 4 deletions(-) diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index 9ad5e07..c494edc 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -66,6 +66,35 @@ function isMinimalDiagram(xml: string): boolean { return !stripped.includes('id="2"') } +// Helper function to replace historical tool call XML with placeholders +// This reduces token usage and forces LLM to rely on the current diagram XML (source of truth) +function replaceHistoricalToolInputs(messages: any[]): any[] { + return messages.map((msg) => { + if (msg.role !== "assistant" || !Array.isArray(msg.content)) { + return msg + } + const replacedContent = msg.content.map((part: any) => { + if (part.type === "tool-call") { + const toolName = part.toolName + if ( + toolName === "display_diagram" || + toolName === "edit_diagram" + ) { + return { + ...part, + input: { + placeholder: + "[XML content replaced - see current diagram XML in system context]", + }, + } + } + } + return part + }) + return { ...msg, content: replacedContent } + }) +} + // Helper function to fix tool call inputs for Bedrock API // Bedrock requires toolUse.input to be a JSON object, not a string function fixToolCallInputs(messages: any[]): any[] { @@ -144,7 +173,7 @@ async function handleChatRequest(req: Request): Promise { } } - const { messages, xml, sessionId } = await req.json() + const { messages, xml, previousXml, sessionId } = await req.json() // Get user IP for Langfuse tracking const forwardedFor = req.headers.get("x-forwarded-for") @@ -242,9 +271,12 @@ ${lastMessageText} // Fix tool call inputs for Bedrock API (requires JSON objects, not strings) const fixedMessages = fixToolCallInputs(modelMessages) + // Replace historical tool call XML with placeholders to reduce tokens and avoid confusion + const placeholderMessages = replaceHistoricalToolInputs(fixedMessages) + // Filter out messages with empty content arrays (Bedrock API rejects these) // This is a safety measure - ideally convertToModelMessages should handle all cases - let enhancedMessages = fixedMessages.filter( + let enhancedMessages = placeholderMessages.filter( (msg: any) => msg.content && Array.isArray(msg.content) && msg.content.length > 0, ) @@ -308,10 +340,10 @@ ${lastMessageText} }, }), }, - // Cache breakpoint 2: Current diagram XML context + // Cache breakpoint 2: Previous and Current diagram XML context { role: "system" as const, - content: `Current diagram XML:\n"""xml\n${xml || ""}\n"""\nWhen using edit_diagram, COPY search patterns exactly from this XML - attribute order matters!`, + content: `${previousXml ? `Previous diagram XML (before user's last message):\n"""xml\n${previousXml}\n"""\n\n` : ""}Current diagram XML (AUTHORITATIVE - the source of truth):\n"""xml\n${xml || ""}\n"""\n\nIMPORTANT: The "Current diagram XML" is the SINGLE SOURCE OF TRUTH for what's on the canvas right now. The user can manually add, delete, or modify shapes directly in draw.io. Always count and describe elements based on the CURRENT XML, not on what you previously generated. If both previous and current XML are shown, compare them to understand what the user changed. When using edit_diagram, COPY search patterns exactly from the CURRENT XML - attribute order matters!`, ...(shouldCache && { providerOptions: { bedrock: { cachePoint: { type: "default" } }, diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index 2b1d611..53f9022 100644 --- a/components/chat-panel.tsx +++ b/components/chat-panel.tsx @@ -764,6 +764,15 @@ Please retry with an adjusted search pattern or use display_diagram if retries a } } + // Get previous XML from the last snapshot (before this message) + const snapshotKeys = Array.from( + xmlSnapshotsRef.current.keys(), + ).sort((a, b) => b - a) + const previousXml = + snapshotKeys.length > 0 + ? xmlSnapshotsRef.current.get(snapshotKeys[0]) || "" + : "" + // Save XML snapshot for this message (will be at index = current messages.length) const messageIndex = messages.length xmlSnapshotsRef.current.set(messageIndex, chartXml) @@ -805,6 +814,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a { body: { xml: chartXml, + previousXml, sessionId, }, headers: { @@ -869,6 +879,15 @@ Please retry with an adjusted search pattern or use display_diagram if retries a return } + // Get previous XML (snapshot before the one being regenerated) + const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys()) + .filter((k) => k < userMessageIndex) + .sort((a, b) => b - a) + const previousXml = + snapshotKeys.length > 0 + ? xmlSnapshotsRef.current.get(snapshotKeys[0]) || "" + : "" + // Restore the diagram to the saved state (skip validation for trusted snapshots) onDisplayChart(savedXml, true) @@ -923,6 +942,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a { body: { xml: savedXml, + previousXml, sessionId, }, headers: { @@ -956,6 +976,15 @@ Please retry with an adjusted search pattern or use display_diagram if retries a return } + // Get previous XML (snapshot before the one being edited) + const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys()) + .filter((k) => k < messageIndex) + .sort((a, b) => b - a) + const previousXml = + snapshotKeys.length > 0 + ? xmlSnapshotsRef.current.get(snapshotKeys[0]) || "" + : "" + // Restore the diagram to the saved state (skip validation for trusted snapshots) onDisplayChart(savedXml, true) @@ -1018,6 +1047,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a { body: { xml: savedXml, + previousXml, sessionId, }, headers: { diff --git a/components/file-preview-list.tsx b/components/file-preview-list.tsx index 1bedbb9..cb61e0a 100644 --- a/components/file-preview-list.tsx +++ b/components/file-preview-list.tsx @@ -48,6 +48,8 @@ export function FilePreviewList({ files, onRemoveFile }: FilePreviewListProps) { imageUrlsRef.current.forEach((url) => { URL.revokeObjectURL(url) }) + // Clear the ref so StrictMode remount creates fresh URLs + imageUrlsRef.current = new Map() } }, []) @@ -83,6 +85,7 @@ export function FilePreviewList({ files, onRemoveFile }: FilePreviewListProps) { width={80} height={80} className="object-cover w-full h-full" + unoptimized /> ) : (
@@ -124,6 +127,7 @@ export function FilePreviewList({ files, onRemoveFile }: FilePreviewListProps) { height={900} className="object-contain max-w-full max-h-[90vh] w-auto h-auto" onClick={(e) => e.stopPropagation()} + unoptimized />