mirror of
https://github.com/microsoft/autogen.git
synced 2025-08-06 15:51:58 +00:00
Update graph_utils.py (#2601)
* Update graph_utils.py Adds save_speaker_transitions_dict * Update graph_utils.py Adds export_path:str which defaults to None into visualize_speaker_transitions_dict. In there, run plt.savefig(path) if export_path is not None * Update graph_utils.py * Update graph_utils.py Adds export_path functionality to visualize_speaker_transitions_dict() --------- Co-authored-by: Chi Wang <wang.chi@microsoft.com>
This commit is contained in:
parent
ca8bf89673
commit
6dd2b59da4
@ -1,5 +1,5 @@
|
||||
import logging
|
||||
from typing import Dict, List
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from autogen.agentchat import Agent
|
||||
|
||||
@ -110,7 +110,9 @@ def invert_disallowed_to_allowed(disallowed_speaker_transitions_dict: dict, agen
|
||||
return allowed_speaker_transitions_dict
|
||||
|
||||
|
||||
def visualize_speaker_transitions_dict(speaker_transitions_dict: dict, agents: List[Agent]):
|
||||
def visualize_speaker_transitions_dict(
|
||||
speaker_transitions_dict: dict, agents: List[Agent], export_path: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
Visualize the speaker_transitions_dict using networkx.
|
||||
"""
|
||||
@ -133,4 +135,8 @@ def visualize_speaker_transitions_dict(speaker_transitions_dict: dict, agents: L
|
||||
|
||||
# Visualize
|
||||
nx.draw(G, with_labels=True, font_weight="bold")
|
||||
plt.show()
|
||||
|
||||
if export_path is not None:
|
||||
plt.savefig(export_path)
|
||||
else:
|
||||
plt.show()
|
||||
|
Loading…
x
Reference in New Issue
Block a user