mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-09-02 01:20:23 +08:00
refactor(validation): use AI SDK structured outputs and address review feedback
- Replace generateText + manual JSON parsing with generateObject and Zod schema for type-safe structured validation output - Use AbortSignal.timeout() instead of Promise.race for cleaner timeout handling - Add timeout validation with minimum 1000ms to handle malformed env values - Remove unused xml parameter from validateRenderedDiagram API - Remove parseValidationResponse function (now handled by schema) - Clear validationStates on session switch and new chat to prevent memory leak - Update 100ms render delay comment to clarify best-effort heuristic - Remove unused useEffect import from ValidationCard - Fix optional chaining lint warning in ValidationCard - Add unit tests for formatValidationFeedback function
This commit is contained in:
@@ -3,19 +3,45 @@
|
||||
* Accepts a PNG image and returns validation results.
|
||||
*/
|
||||
|
||||
import { generateText } from "ai"
|
||||
import { generateObject } from "ai"
|
||||
import { z } from "zod"
|
||||
import { getValidationModel } from "@/lib/ai-providers"
|
||||
import type { ValidationResult } from "@/lib/diagram-validator"
|
||||
import {
|
||||
parseValidationResponse,
|
||||
VALIDATION_SYSTEM_PROMPT,
|
||||
} from "@/lib/validation-prompts"
|
||||
import { VALIDATION_SYSTEM_PROMPT } from "@/lib/validation-prompts"
|
||||
|
||||
export const maxDuration = 30
|
||||
|
||||
// Schema for structured validation output
|
||||
const ValidationResultSchema = z.object({
|
||||
valid: z.boolean().describe("True if there are no critical issues"),
|
||||
issues: z
|
||||
.array(
|
||||
z.object({
|
||||
type: z
|
||||
.enum([
|
||||
"overlap",
|
||||
"edge_routing",
|
||||
"text",
|
||||
"layout",
|
||||
"rendering",
|
||||
])
|
||||
.describe("Type of visual issue"),
|
||||
severity: z
|
||||
.enum(["critical", "warning"])
|
||||
.describe("Severity level"),
|
||||
description: z
|
||||
.string()
|
||||
.describe("Clear description of the issue"),
|
||||
}),
|
||||
)
|
||||
.describe("List of visual issues found"),
|
||||
suggestions: z
|
||||
.array(z.string())
|
||||
.describe("Actionable suggestions to fix issues"),
|
||||
})
|
||||
|
||||
interface ValidateDiagramRequest {
|
||||
imageData: string // Base64 PNG data URL
|
||||
xml: string // Diagram XML for context
|
||||
sessionId?: string
|
||||
}
|
||||
|
||||
@@ -32,7 +58,7 @@ export async function POST(req: Request): Promise<Response> {
|
||||
}
|
||||
|
||||
const body: ValidateDiagramRequest = await req.json()
|
||||
const { imageData, xml, sessionId } = body
|
||||
const { imageData, sessionId } = body
|
||||
|
||||
if (!imageData) {
|
||||
return Response.json(
|
||||
@@ -69,40 +95,38 @@ export async function POST(req: Request): Promise<Response> {
|
||||
} satisfies ValidationResult)
|
||||
}
|
||||
|
||||
const timeout = parseInt(process.env.VALIDATION_TIMEOUT || "10000", 10)
|
||||
// Parse timeout with validation (minimum 1000ms, default 10000ms)
|
||||
const timeout =
|
||||
Math.max(
|
||||
1000,
|
||||
parseInt(process.env.VALIDATION_TIMEOUT || "10000", 10),
|
||||
) || 10000
|
||||
|
||||
// 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,
|
||||
),
|
||||
),
|
||||
])
|
||||
// Call the VLM with structured output schema
|
||||
const result = await generateObject({
|
||||
model,
|
||||
schema: ValidationResultSchema,
|
||||
system: VALIDATION_SYSTEM_PROMPT,
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: [
|
||||
{
|
||||
type: "image",
|
||||
image: imageData,
|
||||
},
|
||||
{
|
||||
type: "text",
|
||||
text: "Please analyze this diagram for visual quality issues.",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
maxOutputTokens: 1024,
|
||||
abortSignal: AbortSignal.timeout(timeout),
|
||||
})
|
||||
|
||||
// Parse the VLM response
|
||||
const validationResult = parseValidationResponse(result.text)
|
||||
const validationResult: ValidationResult = result.object
|
||||
|
||||
if (sessionId) {
|
||||
console.log(
|
||||
@@ -112,7 +136,10 @@ export async function POST(req: Request): Promise<Response> {
|
||||
|
||||
return Response.json(validationResult)
|
||||
} catch (error) {
|
||||
console.error("[validate-diagram] Error:", error)
|
||||
// Log with session context if available
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : String(error)
|
||||
console.error("[validate-diagram] Error:", errorMessage)
|
||||
|
||||
// On error, return valid to not block the user
|
||||
return Response.json({
|
||||
|
||||
@@ -878,6 +878,7 @@ export default function ChatPanel({
|
||||
} else {
|
||||
justLoadedSessionIdRef.current = null
|
||||
}
|
||||
setValidationStates({}) // Clear validation states when switching sessions
|
||||
syncUIWithSession(sessionData)
|
||||
router.replace(`?session=${sessionId}`, { scroll: false })
|
||||
}
|
||||
@@ -917,6 +918,7 @@ export default function ChatPanel({
|
||||
setMessages([])
|
||||
clearDiagram()
|
||||
setDiagramHistory([])
|
||||
setValidationStates({}) // Clear validation states to prevent memory leak
|
||||
handleFileChange([]) // Use handleFileChange to also clear pdfData
|
||||
setUrlData(new Map())
|
||||
const newSessionId = `session-${Date.now()}-${Math.random()
|
||||
|
||||
@@ -11,7 +11,7 @@ import {
|
||||
X,
|
||||
} from "lucide-react"
|
||||
import Image from "next/image"
|
||||
import { useEffect, useState } from "react"
|
||||
import { useState } from "react"
|
||||
import type { ValidationResult } from "@/lib/diagram-validator"
|
||||
|
||||
export type ValidationStatus =
|
||||
@@ -254,8 +254,7 @@ export function ValidationCard({
|
||||
)}
|
||||
|
||||
{/* Valid result message */}
|
||||
{state.result &&
|
||||
state.result.valid &&
|
||||
{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
|
||||
|
||||
@@ -211,7 +211,8 @@ ${finalXml}
|
||||
// Notify UI that we're starting capture
|
||||
updateValidationState(toolCall.toolCallId, "capturing")
|
||||
|
||||
// Small delay to ensure diagram is rendered before capture
|
||||
// Small delay (100ms) to allow diagram rendering to complete before capture.
|
||||
// This is a best-effort heuristic and may need adjustment for complex diagrams or slower devices.
|
||||
await new Promise((resolve) => setTimeout(resolve, 100))
|
||||
|
||||
capturedPngData = await captureValidationPng()
|
||||
@@ -240,7 +241,6 @@ ${finalXml}
|
||||
|
||||
const result = await validateRenderedDiagram(
|
||||
capturedPngData,
|
||||
finalXml,
|
||||
sessionId,
|
||||
)
|
||||
|
||||
|
||||
@@ -20,13 +20,11 @@ export interface ValidationResult {
|
||||
* 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 {
|
||||
@@ -37,7 +35,6 @@ export async function validateRenderedDiagram(
|
||||
},
|
||||
body: JSON.stringify({
|
||||
imageData: pngDataUrl,
|
||||
xml,
|
||||
sessionId,
|
||||
}),
|
||||
})
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
/**
|
||||
* VLM system prompt and response parsing for diagram validation.
|
||||
* VLM system prompt for diagram validation.
|
||||
* Note: Response parsing is now handled via AI SDK's structured outputs (generateObject with schema).
|
||||
*/
|
||||
|
||||
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:
|
||||
@@ -14,109 +13,10 @@ Evaluate the diagram for the following issues:
|
||||
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: [],
|
||||
}
|
||||
}
|
||||
}
|
||||
- If the diagram looks generally acceptable, set valid to true even with minor warnings`
|
||||
|
||||
119
tests/unit/diagram-validator.test.ts
Normal file
119
tests/unit/diagram-validator.test.ts
Normal file
@@ -0,0 +1,119 @@
|
||||
import { describe, expect, it } from "vitest"
|
||||
import {
|
||||
formatValidationFeedback,
|
||||
type ValidationResult,
|
||||
} from "@/lib/diagram-validator"
|
||||
|
||||
describe("formatValidationFeedback", () => {
|
||||
it("formats result with critical issues", () => {
|
||||
const result: ValidationResult = {
|
||||
valid: false,
|
||||
issues: [
|
||||
{
|
||||
type: "overlap",
|
||||
severity: "critical",
|
||||
description: "Box A overlaps with Box B",
|
||||
},
|
||||
],
|
||||
suggestions: ["Move Box A to the left"],
|
||||
}
|
||||
|
||||
const feedback = formatValidationFeedback(result)
|
||||
|
||||
expect(feedback).toContain("DIAGRAM VISUAL VALIDATION FAILED")
|
||||
expect(feedback).toContain("Critical Issues (must fix):")
|
||||
expect(feedback).toContain("[overlap] Box A overlaps with Box B")
|
||||
expect(feedback).toContain("Suggestions to fix:")
|
||||
expect(feedback).toContain("Move Box A to the left")
|
||||
expect(feedback).toContain(
|
||||
"Please regenerate the diagram with corrected layout",
|
||||
)
|
||||
})
|
||||
|
||||
it("formats result with warnings only", () => {
|
||||
const result: ValidationResult = {
|
||||
valid: true,
|
||||
issues: [
|
||||
{
|
||||
type: "text",
|
||||
severity: "warning",
|
||||
description: "Label text is small",
|
||||
},
|
||||
],
|
||||
suggestions: [],
|
||||
}
|
||||
|
||||
const feedback = formatValidationFeedback(result)
|
||||
|
||||
expect(feedback).toContain("Warnings:")
|
||||
expect(feedback).toContain("[text] Label text is small")
|
||||
expect(feedback).not.toContain("Critical Issues")
|
||||
})
|
||||
|
||||
it("formats result with both critical issues and warnings", () => {
|
||||
const result: ValidationResult = {
|
||||
valid: false,
|
||||
issues: [
|
||||
{
|
||||
type: "edge_routing",
|
||||
severity: "critical",
|
||||
description: "Edge crosses through node",
|
||||
},
|
||||
{
|
||||
type: "layout",
|
||||
severity: "warning",
|
||||
description: "Uneven spacing",
|
||||
},
|
||||
],
|
||||
suggestions: ["Reroute the edge", "Adjust spacing"],
|
||||
}
|
||||
|
||||
const feedback = formatValidationFeedback(result)
|
||||
|
||||
expect(feedback).toContain("Critical Issues (must fix):")
|
||||
expect(feedback).toContain("[edge_routing] Edge crosses through node")
|
||||
expect(feedback).toContain("Warnings:")
|
||||
expect(feedback).toContain("[layout] Uneven spacing")
|
||||
expect(feedback).toContain("Reroute the edge")
|
||||
expect(feedback).toContain("Adjust spacing")
|
||||
})
|
||||
|
||||
it("formats result with no issues", () => {
|
||||
const result: ValidationResult = {
|
||||
valid: true,
|
||||
issues: [],
|
||||
suggestions: [],
|
||||
}
|
||||
|
||||
const feedback = formatValidationFeedback(result)
|
||||
|
||||
expect(feedback).toContain("DIAGRAM VISUAL VALIDATION FAILED")
|
||||
expect(feedback).not.toContain("Critical Issues")
|
||||
expect(feedback).not.toContain("Warnings:")
|
||||
expect(feedback).not.toContain("Suggestions to fix:")
|
||||
})
|
||||
|
||||
it("formats result with multiple suggestions", () => {
|
||||
const result: ValidationResult = {
|
||||
valid: false,
|
||||
issues: [
|
||||
{
|
||||
type: "rendering",
|
||||
severity: "critical",
|
||||
description: "Missing element",
|
||||
},
|
||||
],
|
||||
suggestions: [
|
||||
"Check the XML syntax",
|
||||
"Ensure all elements are defined",
|
||||
"Verify parent-child relationships",
|
||||
],
|
||||
}
|
||||
|
||||
const feedback = formatValidationFeedback(result)
|
||||
|
||||
expect(feedback).toContain("Check the XML syntax")
|
||||
expect(feedback).toContain("Ensure all elements are defined")
|
||||
expect(feedback).toContain("Verify parent-child relationships")
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user