zyflow 0.6.4

package/packages/gitdiagram-core/src/generator.ts

@@ -0,0 +1,283 @@
+/**
+ * Diagram Generator
+ *
+ * Implements the 3-stage AI generation process:
+ * 1. Analyze project and create explanation
+ * 2. Map components to files/directories
+ * 3. Generate Mermaid.js diagram
+ */
+
+import type { LLMAdapter, LLMMessage } from './llm-adapter.js';
+import {
+  SYSTEM_FIRST_PROMPT,
+  SYSTEM_SECOND_PROMPT,
+  SYSTEM_THIRD_PROMPT,
+  SYSTEM_MODIFY_PROMPT,
+  ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT,
+  createFirstPromptUserMessage,
+  createSecondPromptUserMessage,
+  createThirdPromptUserMessage,
+  createModifyPromptUserMessage,
+} from './prompts.js';
+import { getProjectContext, type FileTreeOptions } from './file-tree.js';
+
+export interface GeneratorOptions {
+  /** LLM adapter to use */
+  llm: LLMAdapter;
+  /** File tree options */
+  fileTreeOptions?: FileTreeOptions;
+  /** Custom instructions for generation */
+  instructions?: string;
+  /** Callback for progress updates */
+  onProgress?: (stage: GenerationStage, message: string) => void;
+}
+
+export type GenerationStage =
+  | 'analyzing'
+  | 'mapping'
+  | 'generating'
+  | 'complete'
+  | 'error';
+
+export interface GenerationResult {
+  /** Generated Mermaid.js code */
+  mermaidCode: string;
+  /** Architecture explanation */
+  explanation: string;
+  /** Component to file/directory mappings */
+  componentMapping: Record<string, string>;
+  /** Raw component mapping text */
+  rawComponentMapping: string;
+  /** File tree used for generation */
+  fileTree: string;
+  /** README content if found */
+  readme: string | null;
+}
+
+export interface ModifyResult {
+  /** Modified Mermaid.js code */
+  mermaidCode: string;
+}
+
+/**
+ * Parse component mapping from LLM response
+ */
+function parseComponentMapping(text: string): Record<string, string> {
+  const mapping: Record<string, string> = {};
+
+  // Extract content between <component_mapping> tags
+  const match = text.match(
+    /<component_mapping>([\s\S]*?)<\/component_mapping>/
+  );
+  const content = match ? match[1] : text;
+
+  // Parse numbered list format: "1. [Component]: [Path]"
+  const lines = content.split('\n');
+  for (const line of lines) {
+    const lineMatch = line.match(/^\d+\.\s*(.+?):\s*(.+)$/);
+    if (lineMatch) {
+      const [, component, path] = lineMatch;
+      mapping[component.trim()] = path.trim();
+    }
+  }
+
+  return mapping;
+}
+
+/**
+ * Extract explanation from LLM response
+ */
+function extractExplanation(text: string): string {
+  const match = text.match(/<explanation>([\s\S]*?)<\/explanation>/);
+  return match ? match[1].trim() : text.trim();
+}
+
+/**
+ * Clean Mermaid code (remove markdown fences if present)
+ */
+function cleanMermaidCode(code: string): string {
+  let cleaned = code.trim();
+
+  // Remove markdown code fences
+  if (cleaned.startsWith('```mermaid')) {
+    cleaned = cleaned.slice(10);
+  } else if (cleaned.startsWith('```')) {
+    cleaned = cleaned.slice(3);
+  }
+
+  if (cleaned.endsWith('```')) {
+    cleaned = cleaned.slice(0, -3);
+  }
+
+  return cleaned.trim();
+}
+
+/**
+ * Generate a diagram for a project
+ */
+export async function generateDiagram(
+  projectPath: string,
+  options: GeneratorOptions
+): Promise<GenerationResult> {
+  const { llm, fileTreeOptions, instructions, onProgress } = options;
+
+  // Stage 1: Get project context
+  onProgress?.('analyzing', 'Reading project structure...');
+  const { fileTree, readme } = await getProjectContext(
+    projectPath,
+    fileTreeOptions
+  );
+
+  // Stage 1: Generate explanation
+  onProgress?.('analyzing', 'Analyzing architecture...');
+  const firstPromptMessages: LLMMessage[] = [
+    {
+      role: 'system',
+      content: instructions
+        ? `${SYSTEM_FIRST_PROMPT}\n\n${ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT}`
+        : SYSTEM_FIRST_PROMPT,
+    },
+    {
+      role: 'user',
+      content: instructions
+        ? `${createFirstPromptUserMessage(fileTree, readme ?? 'No README found.')}\n\n<instructions>\n${instructions}\n</instructions>`
+        : createFirstPromptUserMessage(fileTree, readme ?? 'No README found.'),
+    },
+  ];
+
+  const explanationResponse = await llm.complete(firstPromptMessages, {
+    maxTokens: 4096,
+    temperature: 0.7,
+  });
+
+  // Check for BAD_INSTRUCTIONS
+  if (explanationResponse.trim() === 'BAD_INSTRUCTIONS') {
+    throw new Error(
+      'Custom instructions were unclear or unrelated to the task'
+    );
+  }
+
+  const explanation = extractExplanation(explanationResponse);
+
+  // Stage 2: Map components to files
+  onProgress?.('mapping', 'Mapping components to files...');
+  const secondPromptMessages: LLMMessage[] = [
+    { role: 'system', content: SYSTEM_SECOND_PROMPT },
+    {
+      role: 'user',
+      content: createSecondPromptUserMessage(explanation, fileTree),
+    },
+  ];
+
+  const mappingResponse = await llm.complete(secondPromptMessages, {
+    maxTokens: 2048,
+    temperature: 0.5,
+  });
+
+  const componentMapping = parseComponentMapping(mappingResponse);
+  const rawComponentMapping = mappingResponse;
+
+  // Stage 3: Generate Mermaid diagram
+  onProgress?.('generating', 'Generating diagram...');
+  const thirdPromptMessages: LLMMessage[] = [
+    { role: 'system', content: SYSTEM_THIRD_PROMPT },
+    {
+      role: 'user',
+      content: createThirdPromptUserMessage(explanation, mappingResponse),
+    },
+  ];
+
+  const diagramResponse = await llm.complete(thirdPromptMessages, {
+    maxTokens: 8192,
+    temperature: 0.7,
+  });
+
+  const mermaidCode = cleanMermaidCode(diagramResponse);
+
+  onProgress?.('complete', 'Diagram generated successfully');
+
+  return {
+    mermaidCode,
+    explanation,
+    componentMapping,
+    rawComponentMapping,
+    fileTree,
+    readme,
+  };
+}
+
+/**
+ * Modify an existing diagram based on instructions
+ */
+export async function modifyDiagram(
+  diagram: string,
+  explanation: string,
+  instructions: string,
+  llm: LLMAdapter
+): Promise<ModifyResult> {
+  const messages: LLMMessage[] = [
+    { role: 'system', content: SYSTEM_MODIFY_PROMPT },
+    {
+      role: 'user',
+      content: createModifyPromptUserMessage(diagram, explanation, instructions),
+    },
+  ];
+
+  const response = await llm.complete(messages, {
+    maxTokens: 8192,
+    temperature: 0.7,
+  });
+
+  // Check for BAD_INSTRUCTIONS
+  if (response.trim() === 'BAD_INSTRUCTIONS') {
+    throw new Error(
+      'Modification instructions were unclear or unrelated to the task'
+    );
+  }
+
+  return {
+    mermaidCode: cleanMermaidCode(response),
+  };
+}
+
+/**
+ * DiagramGenerator class for stateful generation
+ */
+export class DiagramGenerator {
+  private llm: LLMAdapter;
+  private fileTreeOptions: FileTreeOptions;
+
+  constructor(options: { llm: LLMAdapter; fileTreeOptions?: FileTreeOptions }) {
+    this.llm = options.llm;
+    this.fileTreeOptions = options.fileTreeOptions ?? {};
+  }
+
+  /**
+   * Generate a diagram for a project
+   */
+  async generate(
+    projectPath: string,
+    options?: {
+      instructions?: string;
+      onProgress?: (stage: GenerationStage, message: string) => void;
+    }
+  ): Promise<GenerationResult> {
+    return generateDiagram(projectPath, {
+      llm: this.llm,
+      fileTreeOptions: this.fileTreeOptions,
+      instructions: options?.instructions,
+      onProgress: options?.onProgress,
+    });
+  }
+
+  /**
+   * Modify an existing diagram
+   */
+  async modify(
+    diagram: string,
+    explanation: string,
+    instructions: string
+  ): Promise<ModifyResult> {
+    return modifyDiagram(diagram, explanation, instructions, this.llm);
+  }
+}

