Function print_tree_graphviz

visualize_mcts_claude.py:333–429

使用 Graphviz 生成树形图

(nodes: Dict[str, NodeInfo], root_id: str, output_file: str = "mcts_tree_claude")

def print_tree_graphviz(nodes: Dict[str, NodeInfo], root_id: str, output_file: str = "mcts_tree_claude"):
    """
    使用 Graphviz 生成树形图
    """
    try:
        from graphviz import Digraph
    except ImportError:
        print("Error: graphviz not installed. Install with: pip install graphviz")
        return

    # 找到最佳节点
    best_node_id = None
    best_nodes = [n for n in nodes.values() if n.metric is not None]
    if best_nodes:
        best = min(best_nodes, key=lambda n: n.metric)
        best_node_id = best.id

    g = Digraph(
        name='MCTS_Tree_Claude',
        comment='Claude Code MCTS Search Tree',
        format='png',
        graph_attr={
            'rankdir': 'TB',
            'splines': 'ortho',
            'nodesep': '0.5',
            'ranksep': '1.0',
            'bgcolor': '#f5f5f5'
        },
        node_attr={
            'shape': 'box',
            'style': 'rounded,filled',
            'fontsize': '10',
            'fontname': 'Arial'
        },
        edge_attr={
            'arrowhead': 'vee',
            'color': '#666666',
            'penwidth': '1.5'
        }
    )

    def add_node_recursive(node_id: str):
        if node_id not in nodes:
            return

        node = nodes[node_id]

        # 构建标签
        status_sym = "✓" if node.run_status == "SUCCESS" else "✗" if node.run_status == "FAILED" else "?"
        stage_name = {"root": "ROOT", "draft": "DRAFT", "improve": "IMPROVE"}.get(node.stage, node.stage.upper())

        label_lines = [
            f"{status_sym} {stage_name}",
            f"ID: {node.id[:8]}",
        ]

        if node.metric is not None:
            label_lines.append(f"metric: {node.metric:.4f}")