visualize.py

import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
from workflow import WorkflowDAG

def visualize_graph(workflow_dag:WorkflowDAG, filename="workflow", show=True):
    """使用NetworkX可视化DAG结构（从左到右层级布局）"""
    workflow_dag.assign_layers()
    
    G = nx.DiGraph()
    
    for node_id, node in workflow_dag.nodes.items():
        G.add_node(node_id, node_type=node.task_type, layer=node.layer)
        for child in node.children:
            G.add_edge(node_id, child.task_id)
    
    layers = {}
    for node in G.nodes:
        layer = G.nodes[node]['layer']
        if layer not in layers:
            layers[layer] = []
        layers[layer].append(node)
    
    sorted_layers = sorted(layers.keys())
    
    pos = {}
    for layer in sorted_layers:
        nodes_in_layer = layers[layer]
        y_positions = np.linspace(0, 1, len(nodes_in_layer))
        for i, node in enumerate(nodes_in_layer):
            pos[node] = (layer, y_positions[i])
    
    node_colors = []
    node_shapes = []
    for node in G.nodes:
        node_type = G.nodes[node]['node_type']
        if node_type == "prefill":
            node_colors.append("skyblue")
            node_shapes.append("s")  # 正方形
        elif node_type == "decoding":
            node_colors.append("lightgreen")
            node_shapes.append("d")  # 菱形
        elif node_type == "tool_call":
            node_colors.append("salmon")
            node_shapes.append("o")  # 圆形
    
    plt.figure(figsize=(15, 8))
    for i, (node, shape) in enumerate(zip(G.nodes(), node_shapes)):
        nx.draw_networkx_nodes(G, pos, nodelist=[node], node_color=node_colors[i], 
                                node_shape=shape, node_size=1500)
    
    nx.draw_networkx_edges(G, pos, arrowstyle='->', arrowsize=20)
    
    labels = {node: f"{node}\n({G.nodes[node]['node_type']})" for node in G.nodes}
    nx.draw_networkx_labels(G, pos, labels, font_size=10)
    
    legend_elements = [
        plt.Line2D([0], [0], marker='s', color='w', markerfacecolor='skyblue', markersize=15, label='Prefill'),
        plt.Line2D([0], [0], marker='d', color='w', markerfacecolor='lightgreen', markersize=15, label='Decoding'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='salmon', markersize=15, label='Tool Call')
    ]
    plt.legend(handles=legend_elements, loc='best')
    
    for layer in sorted_layers:
        plt.axvline(x=layer, color='gray', linestyle='--', alpha=0.3)
    
    plt.title("Agentic AI Workflow DAG (Left-to-Right Layout)")
    plt.xlabel("Layer")
    plt.grid(True, axis='x', linestyle='--', alpha=0.3)
    plt.xticks(sorted_layers)
    plt.yticks([])
    
    plt.savefig(f"log/{filename}.png", dpi=300, bbox_inches='tight')
    if show:
        plt.show()
    
    return G

if __name__ == "__main__":
    dag = WorkflowDAG.load_from_json("multi_workflows/workflow_001.json")
    visualize_graph(dag)