Visualize the tree of thoughts using graphviz.
PARAMETER | DESCRIPTION |
root | The root node of the tree. TYPE: ThinkNode |
Source code in autogen/agentchat/contrib/reasoning_agent.py
| @export_module("autogen")
def visualize_tree(root: ThinkNode) -> None:
"""Visualize the tree of thoughts using graphviz.
Args:
root (ThinkNode): The root node of the tree.
"""
with optional_import_block() as result:
from graphviz import Digraph
if not result.is_successful:
print("Please install graphviz: pip install graphviz")
return
dot = Digraph(comment="Tree of Thoughts")
dot.attr(rankdir="TB") # Top to Bottom direction
def add_nodes(node: ThinkNode, node_id: str = "0"):
# Truncate long content for better visualization
display_content = (node.content[:50] + "...") if len(node.content) > 50 else node.content
# Add node with stats
label = f"{display_content}\n visits: {node.visits}\n value: {node.value}"
dot.node(node_id, label)
# Recursively add children
for i, child in enumerate(node.children):
child_id = f"{node_id}_{i}"
add_nodes(child, child_id)
dot.edge(node_id, child_id)
add_nodes(root)
# Render the graph
try:
dot.render("tree_of_thoughts", view=False, format="png", cleanup=True)
except Exception as e:
print(f"Error rendering graph: {e}")
print("Make sure graphviz is installed on your system: https://graphviz.org/download/")
|