diff --git a/src/agents/visualizations.py b/src/agents/visualizations.py new file mode 100644 index 0000000..934647f --- /dev/null +++ b/src/agents/visualizations.py @@ -0,0 +1,113 @@ +import graphviz + +from src.agents.agent import Agent + + +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]; + + "__start__" [shape=ellipse, style=filled, fillcolor=lightblue]; + "__end__" [shape=ellipse, style=filled, fillcolor=lightblue]; + """] + parts.append(get_all_nodes(agent)) + parts.append(get_all_edges(agent)) + parts.append("}") + return "".join(parts) + + +def get_all_nodes(agent: Agent, parent: Agent = 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. + """ + parts = [] + + # Ensure parent agent node is colored + if not parent: + parts.append(f""" + "{agent.name}" [label="{agent.name}", shape=box, style=filled, fillcolor=lightyellow, width=1.5, height=0.8];""") + + # Smaller tools (ellipse, green) + for tool in agent.tools: + parts.append(f""" + "{tool.name}" [label="{tool.name}", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];""") + + # Bigger handoffs (rounded box, yellow) + for handoff in agent.handoffs: + parts.append(f""" + "{handoff.name}" [label="{handoff.name}", shape=box, style=filled, style=rounded, fillcolor=lightyellow, width=1.5, height=0.8];""") + parts.append(get_all_nodes(handoff)) + + return "".join(parts) + + +def get_all_edges(agent: Agent, parent: Agent = 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. + """ + parts = [] + + if not parent: + parts.append(f""" + "__start__" -> "{agent.name}";""") + + for tool in agent.tools: + parts.append(f""" + "{agent.name}" -> "{tool.name}" [style=dotted, penwidth=1.5]; + "{tool.name}" -> "{agent.name}" [style=dotted, penwidth=1.5];""") + + if not agent.handoffs: + parts.append(f""" + "{agent.name}" -> "__end__";""") + + for handoff in agent.handoffs: + parts.append(f""" + "{agent.name}" -> "{handoff.name}";""") + parts.append(get_all_edges(handoff, agent)) + + return "".join(parts) + + +def draw_graph(agent: Agent, filename: str = 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') + + return graph \ No newline at end of file