Diagram Forge

Turn natural language into enterprise-grade architecture diagrams. Diagram Forge is an MCP server that combines template-driven prompt engineering with swappable AI image providers to generate professional diagrams from any MCP-compatible client.

Instead of wrestling with diagramming tools or manually crafting image generation prompts, describe your system in plain English and let Diagram Forge handle the rest — template selection, prompt engineering, style application, and cost tracking.

Features

• 13 diagram templates — Architecture (TOGAF), C4 Container, Executive Infographic, data flow, component, sequence, integration, infographic, generic, product roadmap, workstreams, kanban, and brand infographic

• 2 image providers — Google Gemini (recommended), OpenAI (GPT Image)

• Auto provider selection — Each template recommends the best provider/model for its diagram type

• Template-driven prompts — YAML templates with hex-coded color systems, explicit rendering instructions, and layout rules

• Style references — Feed a visual example to guide output consistency (Gemini)

• Cost tracking — SQLite-backed usage and cost reporting

• Cross-client — Works with Claude Code, Claude Desktop, Codex CLI, Gemini CLI via stdio transport

Related MCP server: Azure Diagram MCP Server

Quick Start

1. Install

pip install diagram-forge

Or from source:

git clone https://github.com/jessepike/diagram-forge.git
cd diagram-forge
pip install -e ".[dev]"

2. Configure a provider

Set at least one API key:

export GEMINI_API_KEY="your-key"       # Google Gemini (recommended)
export OPENAI_API_KEY="your-key"       # OpenAI GPT Image

3. Add to your MCP client

Claude Code (.mcp.json in your project):

{
  "diagram-forge": {
    "command": "python",
    "args": ["-m", "diagram_forge.server"]
  }
}

Claude Desktop (claude_desktop_config.json):

{
  "mcpServers": {
    "diagram-forge": {
      "command": "python",
      "args": ["-m", "diagram_forge.server"]
    }
  }
}

Codex CLI / Gemini CLI — same .mcp.json format as Claude Code.

4. Generate a diagram

Ask your AI client naturally:

"Generate an architecture diagram of a three-tier web app with a React frontend, Node.js API layer, and PostgreSQL database"

Or be more specific:

"Create a TOGAF-style architecture diagram showing our microservices. Use the architecture template, Gemini provider, 16:9 aspect ratio."

MCP Tools

Diagram Types

Style References

Feed a visual example to guide output consistency. Gemini supports this natively via multi-image input.

generate_diagram(prompt="...", style_reference="c4-container")

Save your own styles to ~/.diagram-forge/styles/<name>/reference.png with an optional style.yaml for metadata.

Auto Provider Selection

Set provider="auto" (the default) and Diagram Forge picks the best provider based on the diagram type. Each template includes a tested recommendation. Override with provider="openai" or provider="gemini" when you want a specific model.

Claude Code Plugin

This repo includes a Claude Code plugin in diagram-forge-plugin/ that adds a guided UX layer on top of the MCP server:

• /diagram:create — Guided diagram creation with context gathering

• /diagram:iterate — Refine an existing diagram

• /diagram:usage — View cost report

• /diagram:templates — Browse available templates

• Context-gatherer agent — Automatically explores your project to understand what to diagram

• Diagram intelligence skill — Auto-triggers when you mention diagrams

To use, install the plugin or add the .mcp.json from the plugin directory.

How It Works

1. Template selection — Matches your request to one of 13 YAML templates, each encoding proven prompt patterns (color systems, layer organization, legibility rules)

2. Prompt rendering — Merges your description with the template, substituting variables and applying style defaults

3. Provider dispatch — Sends the engineered prompt to your chosen provider (Gemini or OpenAI)

4. Image handling — Saves the generated image, records cost and metadata to SQLite

5. Iteration — Edit existing diagrams with natural language instructions via providers that support image editing

Development

# Install dev dependencies
pip install -e ".[dev]"

# Run tests (52 tests)
python -m pytest tests/ -v --cov=diagram_forge

# Lint
ruff check src/ tests/

# Type check
mypy src/

# Test MCP tools interactively
npx @modelcontextprotocol/inspector python -m diagram_forge.server

# Run low-cost model benchmark (dry-run first)
python scripts/eval_diagram_models.py --dry-run --max-cost-usd 5
python scripts/eval_diagram_models.py --execute --providers gemini,openai --resolution 1K --max-cases 6 --max-cost-usd 5

Benchmark and model-refresh docs:

• docs/evaluation-runbook.md

• docs/model-refresh-process.md

• evals/benchmark_v1.yaml

Architecture

src/diagram_forge/
  server.py              # FastMCP server — 7 tools, stdio transport
  models.py              # Pydantic v2 models
  config.py              # YAML + env var config loading
  template_engine.py     # Template loading and prompt rendering
  style_manager.py       # Style reference image management
  cost_tracker.py        # SQLite usage/cost tracking
  providers/
    base.py              # BaseImageProvider ABC
    gemini.py            # Google Gemini
    openai_provider.py   # OpenAI GPT Image
  templates/             # 13 YAML prompt templates

License

MIT