Visualization

get_main_graph

get_main_graph(agent: Agent) -> str

Generates the main graph structure in DOT format for the given agent.

Parameters:

Name	Type	Description	Default
agent	Agent	The agent for which the graph is to be generated.	required

Returns:

Name	Type	Description
str	str	The DOT format string representing the graph.

Source code in src/agents/extensions/visualization.py

def get_main_graph(agent: Agent) -> str:
    """
    Generates the main graph structure in DOT format for the given agent.

    Args:
        agent (Agent): The agent for which the graph is to be generated.

    Returns:
        str: The DOT format string representing the graph.
    """
    parts = [
        """
    digraph G {
        graph [splines=true];
        node [fontname="Arial"];
        edge [penwidth=1.5];
    """
    ]
    node_ids = _GraphNodeIds(agent)
    parts.append(_get_all_nodes(agent, node_ids=node_ids))
    parts.append(_get_all_edges(agent, node_ids=node_ids))
    parts.append("}")
    return "".join(parts)

get_all_nodes

get_all_nodes(
    agent: Agent,
    parent: Agent | None = None,
    visited: set[str] | None = None,
) -> str

Recursively generates the nodes for the given agent and its handoffs in DOT format.

Parameters:

Name	Type	Description	Default
agent	Agent	The agent for which the nodes are to be generated.	required

Returns:

Name	Type	Description
str	str	The DOT format string representing the nodes.

Source code in src/agents/extensions/visualization.py

def get_all_nodes(
    agent: Agent, parent: Agent | None = None, visited: set[str] | None = None
) -> str:
    """
    Recursively generates the nodes for the given agent and its handoffs in DOT format.

    Args:
        agent (Agent): The agent for which the nodes are to be generated.

    Returns:
        str: The DOT format string representing the nodes.
    """
    visited_names = visited if visited is not None else set()
    initially_visited_names = frozenset(visited_names)
    return _get_all_nodes(
        agent,
        parent=parent,
        visited_names=visited_names,
        initially_visited_names=initially_visited_names,
        node_ids=_GraphNodeIds(agent, initially_visited_names=initially_visited_names),
    )

get_all_edges

get_all_edges(
    agent: Agent,
    parent: Agent | None = None,
    visited: set[str] | None = None,
) -> str

Recursively generates the edges for the given agent and its handoffs in DOT format.

Parameters:

Name	Type	Description	Default
agent	Agent	The agent for which the edges are to be generated.	required
parent	Agent	The parent agent. Defaults to None.	None

Returns:

Name	Type	Description
str	str	The DOT format string representing the edges.

Source code in src/agents/extensions/visualization.py

def get_all_edges(
    agent: Agent, parent: Agent | None = None, visited: set[str] | None = None
) -> str:
    """
    Recursively generates the edges for the given agent and its handoffs in DOT format.

    Args:
        agent (Agent): The agent for which the edges are to be generated.
        parent (Agent, optional): The parent agent. Defaults to None.

    Returns:
        str: The DOT format string representing the edges.
    """
    visited_names = visited if visited is not None else set()
    initially_visited_names = frozenset(visited_names)
    return _get_all_edges(
        agent,
        parent=parent,
        visited_names=visited_names,
        initially_visited_names=initially_visited_names,
        node_ids=_GraphNodeIds(agent, initially_visited_names=initially_visited_names),
    )

draw_graph

draw_graph(
    agent: Agent, filename: str | None = None
) -> Source

Draws the graph for the given agent and optionally saves it as a PNG file.

Parameters:

Name	Type	Description	Default
agent	Agent	The agent for which the graph is to be drawn.	required
filename	str	The name of the file to save the graph as a PNG.	None

Returns:

Type	Description
Source	graphviz.Source: The graphviz Source object representing the graph.

Source code in src/agents/extensions/visualization.py

def draw_graph(agent: Agent, filename: str | None = None) -> graphviz.Source:
    """
    Draws the graph for the given agent and optionally saves it as a PNG file.

    Args:
        agent (Agent): The agent for which the graph is to be drawn.
        filename (str): The name of the file to save the graph as a PNG.

    Returns:
        graphviz.Source: The graphviz Source object representing the graph.
    """
    dot_code = get_main_graph(agent)
    graph = graphviz.Source(dot_code)

    if filename:
        graph.render(filename, format="png", cleanup=True)

    return graph