@eventcatalogtest/studio/utils/prompt-generator.ts

A drag and drop UI for distributed systems that keeps your diagrams where they belong – in your repo

import { ReactFlowJsonObject, Node, Edge } from '@xyflow/react';

export interface GeneratePromptOptions {
  flowData: ReactFlowJsonObject;
  designName: string;
}

interface NodeData extends Record<string, unknown> {
  message?: { name?: string; summary?: string; version?: string; collection?: string };
  service?: { name?: string; summary?: string; version?: string };
  query?: { name?: string; summary?: string; version?: string };
  command?: { name?: string; summary?: string; version?: string };
  data?: { name?: string; summary?: string; version?: string };
  name?: string;
  summary?: string;
  version?: string;
}

interface EdgeData extends Record<string, unknown> {
  message?: { collection?: string };
  label?: string;
}

export const generateMermaidDiagram = (flowData: ReactFlowJsonObject): string => {
  if (!flowData?.nodes || flowData.nodes.length === 0) {
    return 'flowchart TD\n    A[No nodes available]';
  }

  const { nodes, edges } = flowData;

  // Helper function to get node display name
  const getNodeName = (node: Node<NodeData>) => {
    return node.data?.message?.name || node.data?.service?.name || node.data?.query?.name || 
           node.data?.command?.name || node.data?.data?.name || node.data?.name || 
           `${node.type || 'unknown'}_${node.id.slice(-4)}`;
  };

  // Helper function to get mermaid node shape based on type
  const getMermaidNodeShape = (type: string, name: string) => {
    switch (type) {
      case 'service':
        return `${name}[${name}]`;
      case 'event':
        return `${name}((${name}))`;
      case 'command':
        return `${name}>${name}]`;
      case 'query':
        return `${name}{${name}}`;
      case 'domain':
        return `${name}[/${name}/]`;
      case 'external-system':
        return `${name}[[${name}]]`;
      case 'data':
        return `${name}[(${name})]`;
      case 'channel':
        return `${name}[${name}]`;
      case 'view':
        return `${name}[${name}]`;
      case 'actor':
        return `${name}[${name}]`;
      default:
        return `${name}[${name}]`;
    }
  };

  // Build Mermaid diagram
  let mermaidDiagram = 'flowchart TD\n';
  
  // Add all nodes first
  const nodeNameMap: { [key: string]: string } = {};
  (nodes as Node<NodeData>[]).forEach(node => {
    const displayName = getNodeName(node);
    const safeName = displayName.replace(/[^a-zA-Z0-9]/g, '_');
    nodeNameMap[node.id] = safeName;
    mermaidDiagram += `    ${getMermaidNodeShape(node.type || 'unknown', safeName)}\n`;
  });

  // Add edges/connections
  (edges as Edge<EdgeData>[]).forEach(edge => {
    const sourceName = nodeNameMap[edge.source];
    const targetName = nodeNameMap[edge.target];
    const edgeData = edge.data as EdgeData | undefined;
    const label = edgeData?.message?.collection || edge.label || '';
    
    if (sourceName && targetName) {
      if (label) {
        mermaidDiagram += `    ${sourceName} -->|${label}| ${targetName}\n`;
      } else {
        mermaidDiagram += `    ${sourceName} --> ${targetName}\n`;
      }
    }
  });

  // Add styling for different node types
  mermaidDiagram += '\n    %% Styling\n';
  mermaidDiagram += '    classDef service fill:#ff9999\n';
  mermaidDiagram += '    classDef event fill:#99ccff\n';
  mermaidDiagram += '    classDef command fill:#99ff99\n';
  mermaidDiagram += '    classDef query fill:#ffcc99\n';
  mermaidDiagram += '    classDef domain fill:#e6ccff\n';
  mermaidDiagram += '    classDef external fill:#ffff99\n';
  mermaidDiagram += '    classDef data fill:#ccffcc\n';

  // Apply classes to nodes
  (nodes as Node<NodeData>[]).forEach(node => {
    const safeName = nodeNameMap[node.id];
    if (safeName) {
      switch (node.type) {
        case 'service':
          mermaidDiagram += `    class ${safeName} service\n`;
          break;
        case 'event':
          mermaidDiagram += `    class ${safeName} event\n`;
          break;
        case 'command':
          mermaidDiagram += `    class ${safeName} command\n`;
          break;
        case 'query':
          mermaidDiagram += `    class ${safeName} query\n`;
          break;
        case 'domain':
          mermaidDiagram += `    class ${safeName} domain\n`;
          break;
        case 'external-system':
          mermaidDiagram += `    class ${safeName} external\n`;
          break;
        case 'data':
          mermaidDiagram += `    class ${safeName} data\n`;
          break;
      }
    }
  });

  return mermaidDiagram;
};

export const generateAIPrompt = ({ flowData, designName }: GeneratePromptOptions): string => {
  if (!flowData?.nodes || flowData.nodes.length === 0) {
    throw new Error('No design data available. Please add some nodes to your canvas first.');
  }

  const { nodes } = flowData;

  // Helper function to get node display name
  const getNodeName = (node: Node<NodeData>) => {
    return node.data?.message?.name || node.data?.service?.name || node.data?.query?.name || 
           node.data?.command?.name || node.data?.data?.name || node.data?.name || 
           `${node.type || 'unknown'}_${node.id.slice(-4)}`;
  };

  // Generate mermaid diagram using the separate function
  const mermaidDiagram = generateMermaidDiagram(flowData);

  // Build component details
  let componentDetails = '\n## Components\n\n';
  (nodes as Node<NodeData>[]).forEach(node => {
    const name = getNodeName(node);
    const nodeData = node.data as NodeData | undefined;
    const summary = nodeData?.message?.summary || nodeData?.service?.summary || 
                  nodeData?.query?.summary || nodeData?.command?.summary || 
                  nodeData?.data?.summary || nodeData?.summary || '';
    const version = nodeData?.message?.version || nodeData?.service?.version || 
                  nodeData?.query?.version || nodeData?.command?.version || 
                  nodeData?.version || '';
    
    componentDetails += `### ${name} (${node.type || 'unknown'})\n`;
    if (version) componentDetails += `**Version:** ${version}\n\n`;
    if (summary) componentDetails += `**Description:** ${summary}\n\n`;
    componentDetails += '---\n\n';
  });

  const aiPrompt = `# Architecture Design: ${designName}

This is a design created using EventCatalog Studio.

## Architecture Diagram

\`\`\`mermaid
${mermaidDiagram}
\`\`\`

${componentDetails}`;

  return aiPrompt;
};