admin-valentin/openai-agents-python
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor visualization functions to improve formatting and streamline edge generation

Browse source
This commit is contained in:
Martín Bravo 2025-03-25 19:12:40 +01:00
parent 5ad53d8000
commit 351b6074e5
2 changed files with 6 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
src/agents/extensions/visualization.py
View file

@ -45,10 +45,10 @@ def get_all_nodes(agent: Agent, parent: Optional[Agent] = None) -> str:
# Start and end the graph # Start and end the graph
parts.append( parts.append(
f'"__start__" [label="__start__", shape=ellipse, style=filled, ' '"__start__" [label="__start__", shape=ellipse, style=filled, '
f"fillcolor=lightblue, width=0.5, height=0.3];" "fillcolor=lightblue, width=0.5, height=0.3];"
f'"__end__" [label="__end__", shape=ellipse, style=filled, ' '"__end__" [label="__end__", shape=ellipse, style=filled, '
f"fillcolor=lightblue, width=0.5, height=0.3];" "fillcolor=lightblue, width=0.5, height=0.3];"
) )
# Ensure parent agent node is colored # Ensure parent agent node is colored
if not parent: if not parent:
@ -95,9 +95,7 @@ def get_all_edges(agent: Agent, parent: Optional[Agent] = None) -> str:
parts = [] parts = []
if not parent: if not parent:
parts.append( parts.append(f'"__start__" -> "{agent.name}";')
f'"__start__" -> "{agent.name}";'
)
for tool in agent.tools: for tool in agent.tools:
parts.append(f""" parts.append(f"""
@ -114,9 +112,7 @@ def get_all_edges(agent: Agent, parent: Optional[Agent] = None) -> str:
parts.append(get_all_edges(handoff, agent)) parts.append(get_all_edges(handoff, agent))
if not agent.handoffs and not isinstance(agent, Tool): if not agent.handoffs and not isinstance(agent, Tool):
parts.append( parts.append(f'"{agent.name}" -> "__end__";')
f'"{agent.name}" -> "__end__";'
)
return "".join(parts) return "".join(parts)

1
tests/test_visualization.py
View file

@ -103,7 +103,6 @@ def test_get_all_edges(mock_agent):
assert '"Agent1" -> "Handoff1";' in result assert '"Agent1" -> "Handoff1";' in result
def test_draw_graph(mock_agent): def test_draw_graph(mock_agent):
graph = draw_graph(mock_agent) graph = draw_graph(mock_agent)
assert isinstance(graph, graphviz.Source) assert isinstance(graph, graphviz.Source)