i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-09-02 01:20:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Feature] Add VLM-based diagram validation

Browse Source 
Add automatic VLM (Vision Language Model) validation after display_diagram
tool execution. The system captures a screenshot of the rendered diagram,
sends it to a VLM for visual analysis, and uses feedback to improve
diagram quality through the existing retry mechanism.

Changes:
- Add /api/validate-diagram endpoint for VLM validation
- Add diagram-validator.ts for client-side validation orchestration
- Add validation-prompts.ts for VLM system prompts
- Add ValidationCard component to display validation status in chat
- Add PNG capture functionality to diagram context
- Integrate validation into tool handlers with retry support (max 3)
- Add "Improve with Suggestions" button for manual regeneration
- Add settings toggle to enable/disable VLM validation
- Add getValidationModel() helper in ai-providers.ts
This commit is contained in:
Jinze Yu
2026-01-17 17:09:38 +09:00
parent 5007c7bbe4
commit ee4c0149f1
11 changed files with 1071 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

124
app/api/validate-diagram/route.ts Normal file
View File

@@ -0,0 +1,124 @@
/**
* API endpoint for VLM-based diagram validation.
* Accepts a PNG image and returns validation results.
*/
import { generateText } from "ai"
import { getValidationModel } from "@/lib/ai-providers"
import type { ValidationResult } from "@/lib/diagram-validator"
import {
parseValidationResponse,
VALIDATION_SYSTEM_PROMPT,
} from "@/lib/validation-prompts"
export const maxDuration = 30
interface ValidateDiagramRequest {
imageData: string // Base64 PNG data URL
xml: string // Diagram XML for context
sessionId?: string
}
export async function POST(req: Request): Promise<Response> {
try {
// Check if VLM validation is enabled (default: true)
const enableValidation = process.env.ENABLE_VLM_VALIDATION !== "false"
if (!enableValidation) {
return Response.json({
valid: true,
issues: [],
suggestions: [],
} satisfies ValidationResult)
}
const body: ValidateDiagramRequest = await req.json()
const { imageData, xml, sessionId } = body
if (!imageData) {
return Response.json(
{ error: "Missing imageData" },
{ status: 400 },
)
}
// Validate image data format
if (
!imageData.startsWith("data:image/png;base64,") &&
!imageData.startsWith("data:image/")
) {
return Response.json(
{ error: "Invalid image data format" },
{ status: 400 },
)
}
// Get the validation model
let model
try {
model = getValidationModel()
} catch (error) {
console.warn(
"[validate-diagram] Validation model not available:",
error,
)
// Return valid if no vision model is configured
return Response.json({
valid: true,
issues: [],
suggestions: [],
} satisfies ValidationResult)
}
const timeout = parseInt(process.env.VALIDATION_TIMEOUT || "10000", 10)
// Call the VLM with the image
const result = await Promise.race([
generateText({
model,
system: VALIDATION_SYSTEM_PROMPT,
messages: [
{
role: "user",
content: [
{
type: "image",
image: imageData,
},
{
type: "text",
text: "Please analyze this diagram for visual quality issues and return your assessment as JSON.",
},
],
},
],
maxOutputTokens: 1024,
}),
new Promise<never>((_, reject) =>
setTimeout(
() => reject(new Error("Validation timeout")),
timeout,
),
),
])
// Parse the VLM response
const validationResult = parseValidationResponse(result.text)
if (sessionId) {
console.log(
`[validate-diagram] Session ${sessionId}: valid=${validationResult.valid}, issues=${validationResult.issues.length}`,
)
}
return Response.json(validationResult)
} catch (error) {
console.error("[validate-diagram] Error:", error)
// On error, return valid to not block the user
return Response.json({
valid: true,
issues: [],
suggestions: [],
} satisfies ValidationResult)
}
}

79
components/chat-message-display.tsx
View File

