diagram-forge

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.

!Example: Enterprise architecture diagram generated with Diagram Forge

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

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): ``json { "diagram-forge": { "command": "python", "args": ["-m", "diagram_forge.server"] } } ``

Claude Desktop (claude_desktop_config.json): ``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

| Tool | Description | |------|-------------| | generate_diagram | Generate a diagram from a text prompt with template and style support | | edit_diagram | Edit an existing diagram with natural language instructions | | list_templates | List available diagram templates and their variables | | list_providers | Show configured providers, API key status, and supported features | | list_styles | List available style reference images | | get_usage_report | View generation costs and usage stats by provider, type, or day | | configure_provider | Set up an API key for a provider (session-only) |

Diagram Types

| Type | Template | Best For | |------|----------|----------| | architecture | Enterprise Architecture (TOGAF) | System architecture, layered designs | | c4_container | C4 Container Diagram | Software system internals, C4 Level 2 | | exec_infographic | Executive Infographic | Stakeholder presentations, semantic colors + icons | | data_flow | Data Flow / Pipeline | ETL pipelines, data movement | | component | Component Detail View | Service internals, module structure | | sequence | Sequence Diagram | Request flows, protocol interactions | | integration | Integration / Connection Map | System connections, API landscape | | infographic | Infographic / Learning Card | Concept explanations, overviews | | product_roadmap | Product Roadmap | Phase pipelines, gate icons, status badges | | workstreams | Workstreams / Priority Lanes | Swimlane planning with status and dependencies | | kanban | Kanban Board | Three-column task boards with category color bars | | brand_infographic | Brand Infographic | Investor/marketing slides with brand aesthetic | | generic | Custom / Freeform | Anything else |

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