From 6dd2b59da4311d74bea351aeab4e47a84fb523f8 Mon Sep 17 00:00:00 2001 From: "Mohamed A. Fouad" <111041768+moresearch@users.noreply.github.com> Date: Thu, 9 May 2024 21:29:35 -0300 Subject: [PATCH] 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 --- autogen/graph_utils.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/autogen/graph_utils.py b/autogen/graph_utils.py index 88c218fde..d36b47a12 100644 --- a/autogen/graph_utils.py +++ b/autogen/graph_utils.py @@ -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()