[Feature] Add VLM-based diagram validation

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

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>
)
}

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

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>
)
}

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}