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 <[email protected]>
microsoft · May 10, 2024 · 6dd2b59 · 6dd2b59
1 parent ca8bf89
commit 6dd2b59
Showing 1 changed file with 9 additions and 3 deletions.
diff --git 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()