[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

@@ -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,