feat: implement Story 3.2 — AI diagram generation from natural language

Add complete AI-powered diagram generation pipeline: natural language input
→ type inference → graph patch generation → validated canvas render with
ELK.js layout animation. Includes adversarial code review fixes for
diagramType enum validation, duplicate ID detection, tool part history
preservation, PATCH error handling, and graphData structural validation.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Alejandro Gutiérrez
2026-02-28 13:34:46 +00:00
parent 26215d9060
commit 6dcb4dcd6f
13 changed files with 1577 additions and 44 deletions

View File

@@ -3,7 +3,9 @@ import {
createUIMessageStream,
createUIMessageStreamResponse,
smoothStream,
stepCountIs,
streamText,
tool,
} from "ai";
import { eq } from "@turbostarter/db";
@@ -13,6 +15,13 @@ import { db } from "@turbostarter/db/server";
import { modelStrategies } from "../chat/strategies";
import { Model, Role } from "../chat/types";
import {
graphPatchSchema,
validateEdgeReferences,
validateEdgeTypes,
validateNodeTypes,
validateUniqueIds,
} from "./mutation-schema";
import { buildCopilotSystemPrompt } from "./system-prompt";
import type { CopilotMessagePayload } from "./schema";
@@ -25,6 +34,24 @@ import type {
const DEFAULT_MODEL = Model.CLAUDE_4_SONNET;
const generateDiagramTool = tool({
description:
"Generate or modify a diagram with the given graph data. Call this when the user asks to create, restructure, or modify a diagram.",
inputSchema: graphPatchSchema,
execute: async (patch) => {
const errors = [
...validateUniqueIds(patch),
...validateNodeTypes(patch),
...validateEdgeTypes(patch),
...validateEdgeReferences(patch),
];
if (errors.length > 0) {
return { success: false as const, errors };
}
return { success: true as const, data: patch };
},
});
const createCopilotChat = async (data: InsertChat) =>
db
.insert(chat)
@@ -60,6 +87,15 @@ const toChatMessage = (msg: Awaited<ReturnType<typeof getChatMessages>>[number])
role: msg.role as "user" | "assistant",
parts: msg.part.map((p) => {
const details = p.details as Record<string, unknown>;
// Preserve tool invocation parts so the AI retains generation context
if (p.type.startsWith("tool-")) {
return {
type: p.type as "text",
...(details as Record<string, unknown>),
};
}
return {
type: "text" as const,
text: (details.text as string) ?? "",
@@ -76,6 +112,7 @@ export const streamCopilot = async ({
chatId,
diagramId,
diagramType,
graphContext,
userId,
signal,
...msg
@@ -102,7 +139,9 @@ export const streamCopilot = async ({
})),
);
const systemPrompt = buildCopilotSystemPrompt(diagramType as DiagramType);
const systemPrompt = buildCopilotSystemPrompt(diagramType as DiagramType, {
graphContext,
});
const stream = createUIMessageStream({
execute: ({ writer }) => {
@@ -117,6 +156,8 @@ export const streamCopilot = async ({
},
]),
system: systemPrompt,
tools: { generateDiagram: generateDiagramTool },
stopWhen: stepCountIs(2),
abortSignal: signal,
experimental_transform: smoothStream({
chunking: "word",

View File

@@ -0,0 +1,464 @@
import { describe, expect, it } from "vitest";
import {
graphPatchSchema,
validateEdgeReferences,
validateEdgeTypes,
validateGraphPatch,
validateNodeTypes,
validateUniqueIds,
} from "./mutation-schema";
import type { GraphPatch } from "./mutation-schema";
function createTestPatch(overrides: Partial<GraphPatch> = {}): GraphPatch {
return {
meta: { diagramType: "flowchart", title: "Test Diagram" },
nodes: [
{ id: "n1", type: "terminal", label: "Start" },
{ id: "n2", type: "process", label: "Do something" },
{ id: "n3", type: "terminal", label: "End" },
],
edges: [
{ id: "e1", from: "n1", to: "n2" },
{ id: "e2", from: "n2", to: "n3" },
],
...overrides,
};
}
describe("graphPatchSchema", () => {
it("should accept a valid flowchart patch", () => {
const patch = createTestPatch();
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
it("should accept a valid BPMN patch with pools", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "BPMN Process" },
nodes: [
{ id: "n1", type: "start-event", label: "Start", lane: "lane1" },
{ id: "n2", type: "activity", label: "Task", lane: "lane1" },
],
edges: [{ id: "e1", from: "n1", to: "n2", type: "sequence" }],
pools: [
{
id: "pool1",
label: "Customer",
lanes: [{ id: "lane1", label: "Web" }],
},
],
});
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
it("should accept a valid E-R patch with columns", () => {
const patch = createTestPatch({
meta: { diagramType: "er", title: "Schema" },
nodes: [
{
id: "n1",
type: "entity",
label: "User",
columns: [
{ name: "id", type: "uuid", isPrimaryKey: true },
{ name: "email", type: "varchar", isUnique: true },
],
},
],
edges: [],
});
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
it("should accept a valid sequence patch with lifeline", () => {
const patch = createTestPatch({
meta: { diagramType: "sequence", title: "Flow" },
nodes: [
{ id: "n1", type: "participant", label: "Client", lifeline: true },
{ id: "n2", type: "participant", label: "Server", lifeline: true },
],
edges: [{ id: "e1", from: "n1", to: "n2", type: "sync", label: "request" }],
});
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
it("should reject patch without meta", () => {
const result = graphPatchSchema.safeParse({
nodes: [{ id: "n1", type: "process", label: "A" }],
edges: [],
});
expect(result.success).toBe(false);
});
it("should reject patch without nodes", () => {
const result = graphPatchSchema.safeParse({
meta: { diagramType: "flowchart", title: "Test" },
edges: [],
});
expect(result.success).toBe(false);
});
it("should reject node without required fields", () => {
const result = graphPatchSchema.safeParse({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [{ id: "n1" }],
edges: [],
});
expect(result.success).toBe(false);
});
it("should reject edge without from/to", () => {
const result = graphPatchSchema.safeParse({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [{ id: "n1", type: "process", label: "A" }],
edges: [{ id: "e1", from: "n1" }],
});
expect(result.success).toBe(false);
});
it("should reject more than 200 nodes", () => {
const nodes = Array.from({ length: 201 }, (_, i) => ({
id: `n${i}`,
type: "process",
label: `Node ${i}`,
}));
const result = graphPatchSchema.safeParse({
meta: { diagramType: "flowchart", title: "Too Many" },
nodes,
edges: [],
});
expect(result.success).toBe(false);
});
it("should accept optional layoutDirection", () => {
const patch = createTestPatch({
meta: { diagramType: "flowchart", title: "Test", layoutDirection: "RIGHT" },
});
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
it("should accept groups", () => {
const patch = createTestPatch({
groups: [{ id: "g1", label: "Group A", color: "#ff0000" }],
nodes: [
{ id: "n1", type: "process", label: "Task", group: "g1" },
{ id: "n2", type: "terminal", label: "End" },
],
edges: [{ id: "e1", from: "n1", to: "n2" }],
});
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
});
describe("validateNodeTypes", () => {
it("should accept valid BPMN node types", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [
{ id: "n1", type: "activity", label: "Task" },
{ id: "n2", type: "gateway-exclusive", label: "Decision" },
{ id: "n3", type: "start-event", label: "Start" },
],
});
expect(validateNodeTypes(patch)).toEqual([]);
});
it("should reject invalid BPMN node types", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [{ id: "n1", type: "invalid-type", label: "Bad" }],
});
const errors = validateNodeTypes(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("invalid-type");
});
it("should strip prefix before validation", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [{ id: "n1", type: "bpmn:activity", label: "Task" }],
});
expect(validateNodeTypes(patch)).toEqual([]);
});
it("should accept valid architecture node types", () => {
const patch = createTestPatch({
meta: { diagramType: "architecture", title: "Test" },
nodes: [
{ id: "n1", type: "service", label: "API" },
{ id: "n2", type: "database", label: "DB" },
{ id: "n3", type: "queue", label: "Queue" },
{ id: "n4", type: "loadbalancer", label: "LB" },
{ id: "n5", type: "external", label: "Ext" },
],
});
expect(validateNodeTypes(patch)).toEqual([]);
});
it("should accept valid flowchart node types", () => {
const patch = createTestPatch({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [
{ id: "n1", type: "process", label: "Do" },
{ id: "n2", type: "decision", label: "?" },
{ id: "n3", type: "terminal", label: "End" },
{ id: "n4", type: "io", label: "Read" },
{ id: "n5", type: "subprocess", label: "Sub" },
],
});
expect(validateNodeTypes(patch)).toEqual([]);
});
it("should accept valid sequence node types", () => {
const patch = createTestPatch({
meta: { diagramType: "sequence", title: "Test" },
nodes: [
{ id: "n1", type: "participant", label: "Client" },
{ id: "n2", type: "fragment", label: "alt" },
],
});
expect(validateNodeTypes(patch)).toEqual([]);
});
});
describe("validateEdgeTypes", () => {
it("should accept valid BPMN edge types", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [
{ id: "n1", type: "activity", label: "A" },
{ id: "n2", type: "activity", label: "B" },
],
edges: [
{ id: "e1", from: "n1", to: "n2", type: "sequence" },
{ id: "e2", from: "n1", to: "n2", type: "message" },
],
});
expect(validateEdgeTypes(patch)).toEqual([]);
});
it("should reject invalid BPMN edge types", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [
{ id: "n1", type: "activity", label: "A" },
{ id: "n2", type: "activity", label: "B" },
],
edges: [{ id: "e1", from: "n1", to: "n2", type: "invalid" }],
});
const errors = validateEdgeTypes(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("invalid");
});
it("should accept valid sequence edge types", () => {
const patch = createTestPatch({
meta: { diagramType: "sequence", title: "Test" },
nodes: [
{ id: "n1", type: "participant", label: "A" },
{ id: "n2", type: "participant", label: "B" },
],
edges: [
{ id: "e1", from: "n1", to: "n2", type: "sync" },
{ id: "e2", from: "n2", to: "n1", type: "return" },
{ id: "e3", from: "n1", to: "n2", type: "async" },
],
});
expect(validateEdgeTypes(patch)).toEqual([]);
});
it("should skip validation for types with no edge type constraints", () => {
const patch = createTestPatch({
meta: { diagramType: "er", title: "Test" },
nodes: [{ id: "n1", type: "entity", label: "A" }],
edges: [{ id: "e1", from: "n1", to: "n1", type: "anything" }],
});
expect(validateEdgeTypes(patch)).toEqual([]);
});
it("should allow edges without type for constrained diagram types", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [
{ id: "n1", type: "activity", label: "A" },
{ id: "n2", type: "activity", label: "B" },
],
edges: [{ id: "e1", from: "n1", to: "n2" }],
});
expect(validateEdgeTypes(patch)).toEqual([]);
});
});
describe("validateEdgeReferences", () => {
it("should accept edges referencing existing nodes", () => {
const patch = createTestPatch();
expect(validateEdgeReferences(patch)).toEqual([]);
});
it("should reject edges referencing non-existent source", () => {
const patch = createTestPatch({
edges: [{ id: "e1", from: "missing", to: "n1" }],
});
const errors = validateEdgeReferences(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("missing");
});
it("should reject edges referencing non-existent target", () => {
const patch = createTestPatch({
edges: [{ id: "e1", from: "n1", to: "missing" }],
});
const errors = validateEdgeReferences(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("missing");
});
});
describe("validateGraphPatch", () => {
it("should return success for valid patch", () => {
const result = validateGraphPatch(createTestPatch());
expect(result.success).toBe(true);
if (result.success) {
expect(result.data.nodes).toHaveLength(3);
expect(result.data.edges).toHaveLength(2);
}
});
it("should return errors for invalid schema", () => {
const result = validateGraphPatch({ invalid: true });
expect(result.success).toBe(false);
if (!result.success) {
expect(result.errors.length).toBeGreaterThan(0);
}
});
it("should return errors for invalid node types", () => {
const result = validateGraphPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [{ id: "n1", type: "invalid", label: "Bad" }],
edges: [],
});
expect(result.success).toBe(false);
});
it("should return errors for broken edge references", () => {
const result = validateGraphPatch({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [{ id: "n1", type: "process", label: "A" }],
edges: [{ id: "e1", from: "n1", to: "missing" }],
});
expect(result.success).toBe(false);
});
it("should combine node type and edge reference errors", () => {
const result = validateGraphPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [{ id: "n1", type: "invalid", label: "Bad" }],
edges: [{ id: "e1", from: "n1", to: "missing" }],
});
expect(result.success).toBe(false);
if (!result.success) {
expect(result.errors.length).toBe(2);
}
});
it("should reject invalid diagramType", () => {
const result = validateGraphPatch({
meta: { diagramType: "unknown", title: "Test" },
nodes: [{ id: "n1", type: "process", label: "A" }],
edges: [],
});
expect(result.success).toBe(false);
});
it("should reject empty title", () => {
const result = validateGraphPatch({
meta: { diagramType: "flowchart", title: "" },
nodes: [{ id: "n1", type: "process", label: "A" }],
edges: [],
});
expect(result.success).toBe(false);
});
it("should reject duplicate node IDs", () => {
const result = validateGraphPatch({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [
{ id: "n1", type: "process", label: "A" },
{ id: "n1", type: "decision", label: "B" },
],
edges: [],
});
expect(result.success).toBe(false);
if (!result.success) {
expect(result.errors[0]).toContain("Duplicate node ID");
}
});
it("should reject duplicate edge IDs", () => {
const result = validateGraphPatch({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [
{ id: "n1", type: "process", label: "A" },
{ id: "n2", type: "process", label: "B" },
],
edges: [
{ id: "e1", from: "n1", to: "n2" },
{ id: "e1", from: "n2", to: "n1" },
],
});
expect(result.success).toBe(false);
if (!result.success) {
expect(result.errors[0]).toContain("Duplicate edge ID");
}
});
});
describe("validateUniqueIds", () => {
it("should accept unique IDs", () => {
const patch = createTestPatch();
expect(validateUniqueIds(patch)).toEqual([]);
});
it("should detect duplicate node IDs", () => {
const patch = createTestPatch({
nodes: [
{ id: "n1", type: "process", label: "A" },
{ id: "n1", type: "process", label: "B" },
],
});
const errors = validateUniqueIds(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("n1");
});
it("should detect duplicate edge IDs", () => {
const patch = createTestPatch({
edges: [
{ id: "e1", from: "n1", to: "n2" },
{ id: "e1", from: "n2", to: "n3" },
],
});
const errors = validateUniqueIds(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("e1");
});
it("should allow same ID across nodes and edges", () => {
const patch = createTestPatch({
nodes: [
{ id: "x1", type: "process", label: "A" },
{ id: "x2", type: "process", label: "B" },
],
edges: [{ id: "x1", from: "x1", to: "x2" }],
});
expect(validateUniqueIds(patch)).toEqual([]);
});
});

View File

@@ -0,0 +1,229 @@
import * as z from "zod";
const VALID_DIAGRAM_TYPES = [
"bpmn",
"er",
"orgchart",
"architecture",
"sequence",
"flowchart",
] as const;
const VALID_NODE_TYPES: Record<string, string[]> = {
bpmn: [
"activity",
"subprocess",
"start-event",
"end-event",
"event-timer",
"event-message",
"gateway-exclusive",
"gateway-parallel",
"gateway-inclusive",
"data-object",
"annotation",
],
er: ["entity"],
orgchart: ["person"],
architecture: ["service", "database", "queue", "loadbalancer", "external"],
sequence: ["participant", "fragment"],
flowchart: ["process", "decision", "terminal", "io", "subprocess"],
};
const VALID_EDGE_TYPES: Record<string, string[] | null> = {
bpmn: ["sequence", "message", "association"],
er: null,
orgchart: null,
architecture: null,
sequence: ["sync", "async", "return"],
flowchart: null,
};
const MAX_NODES = 200;
const columnSchema = z.object({
name: z.string(),
type: z.string(),
isPrimaryKey: z.boolean().optional(),
isForeignKey: z.boolean().optional(),
isNullable: z.boolean().optional(),
isUnique: z.boolean().optional(),
references: z.string().optional(),
});
const nodeSchema = z.object({
id: z.string(),
type: z.string(),
label: z.string(),
tag: z.string().optional(),
icon: z.string().optional(),
color: z.string().optional(),
w: z.number().optional(),
position: z
.object({ x: z.number(), y: z.number() })
.optional(),
lane: z.string().optional(),
group: z.string().optional(),
columns: z.array(columnSchema).optional(),
lifeline: z.boolean().optional(),
parentId: z.string().optional(),
});
const edgeSchema = z.object({
id: z.string(),
from: z.string(),
to: z.string(),
label: z.string().optional(),
color: z.string().optional(),
type: z.string().optional(),
cardinality: z.string().optional(),
});
const metaSchema = z.object({
diagramType: z.enum(VALID_DIAGRAM_TYPES),
title: z.string().min(1),
description: z.string().optional(),
version: z.string().optional(),
layoutDirection: z.enum(["DOWN", "RIGHT", "LEFT", "UP"]).optional(),
edgeRouting: z.enum(["ORTHOGONAL", "SPLINES", "POLYLINE"]).optional(),
});
const poolSchema = z.object({
id: z.string(),
label: z.string(),
lanes: z.array(
z.object({
id: z.string(),
label: z.string(),
}),
),
});
const groupSchema = z.object({
id: z.string(),
label: z.string(),
color: z.string().optional(),
});
export const graphPatchSchema = z.object({
meta: metaSchema,
nodes: z.array(nodeSchema).max(MAX_NODES),
edges: z.array(edgeSchema),
pools: z.array(poolSchema).optional(),
groups: z.array(groupSchema).optional(),
});
export type GraphPatch = z.infer<typeof graphPatchSchema>;
/**
* Validate that node types match the diagram type's allowed types.
* Returns an array of error messages (empty if valid).
*/
export function validateNodeTypes(patch: GraphPatch): string[] {
const diagramType = patch.meta.diagramType;
const validTypes = VALID_NODE_TYPES[diagramType];
if (!validTypes) return [];
const errors: string[] = [];
for (const node of patch.nodes) {
const bare = node.type.includes(":") ? node.type.split(":")[1]! : node.type;
if (!validTypes.includes(bare)) {
errors.push(
`Invalid node type "${node.type}" for ${diagramType}. Valid types: ${validTypes.join(", ")}`,
);
}
}
return errors;
}
/**
* Validate that edge types match the diagram type's allowed types.
* Returns an array of error messages (empty if valid).
*/
export function validateEdgeTypes(patch: GraphPatch): string[] {
const diagramType = patch.meta.diagramType;
const validTypes = VALID_EDGE_TYPES[diagramType];
if (!validTypes) return [];
const errors: string[] = [];
for (const edge of patch.edges) {
if (edge.type && !validTypes.includes(edge.type)) {
errors.push(
`Invalid edge type "${edge.type}" for ${diagramType}. Valid types: ${validTypes.join(", ")}`,
);
}
}
return errors;
}
/**
* Validate that all edge from/to references point to existing node IDs.
*/
export function validateEdgeReferences(patch: GraphPatch): string[] {
const nodeIds = new Set(patch.nodes.map((n) => n.id));
const errors: string[] = [];
for (const edge of patch.edges) {
if (!nodeIds.has(edge.from)) {
errors.push(`Edge "${edge.id}" references non-existent source node "${edge.from}"`);
}
if (!nodeIds.has(edge.to)) {
errors.push(`Edge "${edge.id}" references non-existent target node "${edge.to}"`);
}
}
return errors;
}
/**
* Validate that all node and edge IDs are unique.
*/
export function validateUniqueIds(patch: GraphPatch): string[] {
const errors: string[] = [];
const nodeIds = new Set<string>();
for (const node of patch.nodes) {
if (nodeIds.has(node.id)) {
errors.push(`Duplicate node ID "${node.id}"`);
}
nodeIds.add(node.id);
}
const edgeIds = new Set<string>();
for (const edge of patch.edges) {
if (edgeIds.has(edge.id)) {
errors.push(`Duplicate edge ID "${edge.id}"`);
}
edgeIds.add(edge.id);
}
return errors;
}
/**
* Full validation: schema + type-specific + referential integrity + unique IDs.
* Returns { success: true, data } or { success: false, errors }.
*/
export function validateGraphPatch(
input: unknown,
): { success: true; data: GraphPatch } | { success: false; errors: string[] } {
const result = graphPatchSchema.safeParse(input);
if (!result.success) {
return {
success: false,
errors: result.error.issues.map((i) => `${i.path.join(".")}: ${i.message}`),
};
}
const patch = result.data;
const typeErrors = [
...validateUniqueIds(patch),
...validateNodeTypes(patch),
...validateEdgeTypes(patch),
...validateEdgeReferences(patch),
];
if (typeErrors.length > 0) {
return { success: false, errors: typeErrors };
}
return { success: true, data: patch };
}

View File

@@ -25,14 +25,15 @@ describe("buildCopilotSystemPrompt", () => {
const prompt = buildCopilotSystemPrompt(type);
expect(prompt).toContain(`Type: ${type.toUpperCase()}`);
expect(prompt).toContain("domaingraph AI copilot");
expect(prompt).toContain("CHAT-ONLY mode");
}
});
it("should mention chat-only constraint", () => {
const prompt = buildCopilotSystemPrompt("er");
expect(prompt).toContain("CHAT-ONLY mode");
expect(prompt).toContain("cannot modify the diagram directly");
it("should include generation instructions with generateDiagram tool", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("generateDiagram");
expect(prompt).toContain("Generate complete diagrams");
expect(prompt).not.toContain("CHAT-ONLY");
expect(prompt).not.toContain("cannot modify the diagram directly");
});
it("should include today's date", () => {
@@ -52,4 +53,124 @@ describe("buildCopilotSystemPrompt", () => {
expect(prompt).toContain("services");
expect(prompt).toContain("databases");
});
it("should include BPMN-specific node types", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("activity");
expect(prompt).toContain("gateway-exclusive");
expect(prompt).toContain("gateway-parallel");
expect(prompt).toContain("gateway-inclusive");
expect(prompt).toContain("start-event");
expect(prompt).toContain("end-event");
expect(prompt).toContain("event-timer");
expect(prompt).toContain("event-message");
expect(prompt).toContain("subprocess");
expect(prompt).toContain("data-object");
expect(prompt).toContain("annotation");
});
it("should include BPMN-specific edge types", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("sequence");
expect(prompt).toContain("message");
expect(prompt).toContain("association");
});
it("should include BPMN pools and lanes guidance", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("pools");
expect(prompt).toContain("lanes");
});
it("should include node types for each diagram type", () => {
const typeNodeMap: Record<DiagramType, string[]> = {
bpmn: ["activity", "gateway-exclusive", "start-event", "end-event"],
er: ["entity"],
orgchart: ["person"],
architecture: ["service", "database", "queue", "loadbalancer", "external"],
sequence: ["participant", "fragment"],
flowchart: ["process", "decision", "terminal", "io"],
};
for (const [type, nodes] of Object.entries(typeNodeMap)) {
const prompt = buildCopilotSystemPrompt(type as DiagramType);
for (const node of nodes) {
expect(prompt).toContain(node);
}
}
});
it("should include sequence-specific edge types", () => {
const prompt = buildCopilotSystemPrompt("sequence");
expect(prompt).toContain("sync");
expect(prompt).toContain("async");
expect(prompt).toContain("return");
});
it("should include ER column format with key fields", () => {
const prompt = buildCopilotSystemPrompt("er");
expect(prompt).toContain("columns");
expect(prompt).toContain("isPrimaryKey");
expect(prompt).toContain("isForeignKey");
expect(prompt).toContain("isNullable");
});
it("should include type inference instructions for all types", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("infer");
expect(prompt).toContain("bpmn");
expect(prompt).toContain("er");
expect(prompt).toContain("orgchart");
expect(prompt).toContain("architecture");
expect(prompt).toContain("sequence");
expect(prompt).toContain("flowchart");
});
it("should include 200-node generation cap", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("200");
});
it("should include graph context when provided", () => {
const graphContext = '{"nodes":[],"edges":[]}';
const prompt = buildCopilotSystemPrompt("bpmn", { graphContext });
expect(prompt).toContain(graphContext);
});
it("should show empty canvas when no graph context", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("Empty canvas");
});
it("should include a JSON example for each diagram type", () => {
const types: DiagramType[] = [
"bpmn",
"er",
"orgchart",
"architecture",
"sequence",
"flowchart",
];
for (const type of types) {
const prompt = buildCopilotSystemPrompt(type);
expect(prompt).toContain("```json");
expect(prompt).toContain(`"diagramType":"${type}"`);
}
});
it("should include GraphData format specification", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("meta");
expect(prompt).toContain("nodes");
expect(prompt).toContain("edges");
expect(prompt).toContain("diagramType");
expect(prompt).toContain("layoutDirection");
});
it("should include ID generation convention", () => {
const prompt = buildCopilotSystemPrompt("flowchart");
expect(prompt).toContain('"n1"');
expect(prompt).toContain('"e1"');
});
});

