Visual enhancement to `visualize_speaker_transitions_dict`

[signalprime]

Getting to know the visualization tool for speakers, my first try resulted in text being cut-off and arrows occasionally being hidden by text labels

I made some updates seen below to fix those issue:

def visualize_speaker_transitions_dict(speaker_transitions_dict: dict, agents: List[Agent]):
    """
    Visualize the speaker_transitions_dict using networkx, with increased figure size and adjusted layout.
    """
    try:
        import networkx as nx
        import matplotlib.pyplot as plt
    except ImportError as e:
        logging.fatal("Failed to import networkx or matplotlib. Try running 'pip install autogen[graphs]'")
        raise e

    G = nx.DiGraph()

    # Add nodes
    G.add_nodes_from([agent.name for agent in agents])

    # Add edges
    for key, value in speaker_transitions_dict.items():
        for agent in value:
            G.add_edge(key.name, agent.name)

    # Set the figure size
    plt.figure(figsize=(12, 8))

    # Visualize
    pos = nx.spring_layout(G)  # For consistent positioning

    # Draw nodes and edges
    nx.draw_networkx_nodes(G, pos)
    nx.draw_networkx_edges(G, pos)

    # Draw labels below the nodes
    label_pos = {k: [v[0], v[1] - 0.1] for k, v in pos.items()}  # Shift labels below the nodes
    nx.draw_networkx_labels(G, label_pos, verticalalignment='top', font_color="darkgreen")

    # Adding margins
    ax = plt.gca()
    ax.margins(0.1)  # Increase the margin value if needed


    # Adding a dynamic title
    total_transitions = sum(len(v) for v in speaker_transitions_dict.values())
    title = f'Agent Interactions: {len(agents)} Agents, {total_transitions} Potential Transitions'
    plt.title(title)

    plt.show()

Any thoughts or feedback?

~ "Til all are one" - Optimus

[joshkyh]

Thanks, it does look better than the original
nx.draw(G, with_labels=True, font_weight="bold")

Maybe a good idea to expose font_color / font_size whilst you're at it? :)

[ekzhu]

I am converting this into a discussion.