@@ -28,6 +28,8 @@ import {
import { ChatLobby } from "@/components/chat/ChatLobby"
import { ToolCallCard } from "@/components/chat/ToolCallCard"
import type { DiagramOperation, ToolPartLike } from "@/components/chat/types"
import type { ValidationState } from "@/components/chat/ValidationCard"
import { ValidationCard } from "@/components/chat/ValidationCard"
import { ScrollArea } from "@/components/ui/scroll-area"
import { useDictionary } from "@/hooks/use-dictionary"
import { getApiEndpoint } from "@/lib/base-path"
@@ -148,6 +150,8 @@ interface ChatMessageDisplayProps {
onSelectSession?: (id: string) => void
onDeleteSession?: (id: string) => void
loadedMessageIdsRef?: MutableRefObject<Set<string>>
validationStates?: Record<string, ValidationState>
onImproveWithSuggestions?: (feedback: string) => void
}
export function ChatMessageDisplay({
@@ -165,6 +169,8 @@ export function ChatMessageDisplay({
onSelectSession,
onDeleteSession,
loadedMessageIdsRef,
validationStates = {},
onImproveWithSuggestions,
}: ChatMessageDisplayProps) {
const dict = useDictionary()
const { chartXML, loadDiagram: onDisplayChart } = useDiagram()
@@ -911,30 +917,59 @@ export function ChatMessageDisplay({
return groups.map(
(group, groupIndex) => {
if (group.type === "tool") {
const toolPart = group
.parts[0] as ToolPartLike
const toolCallId =
toolPart.toolCallId
const isDisplayDiagram =
toolPart.type ===
"tool-display_diagram"
const validationState =
validationStates[
toolCallId
]
return (
<ToolCallCard
<div
key={`${message.id}-tool-${group.startIndex}`}
part={
group
.parts[0] as ToolPartLike
}
expandedTools={
expandedTools
}
setExpandedTools={
setExpandedTools
}
onCopy={
copyMessageToClipboard
}
copiedToolCallId={
copiedToolCallId
}
copyFailedToolCallId={
copyFailedToolCallId
}
dict={dict}
/>
>
<ToolCallCard
part={
toolPart
}
expandedTools={
expandedTools
}
setExpandedTools={
setExpandedTools
}
onCopy={
copyMessageToClipboard
}
copiedToolCallId={
copiedToolCallId
}
copyFailedToolCallId={
copyFailedToolCallId
}
dict={dict}
/>
{/* Show validation card for display_diagram tools */}
{isDisplayDiagram &&
validationState && (
<ValidationCard
state={
validationState
}
dict={
dict
}
onImproveWithSuggestions={
onImproveWithSuggestions
}
/>
)}
</div>
)
}

65
components/chat-panel.tsx
View File

@@ -34,10 +34,12 @@ import { findCachedResponse } from "@/lib/cached-responses"
import { formatMessage } from "@/lib/i18n/utils"
import { isPdfFile, isTextFile } from "@/lib/pdf-utils"
import { sanitizeMessages } from "@/lib/session-storage"
import { STORAGE_KEYS } from "@/lib/storage"
import type { UrlData } from "@/lib/url-utils"
import { type FileData, useFileProcessor } from "@/lib/use-file-processor"
import { useQuotaManager } from "@/lib/use-quota-manager"
import { cn, formatXML, isRealDiagram } from "@/lib/utils"
import type { ValidationState } from "./chat/ValidationCard"
import { ChatMessageDisplay } from "./chat-message-display"
import { DevXmlSimulator } from "./dev-xml-simulator"
@@ -75,7 +77,8 @@ interface ChatPanelProps {
// Constants for tool states
const TOOL_ERROR_STATE = "output-error" as const
const DEBUG = process.env.NODE_ENV === "development"
const MAX_AUTO_RETRY_COUNT = 1
// Increased to 3 to support VLM validation retries (matches MAX_VALIDATION_RETRIES)
const MAX_AUTO_RETRY_COUNT = 3
const MAX_CONTINUATION_RETRY_COUNT = 2 // Limit for truncation continuation retries
@@ -120,6 +123,7 @@ export default function ChatPanel({
latestSvg,
clearDiagram,
getThumbnailSvg,
captureValidationPng,
diagramHistory,
setDiagramHistory,
} = useDiagram()
@@ -173,6 +177,7 @@ export default function ChatPanel({
const [dailyTokenLimit, setDailyTokenLimit] = useState(0)
const [tpmLimit, setTpmLimit] = useState(0)
const [minimalStyle, setMinimalStyle] = useState(false)
const [vlmValidationEnabled, setVlmValidationEnabled] = useState(true)
// Restore input from sessionStorage on mount (when ChatPanel remounts due to key change)
useEffect(() => {
@@ -182,6 +187,14 @@ export default function ChatPanel({
}
}, [])
// Load VLM validation setting from localStorage on mount
useEffect(() => {
const stored = localStorage.getItem(STORAGE_KEYS.vlmValidationEnabled)
if (stored !== null) {
setVlmValidationEnabled(stored === "true")
}
}, [])
// Check config on mount
useEffect(() => {
fetch(getApiEndpoint("/api/config"))
@@ -269,6 +282,43 @@ export default function ChatPanel({
> | null>(null)
const LOCAL_STORAGE_DEBOUNCE_MS = 1000 // Save at most once per second
// Validation state for displaying VLM validation progress
// Key: toolCallId, Value: ValidationState
const [validationStates, setValidationStates] = useState<
Record<string, ValidationState>
>({})
// Callback to update validation state from tool handler
const handleValidationStateChange = useCallback(
(toolCallId: string, state: ValidationState) => {
setValidationStates((prev) => ({
...prev,
[toolCallId]: state,
}))
},
[],
)
// Handler for VLM validation setting change
const handleVlmValidationChange = useCallback((value: boolean) => {
setVlmValidationEnabled(value)
localStorage.setItem(STORAGE_KEYS.vlmValidationEnabled, String(value))
}, [])
// Ref to store the sendMessage function for use in callbacks
const sendMessageRef = useRef<typeof sendMessage | null>(null)
// Callback to improve diagram with validation suggestions
const handleImproveWithSuggestions = useCallback((feedback: string) => {
if (sendMessageRef.current) {
// Send the feedback as a new user message to trigger regeneration
sendMessageRef.current({
role: "user",
parts: [{ type: "text", text: feedback }],
})
}
}, [])
// Diagram tool handlers (display_diagram, edit_diagram, append_diagram)
const { handleToolCall } = useDiagramToolHandlers({
partialXmlRef,
@@ -277,6 +327,10 @@ export default function ChatPanel({
onDisplayChart,
onFetchChart,
onExport,
captureValidationPng,
enableVlmValidation: vlmValidationEnabled,
sessionId,
onValidationStateChange: handleValidationStateChange,
})
const { messages, sendMessage, addToolOutput, status, error, setMessages } =
@@ -425,6 +479,11 @@ export default function ChatPanel({
},
})
// Store sendMessage in ref for use in callbacks (like handleImproveWithSuggestions)
useEffect(() => {
sendMessageRef.current = sendMessage
}, [sendMessage])
// Ref to track latest messages for unload persistence
const messagesRef = useRef(messages)
useEffect(() => {
@@ -1261,6 +1320,8 @@ export default function ChatPanel({
onSelectSession={handleSelectSession}
onDeleteSession={handleDeleteSession}
loadedMessageIdsRef={loadedMessageIdsRef}
validationStates={validationStates}
onImproveWithSuggestions={handleImproveWithSuggestions}
/>
</main>
@@ -1308,6 +1369,8 @@ export default function ChatPanel({
onToggleDarkMode={onToggleDarkMode}
minimalStyle={minimalStyle}
onMinimalStyleChange={setMinimalStyle}
vlmValidationEnabled={vlmValidationEnabled}
onVlmValidationChange={handleVlmValidationChange}
/>
<ModelConfigDialog

303
components/chat/ValidationCard.tsx Normal file
View File

@@ -0,0 +1,303 @@
"use client"
import {
AlertTriangle,
Check,
ChevronDown,
ChevronUp,
Eye,
ImageIcon,
RefreshCw,
X,
} from "lucide-react"
import Image from "next/image"
import { useEffect, useState } from "react"
import type { ValidationResult } from "@/lib/diagram-validator"
export type ValidationStatus =
| "idle"
| "capturing"
| "validating"
| "success"
| "failed"
| "error"
| "skipped"
export interface ValidationState {
status: ValidationStatus
attempt?: number
maxAttempts?: number
result?: ValidationResult
error?: string
imageData?: string // Base64 PNG data URL
}
interface ValidationCardProps {
state: ValidationState
dict: {
tools: { complete: string }
}
onImproveWithSuggestions?: (feedback: string) => void
}
export function ValidationCard({
state,
dict,
onImproveWithSuggestions,
}: ValidationCardProps) {
const [isExpanded, setIsExpanded] = useState(
state.status === "validating" || state.status === "failed",
)
const [hasRequestedImprovement, setHasRequestedImprovement] =
useState(false)
// Generate improvement feedback from validation result
const generateImprovementFeedback = (): string => {
if (!state.result) return ""
const lines: string[] = []
lines.push(
"Please improve the diagram based on the following visual analysis feedback:",
)
lines.push("")
if (state.result.issues.length > 0) {
lines.push("Issues to address:")
for (const issue of state.result.issues) {
lines.push(
` - [${issue.severity}] ${issue.type}: ${issue.description}`,
)
}
lines.push("")
}
if (state.result.suggestions.length > 0) {
lines.push("Suggestions for improvement:")
for (const suggestion of state.result.suggestions) {
lines.push(` - ${suggestion}`)
}
lines.push("")
}
lines.push("Regenerate the diagram with these improvements applied.")
return lines.join("\n")
}
const handleImproveClick = () => {
if (
!onImproveWithSuggestions ||
!state.result ||
hasRequestedImprovement
)
return
setHasRequestedImprovement(true)
const feedback = generateImprovementFeedback()
onImproveWithSuggestions(feedback)
}
// Check if we should show the improve button
const showImproveButton =
onImproveWithSuggestions &&
state.result &&
(state.status === "success" || state.status === "skipped") &&
(state.result.issues.length > 0 || state.result.suggestions.length > 0)
const getStatusDisplay = () => {
switch (state.status) {
case "capturing":
return {
label: "Capturing",
color: "text-blue-600 bg-blue-50",
icon: (
<div className="h-4 w-4 border-2 border-blue-600 border-t-transparent rounded-full animate-spin" />
),
}
case "validating":
return {
label: state.attempt
? `Validating (${state.attempt}/${state.maxAttempts || 3})`
: "Validating",
color: "text-blue-600 bg-blue-50",
icon: (
<div className="h-4 w-4 border-2 border-blue-600 border-t-transparent rounded-full animate-spin" />
),
}
case "success":
return {
label: dict.tools.complete,
color: "text-green-600 bg-green-50",
icon: <Check className="h-4 w-4" />,
}
case "failed":
return {
label: "Issues Found",
color: "text-yellow-600 bg-yellow-50",
icon: <AlertTriangle className="h-4 w-4" />,
}
case "error":
return {
label: "Error",
color: "text-red-600 bg-red-50",
icon: <X className="h-4 w-4" />,
}
case "skipped":
return {
label: "Skipped",
color: "text-gray-600 bg-gray-50",
icon: <Check className="h-4 w-4" />,
}
default:
return null
}
}
const statusDisplay = getStatusDisplay()
if (!statusDisplay || state.status === "idle") return null
return (
<div className="my-3 rounded-xl border border-border/60 bg-muted/30 overflow-hidden">
<div className="flex items-center justify-between px-4 py-3 bg-muted/50">
<div className="flex items-center gap-2">
<div className="w-6 h-6 rounded-md bg-primary/10 flex items-center justify-center">
<Eye className="w-3.5 h-3.5 text-primary" />
</div>
<span className="text-sm font-medium text-foreground/80">
Validate Diagram
</span>
</div>
<div className="flex items-center gap-2">
<span
className={`text-xs font-medium px-2 py-0.5 rounded-full flex items-center gap-1 ${statusDisplay.color}`}
>
{statusDisplay.icon}
<span className="ml-1">{statusDisplay.label}</span>
</span>
{(state.result || state.error) && (
<button
type="button"
onClick={() => setIsExpanded(!isExpanded)}
className="p-1 rounded hover:bg-muted transition-colors"
>
{isExpanded ? (
<ChevronUp className="w-4 h-4 text-muted-foreground" />
) : (
<ChevronDown className="w-4 h-4 text-muted-foreground" />
)}
</button>
)}
</div>
</div>
{/* Validation details when expanded */}
{isExpanded && (state.result || state.imageData) && (
<div className="px-4 py-3 border-t border-border/40 bg-muted/20 space-y-3">
{/* Captured image */}
{state.imageData && (
<div>
<div className="text-xs font-medium text-foreground/70 mb-2 flex items-center gap-1">
<ImageIcon className="h-3 w-3" />
Captured Screenshot:
</div>
<div className="rounded-lg border border-border/50 overflow-hidden bg-white">
<Image
src={state.imageData}
alt="Captured diagram for validation"
width={400}
height={300}
className="w-full h-auto max-h-48 object-contain"
unoptimized
/>
</div>
</div>
)}
{/* Issues */}
{state.result && state.result.issues.length > 0 && (
<div>
<div className="text-xs font-medium text-foreground/70 mb-2">
Issues Found:
</div>
<div className="space-y-2">
{state.result.issues.map((issue, index) => (
<div
key={index}
className={`text-xs px-3 py-2 rounded-lg border ${
issue.severity === "critical"
? "bg-red-50 border-red-200 text-red-700 dark:bg-red-950 dark:border-red-800 dark:text-red-300"
: "bg-yellow-50 border-yellow-200 text-yellow-700 dark:bg-yellow-950 dark:border-yellow-800 dark:text-yellow-300"
}`}
>
<span className="font-medium uppercase text-[10px] mr-2">
[{issue.type}]
</span>
{issue.description}
</div>
))}
</div>
</div>
)}
{/* Suggestions */}
{state.result && state.result.suggestions.length > 0 && (
<div>
<div className="text-xs font-medium text-foreground/70 mb-2">
Suggestions:
</div>
<ul className="text-xs text-foreground/60 space-y-1 list-disc list-inside">
{state.result.suggestions.map(
(suggestion, index) => (
<li key={index}>{suggestion}</li>
),
)}
</ul>
</div>
)}
{/* Valid result message */}
{state.result &&
state.result.valid &&
state.result.issues.length === 0 && (
<div className="text-xs text-green-600 dark:text-green-400">
Diagram passed visual validation - no issues
detected.
</div>
)}
</div>
)}
{/* Improve with Suggestions button - shown when validation passed but has suggestions */}
{showImproveButton && (
<div className="px-4 py-3 border-t border-border/40 bg-muted/10">
{hasRequestedImprovement ? (
<div className="flex items-center justify-center gap-2 px-4 py-2 text-sm font-medium text-green-600 dark:text-green-400">
<Check className="h-4 w-4" />
Improvement requested - check the new diagram below
</div>
) : (
<>
<button
type="button"
onClick={handleImproveClick}
className="w-full flex items-center justify-center gap-2 px-4 py-2 text-sm font-medium text-primary bg-primary/10 hover:bg-primary/20 rounded-lg transition-colors"
>
<RefreshCw className="h-4 w-4" />
Improve with Suggestions
</button>
<p className="text-xs text-muted-foreground mt-2 text-center">
Regenerate the diagram using the validation
feedback
</p>
</>
)}
</div>
)}
{/* Error details when expanded */}
{isExpanded && state.error && (
<div className="px-4 py-3 border-t border-border/40 bg-red-50/50">
<div className="text-xs text-red-600">{state.error}</div>
</div>
)}
</div>
)
}

21
components/settings-dialog.tsx
View File

@@ -67,6 +67,8 @@ interface SettingsDialogProps {
onToggleDarkMode: () => void
minimalStyle?: boolean
onMinimalStyleChange?: (value: boolean) => void
vlmValidationEnabled?: boolean
onVlmValidationChange?: (value: boolean) => void
}
export const STORAGE_ACCESS_CODE_KEY = "next-ai-draw-io-access-code"
@@ -88,6 +90,8 @@ function SettingsContent({
onToggleDarkMode,
minimalStyle = false,
onMinimalStyleChange = () => {},
vlmValidationEnabled = true,
onVlmValidationChange = () => {},
}: SettingsDialogProps) {
const dict = useDictionary()
const router = useRouter()
@@ -403,6 +407,23 @@ function SettingsContent({
</div>
</SettingItem>
{/* VLM Diagram Validation */}
<SettingItem
label="Diagram Validation"
description="Use AI vision to validate and improve generated diagrams"
>
<div className="flex items-center gap-2">
<Switch
id="vlm-validation"
checked={vlmValidationEnabled}
onCheckedChange={onVlmValidationChange}
/>
<span className="text-sm text-muted-foreground">
{vlmValidationEnabled ? "Enabled" : "Disabled"}
</span>
</div>
</SettingItem>
{/* Send Shortcut */}
<SettingItem
label={dict.settings.sendShortcut}

42
contexts/diagram-context.tsx
View File

@@ -31,6 +31,7 @@ interface DiagramContextType {
successMessage?: string,
) => void
getThumbnailSvg: () => Promise<string | null>
captureValidationPng: () => Promise<string | null>
isDrawioReady: boolean
onDrawioLoad: () => void
resetDrawioReady: () => void
@@ -51,6 +52,8 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
const hasCalledOnLoadRef = useRef(false)
const drawioRef = useRef<DrawIoEmbedRef | null>(null)
const resolverRef = useRef<((value: string) => void) | null>(null)
// Resolver for PNG export (used for VLM validation)
const pngResolverRef = useRef<((value: string) => void) | null>(null)
// Track if we're expecting an export for history (user-initiated)
const expectHistoryExportRef = useRef<boolean>(false)
// Track if diagram has been restored after DrawIO remount (e.g., theme change)
@@ -147,6 +150,37 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
}
}
// Capture current diagram as PNG for VLM validation
const captureValidationPng = async (): Promise<string | null> => {
if (!drawioRef.current) return null
// Don't export if diagram is empty
if (!isRealDiagram(chartXML)) return null
try {
const pngData = await Promise.race([
new Promise<string>((resolve) => {
pngResolverRef.current = resolve
drawioRef.current?.exportDiagram({ format: "png" })
}),
new Promise<string>((_, reject) =>
setTimeout(
() => reject(new Error("PNG export timeout")),
5000,
),
),
])
// PNG data should be a base64 data URL
if (pngData?.startsWith("data:image/png")) {
return pngData
}
return null
} catch {
// Timeout is expected occasionally - don't log as error
return null
}
}
const loadDiagram = (
chart: string,
skipValidation?: boolean,
@@ -186,6 +220,13 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
}
const handleDiagramExport = (data: any) => {
// Handle PNG export for VLM validation
if (pngResolverRef.current && data.data?.startsWith("data:image/png")) {
pngResolverRef.current(data.data)
pngResolverRef.current = null
return
}
// Handle save to file if requested (process raw data before extraction)
if (saveResolverRef.current.resolver) {
const format = saveResolverRef.current.format
@@ -353,6 +394,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
clearDiagram,
saveDiagramToFile,
getThumbnailSvg,
captureValidationPng,
isDrawioReady,
onDrawioLoad,
resetDrawioReady,

191
hooks/use-diagram-tool-handlers.ts
View File

@@ -1,5 +1,15 @@
import type { MutableRefObject } from "react"
import { useRef } from "react"
import type { DiagramOperation } from "@/components/chat/types"
import type {
ValidationState,
ValidationStatus,
} from "@/components/chat/ValidationCard"
import type { ValidationResult } from "@/lib/diagram-validator"
import {
formatValidationFeedback,
validateRenderedDiagram,
} from "@/lib/diagram-validator"
import { isMxCellXmlComplete, wrapWithMxFile } from "@/lib/utils"
const DEBUG = process.env.NODE_ENV === "development"
@@ -30,6 +40,8 @@ type AddToolOutputParams = AddToolOutputSuccess | AddToolOutputError
type AddToolOutputFn = (params: AddToolOutputParams) => void
const MAX_VALIDATION_RETRIES = 3
interface UseDiagramToolHandlersParams {
partialXmlRef: MutableRefObject<string>
editDiagramOriginalXmlRef: MutableRefObject<Map<string, string>>
@@ -37,6 +49,13 @@ interface UseDiagramToolHandlersParams {
onDisplayChart: (xml: string, skipValidation?: boolean) => string | null
onFetchChart: (saveToHistory?: boolean) => Promise<string>
onExport: () => void
captureValidationPng?: () => Promise<string | null>
enableVlmValidation?: boolean
sessionId?: string
onValidationStateChange?: (
toolCallId: string,
state: ValidationState,
) => void
}
/**
@@ -53,7 +72,33 @@ export function useDiagramToolHandlers({
onDisplayChart,
onFetchChart,
onExport,
captureValidationPng,
enableVlmValidation = true,
sessionId,
onValidationStateChange,
}: UseDiagramToolHandlersParams) {
// Track validation retry count per tool call
const validationRetryCountRef = useRef<Map<string, number>>(new Map())
// Helper to update validation state
const updateValidationState = (
toolCallId: string,
status: ValidationStatus,
options?: {
attempt?: number
maxAttempts?: number
result?: ValidationResult
error?: string
imageData?: string
},
) => {
if (onValidationStateChange) {
onValidationStateChange(toolCallId, {
status,
...options,
})
}
}
const handleToolCall = async (
{ toolCall }: { toolCall: ToolCall },
addToolOutput: AddToolOutputFn,
@@ -155,7 +200,151 @@ ${finalXml}
// Success - diagram will be rendered by chat-message-display
if (DEBUG) {
console.log(
"[display_diagram] Success! Adding tool output with state: output-available",
"[display_diagram] Success! Checking if VLM validation is enabled...",
)
}
// VLM validation after successful display
if (enableVlmValidation && captureValidationPng) {
let capturedPngData: string | null = null
try {
// Notify UI that we're starting capture
updateValidationState(toolCall.toolCallId, "capturing")
// Small delay to ensure diagram is rendered before capture
await new Promise((resolve) => setTimeout(resolve, 100))
capturedPngData = await captureValidationPng()
if (capturedPngData) {
if (DEBUG) {
console.log(
"[display_diagram] Captured PNG for validation",
)
}
const retryCount =
validationRetryCountRef.current.get(
toolCall.toolCallId,
) || 0
// Notify UI that we're validating (include the image)
updateValidationState(
toolCall.toolCallId,
"validating",
{
attempt: retryCount + 1,
maxAttempts: MAX_VALIDATION_RETRIES,
imageData: capturedPngData,
},
)
const result = await validateRenderedDiagram(
capturedPngData,
finalXml,
sessionId,
)
if (!result.valid) {
if (retryCount < MAX_VALIDATION_RETRIES) {
validationRetryCountRef.current.set(
toolCall.toolCallId,
retryCount + 1,
)
const feedback =
formatValidationFeedback(result)
if (DEBUG) {
console.log(
`[display_diagram] Validation failed (attempt ${retryCount + 1}/${MAX_VALIDATION_RETRIES}):`,
result.issues,
)
}
// Notify UI of validation failure (include the image)
updateValidationState(
toolCall.toolCallId,
"failed",
{
attempt: retryCount + 1,
maxAttempts: MAX_VALIDATION_RETRIES,
result,
imageData: capturedPngData,
},
)
addToolOutput({
tool: "display_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `[Validation attempt ${retryCount + 1}/${MAX_VALIDATION_RETRIES}]\n${feedback}`,
})
return
} else {
// Max retries reached - accept the diagram with warning
if (DEBUG) {
console.log(
"[display_diagram] Max validation retries reached, accepting diagram",
)
}
validationRetryCountRef.current.delete(
toolCall.toolCallId,
)
// Notify UI that we're accepting with issues (include the image)
updateValidationState(
toolCall.toolCallId,
"skipped",
{ result, imageData: capturedPngData },
)
addToolOutput({
tool: "display_diagram",
toolCallId: toolCall.toolCallId,
output: "Diagram displayed (validation issues noted but max retries reached).",
})
return
}
} else {
// Validation passed - clean up retry count
validationRetryCountRef.current.delete(
toolCall.toolCallId,
)
if (DEBUG) {
console.log(
"[display_diagram] Validation passed!",
)
}
// Notify UI of success (include the image)
updateValidationState(
toolCall.toolCallId,
"success",
{ result, imageData: capturedPngData },
)
}
} else {
// PNG capture failed - skip validation
updateValidationState(toolCall.toolCallId, "skipped")
}
} catch (error) {
// VLM validation error - log but don't block the user
console.warn(
"[display_diagram] VLM validation error:",
error,
)
updateValidationState(toolCall.toolCallId, "error", {
error:
error instanceof Error
? error.message
: "Validation failed",
imageData: capturedPngData || undefined,
})
}
}
if (DEBUG) {
console.log(
"[display_diagram] Adding tool output with state: output-available",
)
}
addToolOutput({

24
lib/ai-providers.ts
View File

@@ -1177,3 +1177,27 @@ export function supportsImageInput(modelId: string): boolean {
// Default: assume model supports images
return true
}
/**
* Get the AI model for diagram validation.
* Uses VALIDATION_MODEL env var if set, otherwise falls back to AI_MODEL.
* Throws if the model doesn't support image input.
*/
export function getValidationModel(): ReturnType<typeof getAIModel>["model"] {
const modelId = process.env.VALIDATION_MODEL || process.env.AI_MODEL
if (!modelId) {
throw new Error(
"No validation model configured. Set VALIDATION_MODEL or AI_MODEL.",
)
}
if (!supportsImageInput(modelId)) {
throw new Error(
`Validation requires a vision-capable model. Model "${modelId}" does not support image input.`,
)
}
const { model } = getAIModel({ modelId })
return model
}

121
lib/diagram-validator.ts Normal file
View File

@@ -0,0 +1,121 @@
/**
* Client-side validation orchestration for VLM-based diagram validation.
*/
import { getApiEndpoint } from "./base-path"
export interface ValidationIssue {
type: "overlap" | "edge_routing" | "text" | "layout" | "rendering"
severity: "critical" | "warning"
description: string
}
export interface ValidationResult {
valid: boolean
issues: ValidationIssue[]
suggestions: string[]
}
/**
* Validate a rendered diagram by sending its PNG to the VLM validation API.
*
* @param pngDataUrl - Base64 PNG data URL from draw.io export
* @param xml - The diagram XML (for context in error messages)
* @param sessionId - Optional session ID for logging
* @returns ValidationResult with issues and suggestions
*/
export async function validateRenderedDiagram(
pngDataUrl: string,
xml: string,
sessionId?: string,
): Promise<ValidationResult> {
try {
const response = await fetch(getApiEndpoint("/api/validate-diagram"), {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
imageData: pngDataUrl,
xml,
sessionId,
}),
})
if (!response.ok) {
const errorData = await response.json().catch(() => ({}))
console.error(
"[validateRenderedDiagram] API error:",
response.status,
errorData,
)
// Return valid on API error to not block the user
return {
valid: true,
issues: [],
suggestions: [],
}
}
const result: ValidationResult = await response.json()
return result
} catch (error) {
console.error("[validateRenderedDiagram] Failed:", error)
// Return valid on network error to not block the user
return {
valid: true,
issues: [],
suggestions: [],
}
}
}
/**
* Format validation feedback for display to the AI model.
* This creates a human-readable error message that guides the AI to fix issues.
*
* @param result - The validation result from VLM
* @returns Formatted string for tool error output
*/
export function formatValidationFeedback(result: ValidationResult): string {
const lines: string[] = []
lines.push("DIAGRAM VISUAL VALIDATION FAILED")
lines.push("")
// Group issues by severity
const criticalIssues = result.issues.filter(
(i) => i.severity === "critical",
)
const warnings = result.issues.filter((i) => i.severity === "warning")
if (criticalIssues.length > 0) {
lines.push("Critical Issues (must fix):")
for (const issue of criticalIssues) {
lines.push(` - [${issue.type}] ${issue.description}`)
}
lines.push("")
}
if (warnings.length > 0) {
lines.push("Warnings:")
for (const issue of warnings) {
lines.push(` - [${issue.type}] ${issue.description}`)
}
lines.push("")
}
if (result.suggestions.length > 0) {
lines.push("Suggestions to fix:")
for (const suggestion of result.suggestions) {
lines.push(` - ${suggestion}`)
}
lines.push("")
}
lines.push(
"Please regenerate the diagram with corrected layout to fix these visual issues.",
)
return lines.join("\n")
}

3
lib/storage.ts
View File

@@ -24,4 +24,7 @@ export const STORAGE_KEYS = {
// Chat input preferences
sendShortcut: "next-ai-draw-io-send-shortcut",
// Diagram validation
vlmValidationEnabled: "next-ai-draw-io-vlm-validation-enabled",
} as const

122
lib/validation-prompts.ts Normal file
View File

@@ -0,0 +1,122 @@
/**
* VLM system prompt and response parsing for diagram validation.
*/
import type { ValidationResult } from "./diagram-validator"
export const VALIDATION_SYSTEM_PROMPT = `You are a diagram quality validator. Analyze the rendered diagram image for visual issues.
Evaluate the diagram for the following issues:
1. **Overlapping elements** (critical): Shapes covering each other inappropriately, making content unreadable
2. **Edge routing issues** (critical): Lines/arrows crossing through shapes that are not their source or target
3. **Text readability** (warning): Labels cut off, overlapping, or too small to read
4. **Layout quality** (warning): Poor spacing, misalignment, or cramped elements
5. **Rendering errors** (critical): Incomplete, corrupted, or missing visual elements
Return your analysis as a JSON object with this structure:
{
"valid": boolean,
"issues": [
{
"type": "overlap" | "edge_routing" | "text" | "layout" | "rendering",
"severity": "critical" | "warning",
"description": "Clear description of the issue and where it occurs"
}
],
"suggestions": ["Specific actionable suggestions to fix the issues"]
}
Rules:
- Set "valid" to true ONLY if there are no critical issues
- Be specific about which elements have problems (e.g., "The 'Login' box overlaps with 'Register' box")
- Provide actionable suggestions (e.g., "Move the Login box 50 pixels to the left")
- Minor cosmetic issues (slight misalignment, non-uniform spacing) should be warnings, not critical
- Empty diagrams or diagrams with only 1-2 elements should pass unless they have obvious errors
- If the diagram looks generally acceptable, set valid to true even with minor warnings
Return ONLY the JSON object, no additional text.`
/**
* Parse the VLM response text into a ValidationResult.
* Handles various response formats and edge cases.
*/
export function parseValidationResponse(text: string): ValidationResult {
try {
// Try to extract JSON from the response
// The VLM might wrap the JSON in markdown code blocks
let jsonStr = text.trim()
// Remove markdown code block if present
const jsonMatch = jsonStr.match(/```(?:json)?\s*([\s\S]*?)```/)
if (jsonMatch) {
jsonStr = jsonMatch[1].trim()
}
// Try to find JSON object if there's extra text
const objectMatch = jsonStr.match(/\{[\s\S]*\}/)
if (objectMatch) {
jsonStr = objectMatch[0]
}
const parsed = JSON.parse(jsonStr)
// Validate the structure
if (typeof parsed.valid !== "boolean") {
console.warn(
"[parseValidationResponse] Missing or invalid 'valid' field",
)
return {
valid: true, // Default to valid if parsing fails
issues: [],
suggestions: [],
}
}
// Normalize issues array
const issues = Array.isArray(parsed.issues)
? parsed.issues
.filter(
(issue: any) =>
issue &&
typeof issue.type === "string" &&
typeof issue.severity === "string" &&
typeof issue.description === "string",
)
.map((issue: any) => ({
type: issue.type as
| "overlap"
| "edge_routing"
| "text"
| "layout"
| "rendering",
severity: issue.severity as "critical" | "warning",
description: issue.description,
}))
: []
// Normalize suggestions array
const suggestions = Array.isArray(parsed.suggestions)
? parsed.suggestions.filter((s: any) => typeof s === "string")
: []
return {
valid: parsed.valid,
issues,
suggestions,
}
} catch (error) {
console.error(
"[parseValidationResponse] Failed to parse VLM response:",
error,
)
console.error("[parseValidationResponse] Raw response:", text)
// If parsing fails, default to valid to avoid blocking the user
return {
valid: true,
issues: [],
suggestions: [],
}
}
}