package/packages/gitdiagram-core/src/index.ts

@@ -0,0 +1,78 @@
+/**
+ * @zyflow/gitdiagram-core
+ *
+ * AI-powered repository architecture diagram generator
+ * Based on GitDiagram by Ahmed Khaleel
+ *
+ * @example
+ * ```typescript
+ * import { DiagramGenerator, ClaudeAdapter } from '@zyflow/gitdiagram-core';
+ *
+ * const llm = new ClaudeAdapter({ apiKey: process.env.ANTHROPIC_API_KEY });
+ * const generator = new DiagramGenerator({ llm });
+ *
+ * const result = await generator.generate('/path/to/project', {
+ *   onProgress: (stage, message) => console.log(`[${stage}] ${message}`),
+ * });
+ *
+ * console.log(result.mermaidCode);
+ * ```
+ */
+
+// Prompts
+export {
+  SYSTEM_FIRST_PROMPT,
+  SYSTEM_SECOND_PROMPT,
+  SYSTEM_THIRD_PROMPT,
+  ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT,
+  SYSTEM_MODIFY_PROMPT,
+  createFirstPromptUserMessage,
+  createSecondPromptUserMessage,
+  createThirdPromptUserMessage,
+  createModifyPromptUserMessage,
+  withAdditionalInstructions,
+} from './prompts.js';
+
+// File Tree
+export {
+  generateFileTree,
+  readReadme,
+  getProjectContext,
+  type FileTreeOptions,
+} from './file-tree.js';
+
+// LLM Adapter
+export {
+  ClaudeAdapter,
+  OpenAIAdapter,
+  GeminiAdapter,
+  createLLMAdapter,
+  getApiKeyFromEnv,
+  type LLMAdapter,
+  type LLMMessage,
+  type LLMCompletionOptions,
+} from './llm-adapter.js';
+
+// Generator
+export {
+  generateDiagram,
+  modifyDiagram,
+  DiagramGenerator,
+  type GeneratorOptions,
+  type GenerationStage,
+  type GenerationResult,
+  type ModifyResult,
+} from './generator.js';
+
+// Mermaid Utilities
+export {
+  validateMermaidSyntax,
+  extractClickEvents,
+  updateClickEvents,
+  pathsToGitHubUrls,
+  extractNodeIds,
+  ensureColorStyles,
+  formatMermaidCode,
+  type ClickEvent,
+  type ValidationResult,
+} from './mermaid-utils.js';

