nanocode/skills/architecture-diagram/generate.py

#!/usr/bin/env python3
"""
Architecture Diagram Generator
生成精美的深色主题系统架构图

⚠️ 本脚本消耗 AI token，请确保用户明确要求使用
"""

import os
import sys
import json
import re
from datetime import datetime
from pathlib import Path
from typing import Optional

try:
    from openai import OpenAI
except ImportError:
    print("Error: openai package not installed")
    sys.exit(1)

# 配置
SKILL_DIR = Path(__file__).parent
TEMPLATE_PATH = SKILL_DIR / "assets" / "template.html"
OUTPUT_DIR = Path("/tmp/architecture-diagrams")

# 配色方案
COLORS = {
    "frontend": ("rgba(8, 51, 68, 0.4)", "#22d3ee"),
    "backend": ("rgba(6, 78, 59, 0.4)", "#34d399"),
    "database": ("rgba(76, 29, 149, 0.4)", "#a78bfa"),
    "cloud": ("rgba(120, 53, 15, 0.3)", "#fbbf24"),
    "security": ("rgba(136, 19, 55, 0.4)", "#fb7185"),
    "message": ("rgba(251, 146, 60, 0.3)", "#fb923c"),
    "generic": ("rgba(30, 41, 59, 0.5)", "#94a3b8"),
}


def get_api_client():
    """获取 API 客户端（优先 DeepSeek，备选 Moonshot）"""
    # 优先使用 DeepSeek
    api_key = os.getenv("DEEPSEEK_API_KEY")
    if api_key:
        return OpenAI(api_key=api_key, base_url="https://api.deepseek.com")
    
    # 备选 Moonshot
    api_key = os.getenv("MOONSHOT_API_KEY")
    if api_key:
        return OpenAI(api_key=api_key, base_url="https://api.moonshot.cn/v1")
    
    raise ValueError("未找到 DEEPSEEK_API_KEY 或 MOONSHOT_API_KEY")


def build_prompt(project_name: str, description: str, components: list) -> str:
    """构建生成 SVG 的 prompt"""
    
    components_str = "\n".join([
        f"- {c['name']}: {c['type']} - {c.get('description', '')}"
        for c in components
    ])
    
    return f"""你是一个专业的架构图设计师。请根据以下系统描述，生成 SVG 代码。

项目: {project_name}
描述: {description}

组件列表:
{components_str}

要求:
1. 生成完整的 SVG 代码，可直接嵌入 HTML
2. 使用深色主题，背景透明（由外部容器控制）
3. 组件类型对应的配色:
   - frontend: 填充 rgba(8, 51, 68, 0.4), 边框 #22d3ee (青色)
   - backend: 填充 rgba(6, 78, 59, 0.4), 边框 #34d399 (绿色)
   - database: 填充 rgba(76, 29, 149, 0.4), 边框 #a78bfa (紫色)
   - cloud: 填充 rgba(120, 53, 15, 0.3), 边框 #fbbf24 (琥珀色)
   - security: 填充 rgba(136, 19, 55, 0.4), 边框 #fb7185 (玫瑰色)
   - message: 填充 rgba(251, 146, 60, 0.3), 边框 #fb923c (橙色)
   - generic: 填充 rgba(30, 41, 59, 0.5), 边框 #94a3b8 (灰色)

4. 组件使用圆角矩形 (rx="6")，半透明填充，1.5px 边框
5. 组件之间用灰色箭头连接 (#64748b)，箭头使用 marker
6. 组件标签使用白色文字，字号 11px，居中
7. 组件大小约 120x60px，间距合理避免重叠
8. 图例放在右下角，说明各颜色代表的组件类型

只返回 SVG 代码，不要返回其他内容。SVG 的 viewBox 应该是 "0 0 900 500" 左右。
"""


def generate_svg(prompt: str) -> str:
    """调用 AI API 生成 SVG"""
    client = get_api_client()
    
    response = client.chat.completions.create(
        model="deepseek-chat",  # 或 moonshot 模型
        messages=[
            {"role": "system", "content": "你是一个专业的架构图 SVG 生成器。只输出 SVG 代码，不输出任何解释。"},
            {"role": "user", "content": prompt}
        ],
        temperature=0.3,
        max_tokens=4000
    )
    
    return response.choices[0].message.content


def extract_svg_code(content: str) -> str:
    """从 AI 响应中提取 SVG 代码"""
    # 尝试提取 <svg> 标签内容
    svg_match = re.search(r'<svg[^>]*>.*?</svg>', content, re.DOTALL)
    if svg_match:
        return svg_match.group(0)
    
    # 如果没有标签，可能是纯 SVG 路径
    if '<' in content and '>' in content:
        return f'<svg viewBox="0 0 900 500" xmlns="http://www.w3.org/2000/svg">{content}</svg>'
    
    return content


def generate_diagram(
    project_name: str,
    description: str,
    components: list,
    output_path: Optional[str] = None
) -> str:
    """
    生成架构图
    
    Args:
        project_name: 项目名称
        description: 项目描述
        components: 组件列表，每项包含 name, type, description
        output_path: 输出路径，默认保存到 /tmp/architecture-diagrams/
    
    Returns:
        生成的 HTML 文件路径
    """
    # 确保输出目录存在
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    
    # 生成文件名
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    safe_name = re.sub(r'[^\w\-]', '_', project_name.lower())
    
    if output_path is None:
        output_path = OUTPUT_DIR / f"{safe_name}_{timestamp}.html"
    else:
        output_path = Path(output_path)
    
    # 读取模板
    template = TEMPLATE_PATH.read_text(encoding='utf-8')
    
    # 构建 prompt 并生成 SVG
    prompt = build_prompt(project_name, description, components)
    svg_content = generate_svg(prompt)
    svg_content = extract_svg_code(svg_content)
    
    # 生成卡片内容
    frontend_items = [f"<li>{c['name']}</li>" for c in components if c['type'] == 'frontend']
    backend_items = [f"<li>{c['name']}</li>" for c in components if c['type'] == 'backend']
    database_items = [f"<li>{c['name']}</li>" for c in components if c['type'] == 'database']
    
    # 填充模板
    html = template.replace("{{PROJECT_NAME}}", project_name)
    html = html.replace("{{SVG_CONTENT}}", svg_content)
    html = html.replace("{{FRONTEND_ITEMS}}", "\n".join(frontend_items) if frontend_items else "<li>无</li>")
    html = html.replace("{{BACKEND_ITEMS}}", "\n".join(backend_items) if backend_items else "<li>无</li>")
    html = html.replace("{{DATABASE_ITEMS}}", "\n".join(database_items) if database_items else "<li>无</li>")
    html = html.replace("{{TIMESTAMP}}", datetime.now().strftime("%Y-%m-%d %H:%M"))
    
    # 保存
    output_path.write_text(html, encoding='utf-8')
    
    return str(output_path)


if __name__ == "__main__":
    # 测试用例
    test_components = [
        {"name": "React 前端", "type": "frontend", "description": "用户界面"},
        {"name": "FastAPI", "type": "backend", "description": "API 服务"},
        {"name": "SQLite", "type": "database", "description": "数据存储"},
        {"name": "Redis", "type": "database", "description": "缓存"},
    ]
    
    result = generate_diagram(
        project_name="Test Project",
        description="一个简单的测试项目",
        components=test_components
    )
    print(f"Generated: {result}")