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:
Mohamed A. Fouad 2024-05-09 21:29:35 -03:00 committed by GitHub
parent ca8bf89673
commit 6dd2b59da4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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()