package/packages/gitdiagram-core/src/llm-adapter.ts

@@ -0,0 +1,235 @@
+/**
+ * LLM Adapter
+ *
+ * Abstraction layer for different LLM providers (Claude, OpenAI, etc.)
+ */
+
+export interface LLMMessage {
+  role: 'system' | 'user' | 'assistant';
+  content: string;
+}
+
+export interface LLMCompletionOptions {
+  /** Maximum tokens to generate */
+  maxTokens?: number;
+  /** Temperature for generation */
+  temperature?: number;
+  /** Stop sequences */
+  stopSequences?: string[];
+}
+
+export interface LLMAdapter {
+  /** Provider name */
+  name: string;
+
+  /** Generate a completion */
+  complete(
+    messages: LLMMessage[],
+    options?: LLMCompletionOptions
+  ): Promise<string>;
+}
+
+/**
+ * Claude API Adapter using Anthropic SDK
+ */
+export class ClaudeAdapter implements LLMAdapter {
+  name = 'claude';
+  private apiKey: string;
+  private model: string;
+
+  constructor(options: { apiKey: string; model?: string }) {
+    this.apiKey = options.apiKey;
+    this.model = options.model ?? 'claude-sonnet-4-20250514';
+  }
+
+  async complete(
+    messages: LLMMessage[],
+    options: LLMCompletionOptions = {}
+  ): Promise<string> {
+    // Extract system message
+    const systemMessage = messages.find((m) => m.role === 'system');
+    const nonSystemMessages = messages.filter((m) => m.role !== 'system');
+
+    const response = await fetch('https://api.anthropic.com/v1/messages', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'x-api-key': this.apiKey,
+        'anthropic-version': '2023-06-01',
+      },
+      body: JSON.stringify({
+        model: this.model,
+        max_tokens: options.maxTokens ?? 4096,
+        temperature: options.temperature ?? 0.7,
+        system: systemMessage?.content,
+        messages: nonSystemMessages.map((m) => ({
+          role: m.role,
+          content: m.content,
+        })),
+        stop_sequences: options.stopSequences,
+      }),
+    });
+
+    if (!response.ok) {
+      const error = await response.text();
+      throw new Error(`Claude API error: ${response.status} - ${error}`);
+    }
+
+    const data = (await response.json()) as {
+      content?: Array<{ type: string; text?: string }>;
+    };
+    const textContent = data.content?.find((c) => c.type === 'text');
+    return textContent?.text ?? '';
+  }
+}
+
+/**
+ * Google Gemini API Adapter
+ */
+export class GeminiAdapter implements LLMAdapter {
+  name = 'gemini';
+  private apiKey: string;
+  private model: string;
+
+  constructor(options: { apiKey: string; model?: string }) {
+    this.apiKey = options.apiKey;
+    this.model = options.model ?? 'gemini-2.0-flash';
+  }
+
+  async complete(
+    messages: LLMMessage[],
+    options: LLMCompletionOptions = {}
+  ): Promise<string> {
+    // Extract system message and combine with first user message
+    const systemMessage = messages.find((m) => m.role === 'system');
+    const nonSystemMessages = messages.filter((m) => m.role !== 'system');
+
+    // Gemini uses 'contents' format with 'parts'
+    const contents = nonSystemMessages.map((m) => ({
+      role: m.role === 'assistant' ? 'model' : 'user',
+      parts: [{ text: m.content }],
+    }));
+
+    // Prepend system instruction to first user message if exists
+    if (systemMessage && contents.length > 0 && contents[0].role === 'user') {
+      contents[0].parts[0].text = `${systemMessage.content}\n\n${contents[0].parts[0].text}`;
+    }
+
+    const response = await fetch(
+      `https://generativelanguage.googleapis.com/v1beta/models/${this.model}:generateContent?key=${this.apiKey}`,
+      {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+        },
+        body: JSON.stringify({
+          contents,
+          generationConfig: {
+            maxOutputTokens: options.maxTokens ?? 4096,
+            temperature: options.temperature ?? 0.7,
+            stopSequences: options.stopSequences,
+          },
+        }),
+      }
+    );
+
+    if (!response.ok) {
+      const error = await response.text();
+      throw new Error(`Gemini API error: ${response.status} - ${error}`);
+    }
+
+    const data = (await response.json()) as {
+      candidates?: Array<{
+        content?: { parts?: Array<{ text?: string }> };
+      }>;
+    };
+
+    return data.candidates?.[0]?.content?.parts?.[0]?.text ?? '';
+  }
+}
+
+/**
+ * OpenAI API Adapter
+ */
+export class OpenAIAdapter implements LLMAdapter {
+  name = 'openai';
+  private apiKey: string;
+  private model: string;
+  private baseUrl: string;
+
+  constructor(options: { apiKey: string; model?: string; baseUrl?: string }) {
+    this.apiKey = options.apiKey;
+    this.model = options.model ?? 'gpt-4o';
+    this.baseUrl = options.baseUrl ?? 'https://api.openai.com/v1';
+  }
+
+  async complete(
+    messages: LLMMessage[],
+    options: LLMCompletionOptions = {}
+  ): Promise<string> {
+    const response = await fetch(`${this.baseUrl}/chat/completions`, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: `Bearer ${this.apiKey}`,
+      },
+      body: JSON.stringify({
+        model: this.model,
+        max_tokens: options.maxTokens ?? 4096,
+        temperature: options.temperature ?? 0.7,
+        messages: messages.map((m) => ({
+          role: m.role,
+          content: m.content,
+        })),
+        stop: options.stopSequences,
+      }),
+    });
+
+    if (!response.ok) {
+      const error = await response.text();
+      throw new Error(`OpenAI API error: ${response.status} - ${error}`);
+    }
+
+    const data = (await response.json()) as {
+      choices?: Array<{ message?: { content?: string } }>;
+    };
+    return data.choices?.[0]?.message?.content ?? '';
+  }
+}
+
+/**
+ * Create an LLM adapter based on provider name
+ */
+export function createLLMAdapter(
+  provider: 'claude' | 'openai' | 'gemini',
+  options: { apiKey: string; model?: string; baseUrl?: string }
+): LLMAdapter {
+  switch (provider) {
+    case 'claude':
+      return new ClaudeAdapter(options);
+    case 'openai':
+      return new OpenAIAdapter(options);
+    case 'gemini':
+      return new GeminiAdapter(options);
+    default:
+      throw new Error(`Unknown LLM provider: ${provider}`);
+  }
+}
+
+/**
+ * Get API key from environment variable
+ */
+export function getApiKeyFromEnv(provider: 'claude' | 'openai' | 'gemini'): string | null {
+  switch (provider) {
+    case 'claude':
+      return (
+        process.env.ANTHROPIC_API_KEY ?? process.env.CLAUDE_API_KEY ?? null
+      );
+    case 'openai':
+      return process.env.OPENAI_API_KEY ?? null;
+    case 'gemini':
+      return process.env.GOOGLE_API_KEY ?? process.env.GEMINI_API_KEY ?? null;
+    default:
+      return null;
+  }
+}