Files
next-ai-draw-io/app/api/validate-diagram/route.ts

137 lines
4.3 KiB
TypeScript
Raw Normal View History

Add VLM-based diagram validation (#602) * [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 * 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 * refactor(validation): use AI SDK experimental_useObject hook instead of raw fetch - Change API endpoint from generateObject to streamObject for useObject compatibility - Create useValidateDiagram hook using AI SDK's experimental_useObject for reactive validation - Update useDiagramToolHandlers to accept validation function as parameter - Update chat-panel to use new useValidateDiagram hook - Remove validateRenderedDiagram function from lib/diagram-validator.ts (now in hook) - Export ValidationResultSchema from API route for client-side use * fix(validation): extract schema to shared file for client/server compatibility Move ValidationResultSchema to lib/validation-schema.ts to avoid importing server-side modules (ai-providers) into client-side code. This fixes the Turbopack build error caused by the hook importing from the API route. * fix(validation): use 'Valid' instead of 'Complete' for validation success Change ValidationCard success label from 'Complete' to 'Valid' to avoid conflicting with ToolCallCard's 'Complete' badge in E2E tests. This fixes the diagram-generation E2E test that expects a specific count of 'Complete' badges. * fix(validation): add aria-hidden to icons to prevent duplicate ID warning * fix: improve VLM validation with bug fixes and i18n - Fix race condition in pendingValidationRef (reject previous pending validation) - Fix response format consistency (use streaming for all responses) - Remove dead code (unused lastRequestRef and ValidationRequest interface) - Consolidate duplicate types (re-export from validation-schema.ts) - Add 'success_with_warnings' status for valid diagrams with warnings - Fix tool card auto-collapse (only collapse once, respect user toggle) - Set VLM validation default to disabled - Add i18n support for diagram validation settings (en/zh/ja) - Mark feature as experimental in settings UI * fix: resolve TypeScript errors in electron-standalone - Add forwardRef support to ChatInput component with ChatInputRef type - Copy electron.d.ts to electron-standalone/electron folder - Exclude electron-standalone from root tsconfig type checking * fix: return empty string for valid result with no issues in formatValidationFeedback * feat(i18n): add validation strings for ValidationCard component - Add validation section to en.json, zh.json, ja.json dictionaries - Update ValidationCard to use useDictionary hook - Replace all hardcoded English strings with i18n keys --------- Co-authored-by: dayuan.jiang <jdy.toh@gmail.com>
2026-01-20 20:52:04 +09:00
/**
* API endpoint for VLM-based diagram validation.
* Accepts a PNG image and streams validation results using useObject-compatible format.
*/
import { streamObject } from "ai"
import { getValidationModel } from "@/lib/ai-providers"
import { VALIDATION_SYSTEM_PROMPT } from "@/lib/validation-prompts"
import {
type ValidationResult,
ValidationResultSchema,
} from "@/lib/validation-schema"
export const maxDuration = 30
interface ValidateDiagramRequest {
imageData: string // Base64 PNG data URL
sessionId?: string
}
// Default valid result for disabled/error cases
const DEFAULT_VALID_RESULT: ValidationResult = {
valid: true,
issues: [],
suggestions: [],
}
/**
* Create a streaming response for useObject compatibility.
* useObject expects text stream format, not plain JSON.
*/
function createStreamingResponse(result: ValidationResult): Response {
const encoder = new TextEncoder()
const stream = new ReadableStream({
start(controller) {
// Stream the JSON as text (useObject parses this)
controller.enqueue(encoder.encode(JSON.stringify(result)))
controller.close()
},
})
return new Response(stream, {
headers: { "Content-Type": "text/plain; charset=utf-8" },
})
}
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 createStreamingResponse(DEFAULT_VALID_RESULT)
}
const body: ValidateDiagramRequest = await req.json()
const { imageData, 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 createStreamingResponse(DEFAULT_VALID_RESULT)
}
// Parse timeout with validation (minimum 1000ms, default 10000ms)
const timeout =
Math.max(
1000,
parseInt(process.env.VALIDATION_TIMEOUT || "10000", 10),
) || 10000
// Stream the VLM response for useObject consumption
const result = streamObject({
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),
onFinish: ({ object }) => {
if (sessionId && object) {
console.log(
`[validate-diagram] Session ${sessionId}: valid=${object.valid}, issues=${object.issues?.length ?? 0}`,
)
}
},
})
return result.toTextStreamResponse()
} catch (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 createStreamingResponse(DEFAULT_VALID_RESULT)
}
}