ag2 0.10.2__py3-none-any.whl

autogen/graph_utils.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+
+import logging
+
+from .agentchat import Agent
+from .import_utils import optional_import_block, require_optional_import
+
+with optional_import_block():
+    import matplotlib.pyplot as plt
+    import networkx as nx
+
+
+def has_self_loops(allowed_speaker_transitions: dict[str, list[Agent]]) -> bool:
+    """Check if there are self loops in the allowed_speaker_transitions.
+
+    Args:
+        allowed_speaker_transitions (dict[str, list[Agent]]): A dictionary of keys and list as values. The keys are the names of the agents, and the values are the names of the agents that the key agent can transition
+
+    Returns:
+        True if there are self loops in the allowed_speaker_transitions_Dict.
+    """
+    return any(key in value for key, value in allowed_speaker_transitions.items())
+
+
+def check_graph_validity(
+    allowed_speaker_transitions_dict: dict[str, list[Agent]],
+    agents: list[Agent],
+) -> None:
+    """Check the validity of the allowed_speaker_transitions_dict.
+
+    Args:
+        allowed_speaker_transitions_dict (dict[str, list[Agent]]):
+            A dictionary of keys and list as values. The keys are the names of the agents, and the values are the names of the agents that the key agent can transition to.
+        agents (list[Agent]): A list of Agents
+
+    agents: A list of Agents
+
+    Checks for the following:
+        Errors
+        1. The dictionary must have a structure of keys and list as values
+        2. Every key exists in agents.
+        3. Every value is a list of Agents (not string).
+
+    Warnings:
+        1. Warning if there are isolated agent nodes
+        2. Warning if the set of agents in allowed_speaker_transitions do not match agents
+        3. Warning if there are duplicated agents in any values of `allowed_speaker_transitions_dict`
+    """
+    # Errors
+
+    # Check 1. The dictionary must have a structure of keys and list as values
+    if not isinstance(allowed_speaker_transitions_dict, dict):
+        raise ValueError("allowed_speaker_transitions_dict must be a dictionary.")
+
+    # All values must be lists of Agent or empty
+    if not all(isinstance(value, list) for value in allowed_speaker_transitions_dict.values()):
+        raise ValueError("allowed_speaker_transitions_dict must be a dictionary with lists as values.")
+
+    # Check 2. Every key exists in agents
+    if not all(key in agents for key in allowed_speaker_transitions_dict):
+        raise ValueError("allowed_speaker_transitions_dict has keys not in agents.")
+
+    # Check 3. Every value is a list of Agents or empty list (not string).
+    if not all(all(isinstance(agent, Agent) for agent in value) for value in allowed_speaker_transitions_dict.values()):
+        raise ValueError("allowed_speaker_transitions_dict has values that are not lists of Agents.")
+
+    # Warnings
+    # Warning 1. Warning if there are isolated agent nodes, there are not incoming nor outgoing edges
+    # Concat keys if len(value) is positive
+    has_outgoing_edge = []
+    for key, agent_list in allowed_speaker_transitions_dict.items():
+        if len(agent_list) > 0:
+            has_outgoing_edge.append(key)
+    no_outgoing_edges = [agent for agent in agents if agent not in has_outgoing_edge]
+
+    # allowed_speaker_transitions_dict.values() is a list of list of Agents
+    # values_all_agents is a list of all agents in allowed_speaker_transitions_dict.values()
+    has_incoming_edge = []
+    for agent_list in allowed_speaker_transitions_dict.values():
+        if len(agent_list) > 0:
+            has_incoming_edge.extend(agent_list)
+
+    no_incoming_edges = [agent for agent in agents if agent not in has_incoming_edge]
+
+    isolated_agents = set(no_incoming_edges).intersection(set(no_outgoing_edges))
+    if len(isolated_agents) > 0:
+        logging.warning(
+            f"""Warning: There are isolated agent nodes, there are not incoming nor outgoing edges. Isolated agents: {[agent.name for agent in isolated_agents]}"""
+        )
+
+    # Warning 2. Warning if the set of agents in allowed_speaker_transitions do not match agents
+    # Get set of agents
+    agents_in_allowed_speaker_transitions = set(has_incoming_edge).union(set(has_outgoing_edge))
+    full_anti_join = set(agents_in_allowed_speaker_transitions).symmetric_difference(set(agents))
+    if len(full_anti_join) > 0:
+        logging.warning(
+            f"""Warning: The set of agents in allowed_speaker_transitions do not match agents. Offending agents: {[agent.name for agent in full_anti_join]}"""
+        )
+
+    # Warning 3. Warning if there are duplicated agents in any values of `allowed_speaker_transitions_dict`
+    for key, values in allowed_speaker_transitions_dict.items():
+        duplicates = [item for item in values if values.count(item) > 1]
+        unique_duplicates = list(set(duplicates))
+        if unique_duplicates:
+            logging.warning(
+                f"Agent '{key.name}' has duplicate elements: {[agent.name for agent in unique_duplicates]}. Please remove duplicates manually."
+            )
+
+
+def invert_disallowed_to_allowed(
+    disallowed_speaker_transitions_dict: dict[str, list[Agent]], agents: list[Agent]
+) -> dict[str, list[Agent]]:
+    """Invert the disallowed_speaker_transitions_dict to form the allowed_speaker_transitions_dict.
+
+    Start with a fully connected allowed_speaker_transitions_dict of all agents. Remove edges from the fully connected allowed_speaker_transitions_dict according to the disallowed_speaker_transitions_dict to form the allowed_speaker_transitions_dict.
+
+    Args:
+        disallowed_speaker_transitions_dict: A dictionary of keys and list as values. The keys are the names of the agents, and the values are the names of the agents that the key agent cannot transition to.
+        agents: A list of Agents
+
+    Returns:
+        allowed_speaker_transitions_dict: A dictionary of keys and list as values. The keys are the names of the agents, and the values are the names of the agents that the key agent can transition to.
+    """
+    # Create a fully connected allowed_speaker_transitions_dict of all agents
+    allowed_speaker_transitions_dict = {agent: list(agents) for agent in agents}
+
+    # Remove edges from allowed_speaker_transitions_dict according to the disallowed_speaker_transitions_dict
+    for key, value in disallowed_speaker_transitions_dict.items():
+        allowed_speaker_transitions_dict[key] = [
+            agent for agent in allowed_speaker_transitions_dict[key] if agent not in value
+        ]
+
+    return allowed_speaker_transitions_dict
+
+
+@require_optional_import(["matplotlib", "networkx"], "graph")
+def visualize_speaker_transitions_dict(
+    speaker_transitions_dict: dict[str, list[Agent]], agents: list[Agent], export_path: str | None = None
+) -> None:
+    """Visualize the speaker_transitions_dict using networkx.
+
+    Args:
+        speaker_transitions_dict: A dictionary of keys and list as values. The keys are the names of the agents, and the values are the names of the agents that the key agent can transition to.
+        agents: A list of Agents
+        export_path: The path to export the graph. If None, the graph will be shown.
+
+    Returns:
+        None
+
+
+    """
+    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)
+
+    # Visualize
+    nx.draw(g, with_labels=True, font_weight="bold")
+
+    if export_path is not None:
+        plt.savefig(export_path)
+    else:
+        plt.show()