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

refactor(validation): use AI SDK structured outputs and address review feedback

Browse Source 
- 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:
Jinze Yu
2026-01-19 00:47:49 +09:00
parent 07c9e7d758
commit 24270e2622
7 changed files with 195 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
app/api/validate-diagram/route.ts
View File

@@ -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({

2
components/chat-panel.tsx
View File

@@ -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()

5
components/chat/ValidationCard.tsx
View File

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

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

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

3
lib/diagram-validator.ts
View File

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

106
lib/validation-prompts.ts
View File

@@ -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
View 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")
})
})