View File

@@ -9,25 +9,137 @@ const DIAGRAM_DESCRIPTIONS: Record<DiagramType, string> = {
flowchart: "Flowchart — decision flows with processes, decisions, terminals, I/O, and subprocesses",
};
export function buildCopilotSystemPrompt(diagramType: DiagramType): string {
const description = DIAGRAM_DESCRIPTIONS[diagramType];
const NODE_TYPE_REFERENCE: Record<DiagramType, string> = {
bpmn: `- activity: Task or action step
- subprocess: Collapsed subprocess
- start-event: Process start
- end-event: Process end
- event-timer: Timer event
- event-message: Message event
- gateway-exclusive: XOR decision (one path)
- gateway-parallel: AND split/join (all paths)
- gateway-inclusive: OR split/join (one or more)
- data-object: Data artifact
- annotation: Text note`,
er: `- entity: Database table (use columns[] for attributes)`,
orgchart: `- person: Team member or role`,
architecture: `- service: Application or microservice
- database: Data store
- queue: Message queue or event bus
- loadbalancer: Load balancer or API gateway
- external: External system`,
sequence: `- participant: Actor or system (set lifeline: true)
- fragment: Combined fragment (loop, alt, opt)`,
flowchart: `- process: Action or operation
- decision: Yes/no branch
- terminal: Start or end point
- io: Input/output step
- subprocess: Nested process`,
};
return `You are the domaingraph AI copilot — a diagram design assistant.
const EDGE_TYPE_REFERENCE: Record<DiagramType, string> = {
bpmn: `- sequence (default): Flow between activities
- message: Cross-pool message flow
- association: Link to annotations/data`,
er: `- (default): Relationship. Set cardinality: "1:1", "1:N", "N:1", or "M:N"`,
orgchart: `- (default): Reports-to relationship`,
architecture: `- (default): Connection. Use label for protocol/description`,
sequence: `- sync (default): Synchronous call (solid arrow)
- async: Asynchronous message (open arrow)
- return: Response (dashed arrow)`,
flowchart: `- (default): Flow. Use label for conditions ("Yes", "No")`,
};
const TYPE_INFERENCE_RULES = `If the diagram type is not established, infer from the user's description:
- Business processes, workflows, approvals, order handling → bpmn
- Database schemas, tables, entities, data models → er
- Team structures, org hierarchies, reporting lines → orgchart
- System design, microservices, infrastructure, APIs → architecture
- Interactions between actors over time, API calls, request/response → sequence
- Decision logic, algorithms, if/else flows → flowchart`;
const DIAGRAM_EXAMPLES: Record<DiagramType, string> = {
bpmn: `{"meta":{"diagramType":"bpmn","title":"Order Process","layoutDirection":"RIGHT"},"nodes":[{"id":"n1","type":"start-event","label":"Order received"},{"id":"n2","type":"activity","label":"Validate order"},{"id":"n3","type":"gateway-exclusive","label":"Valid?"},{"id":"n4","type":"activity","label":"Process payment"},{"id":"n5","type":"end-event","label":"Complete"}],"edges":[{"id":"e1","from":"n1","to":"n2"},{"id":"e2","from":"n2","to":"n3"},{"id":"e3","from":"n3","to":"n4","label":"Yes"},{"id":"e4","from":"n4","to":"n5"}]}`,
er: `{"meta":{"diagramType":"er","title":"Blog Schema"},"nodes":[{"id":"n1","type":"entity","label":"User","columns":[{"name":"id","type":"uuid","isPrimaryKey":true},{"name":"email","type":"varchar","isUnique":true}]},{"id":"n2","type":"entity","label":"Post","columns":[{"name":"id","type":"uuid","isPrimaryKey":true},{"name":"author_id","type":"uuid","isForeignKey":true,"references":"User.id"},{"name":"title","type":"varchar"}]}],"edges":[{"id":"e1","from":"n1","to":"n2","cardinality":"1:N","label":"writes"}]}`,
orgchart: `{"meta":{"diagramType":"orgchart","title":"Engineering Team","layoutDirection":"DOWN"},"nodes":[{"id":"n1","type":"person","label":"VP Engineering"},{"id":"n2","type":"person","label":"Frontend Lead"},{"id":"n3","type":"person","label":"Backend Lead"}],"edges":[{"id":"e1","from":"n1","to":"n2"},{"id":"e2","from":"n1","to":"n3"}]}`,
architecture: `{"meta":{"diagramType":"architecture","title":"Web Stack"},"nodes":[{"id":"n1","type":"loadbalancer","label":"API Gateway"},{"id":"n2","type":"service","label":"Auth Service"},{"id":"n3","type":"service","label":"User Service"},{"id":"n4","type":"database","label":"PostgreSQL"}],"edges":[{"id":"e1","from":"n1","to":"n2","label":"REST"},{"id":"e2","from":"n1","to":"n3","label":"REST"},{"id":"e3","from":"n2","to":"n4"},{"id":"e4","from":"n3","to":"n4"}]}`,
sequence: `{"meta":{"diagramType":"sequence","title":"Login Flow"},"nodes":[{"id":"n1","type":"participant","label":"Browser","lifeline":true},{"id":"n2","type":"participant","label":"Auth API","lifeline":true},{"id":"n3","type":"participant","label":"Database","lifeline":true}],"edges":[{"id":"e1","from":"n1","to":"n2","type":"sync","label":"POST /login"},{"id":"e2","from":"n2","to":"n3","type":"sync","label":"SELECT user"},{"id":"e3","from":"n3","to":"n2","type":"return","label":"user record"},{"id":"e4","from":"n2","to":"n1","type":"return","label":"JWT token"}]}`,
flowchart: `{"meta":{"diagramType":"flowchart","title":"Validation Flow","layoutDirection":"DOWN"},"nodes":[{"id":"n1","type":"terminal","label":"Start"},{"id":"n2","type":"io","label":"Read input"},{"id":"n3","type":"decision","label":"Valid?"},{"id":"n4","type":"process","label":"Process data"},{"id":"n5","type":"terminal","label":"End"}],"edges":[{"id":"e1","from":"n1","to":"n2"},{"id":"e2","from":"n2","to":"n3"},{"id":"e3","from":"n3","to":"n4","label":"Yes"},{"id":"e4","from":"n3","to":"n2","label":"No"},{"id":"e5","from":"n4","to":"n5"}]}`,
};
export function buildCopilotSystemPrompt(
diagramType: DiagramType,
options?: { graphContext?: string },
): string {
const description = DIAGRAM_DESCRIPTIONS[diagramType];
const graphContext = options?.graphContext;
const date = new Date().toLocaleDateString("en-US", {
year: "numeric",
month: "short",
day: "2-digit",
weekday: "short",
});
const typeSpecificSections: string[] = [];
if (diagramType === "er") {
typeSpecificSections.push(`### E-R column format
Each entity node has a \`columns\` array: [{ name, type, isPrimaryKey?, isForeignKey?, isNullable?, isUnique?, references? }]`);
}
if (diagramType === "bpmn") {
typeSpecificSections.push(`### BPMN pools and lanes
Use \`pools\` to define swim lanes. Assign nodes to lanes via the \`lane\` field matching a lane ID.`);
}
return `You are the domaingraph AI copilot — a diagram design assistant that can discuss, generate, and modify diagrams.
## Current diagram
Type: ${diagramType.toUpperCase()}${description}
## Your role
- Help users think through their diagram design
- Explain diagram concepts and best practices for ${diagramType.toUpperCase()} diagrams
- Suggest improvements, missing elements, or structural changes
- Answer questions about the current diagram or diagram type
- Keep responses concise and diagram-focused
## Capabilities
- Answer questions about diagram design and best practices
- Generate complete diagrams from natural language descriptions
- Modify existing diagrams based on user requests
- Analyze the current diagram and suggest improvements
## Important constraints
- You are in CHAT-ONLY mode: you can discuss and advise, but you cannot modify the diagram directly yet
- When users ask you to add or change elements, explain what you would do and tell them this capability is coming soon
- Use markdown formatting for clarity (bold, lists, code blocks)
- Do not use h1 headings in responses
- Today's date is ${new Date().toLocaleDateString("en-US", { year: "numeric", month: "short", day: "2-digit", weekday: "short" })}`;
## Generating and modifying diagrams
When the user describes what they want to build, asks you to create a diagram, or requests changes, use the \`generateDiagram\` tool.
1. Briefly explain what you'll create or change (1-2 sentences)
2. Call the \`generateDiagram\` tool with the complete GraphData object
When modifying an existing diagram, include ALL nodes and edges in the output (not just changes). The tool output replaces the entire graph.
### GraphData format
The tool accepts:
- \`meta\`: { diagramType, title, description?, layoutDirection?: "DOWN"|"RIGHT"|"LEFT"|"UP" }
- \`nodes\`: [{ id, type, label, tag?, columns?, lane?, group?, lifeline? }]
- \`edges\`: [{ id, from, to, label?, type?, cardinality? }]
- \`pools\`: (BPMN only) [{ id, label, lanes: [{ id, label }] }]
- \`groups\`: [{ id, label, color? }]
### Node types for ${diagramType.toUpperCase()}
${NODE_TYPE_REFERENCE[diagramType]}
### Edge types for ${diagramType.toUpperCase()}
${EDGE_TYPE_REFERENCE[diagramType]}
${typeSpecificSections.length > 0 ? typeSpecificSections.join("\n\n") + "\n" : ""}## Example
\`\`\`json
${DIAGRAM_EXAMPLES[diagramType]}
\`\`\`
## Type inference
${TYPE_INFERENCE_RULES}
## Current graph state
${graphContext ? `The diagram currently contains:\n\`\`\`json\n${graphContext}\n\`\`\`` : "Empty canvas — no nodes or edges yet."}
## Constraints
- Maximum 200 nodes per generation
- Generate IDs as "n1", "n2", ... for nodes and "e1", "e2", ... for edges
- Keep text responses concise and diagram-focused
- Use markdown formatting (bold, lists, code blocks) — no h1 headings
- Today's date is ${date}`;
}

View File

@@ -3,3 +3,6 @@ import { diagramTypeEnum } from "@turbostarter/db/schema/diagram";
export type DiagramType = (typeof diagramTypeEnum.enumValues)[number];
export const DIAGRAM_TYPES = diagramTypeEnum.enumValues;
export type { GraphPatch } from "./mutation-schema";
export { graphPatchSchema, validateGraphPatch, validateUniqueIds } from "./mutation-schema";

View File

@@ -32,12 +32,22 @@ export const updateDiagramBodySchema = z
title: z.string().min(1).max(255).optional(),
projectId: z.string().nullable().optional(),
sortOrder: z.number().int().min(0).optional(),
graphData: z
.object({
meta: z.record(z.string(), z.unknown()).optional(),
nodes: z.array(z.record(z.string(), z.unknown())),
edges: z.array(z.record(z.string(), z.unknown())),
pools: z.array(z.record(z.string(), z.unknown())).optional(),
groups: z.array(z.record(z.string(), z.unknown())).optional(),
})
.optional(),
})
.refine(
(data) =>
data.title !== undefined ||
data.projectId !== undefined ||
data.sortOrder !== undefined,
data.sortOrder !== undefined ||
data.graphData !== undefined,
{ message: "At least one field must be provided" },
);