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

Linting

Browse source
This commit is contained in:
Martín Bravo 2025-03-13 18:49:12 +01:00
parent 2993d26d58
commit 29e9983ae8
2 changed files with 81 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
src/agents/visualizations.py
View file

@ -13,7 +13,8 @@ def get_main_graph(agent: Agent) -> str:
Returns:
str: The DOT format string representing the graph.
"""
parts = ["""
parts = [
"""
digraph G {
graph [splines=true];
node [fontname="Arial"];
@ -21,7 +22,8 @@ def get_main_graph(agent: Agent) -> str:
"__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("}")
@ -108,6 +110,6 @@ def draw_graph(agent: Agent, filename: str = None) -> graphviz.Source:
graph = graphviz.Source(dot_code)
if filename:
graph.render(filename, format='png')
graph.render(filename, format="png")
return graph

81
tests/test_visualizations.py
View file

@ -1,8 +1,11 @@
import pytest
from unittest.mock import Mock
from src.agents.visualizations import get_main_graph, get_all_nodes, get_all_edges, draw_graph
from src.agents.agent import Agent
import graphviz
import pytest
from src.agents.agent import Agent
from src.agents.visualizations import draw_graph, get_all_edges, get_all_nodes, get_main_graph
@pytest.fixture
def mock_agent():
@ -23,25 +26,52 @@ def mock_agent():
return agent
def test_get_main_graph(mock_agent):
result = get_main_graph(mock_agent)
assert "digraph G" in result
assert 'graph [splines=true];' in result
assert "graph [splines=true];" in result
assert 'node [fontname="Arial"];' in result
assert 'edge [penwidth=1.5];' in result
assert "edge [penwidth=1.5];" in result
assert '"__start__" [shape=ellipse, style=filled, fillcolor=lightblue];' in result
assert '"__end__" [shape=ellipse, style=filled, fillcolor=lightblue];' in result
assert '"Agent1" [label="Agent1", shape=box, style=filled, fillcolor=lightyellow, width=1.5, height=0.8];' in result
assert '"Tool1" [label="Tool1", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in result
assert '"Tool2" [label="Tool2", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in result
assert '"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, fillcolor=lightyellow, width=1.5, height=0.8];' in result
assert (
'"Agent1" [label="Agent1", shape=box, style=filled, fillcolor=lightyellow, width=1.5, height=0.8];'
in result
)
assert (
'"Tool1" [label="Tool1", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];'
in result
)
assert (
'"Tool2" [label="Tool2", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];'
in result
)
assert (
'"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, fillcolor=lightyellow, width=1.5, height=0.8];'
in result
)
def test_get_all_nodes(mock_agent):
result = get_all_nodes(mock_agent)
assert '"Agent1" [label="Agent1", shape=box, style=filled, fillcolor=lightyellow, width=1.5, height=0.8];' in result
assert '"Tool1" [label="Tool1", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in result
assert '"Tool2" [label="Tool2", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in result
assert '"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, fillcolor=lightyellow, width=1.5, height=0.8];' in result
assert (
'"Agent1" [label="Agent1", shape=box, style=filled, fillcolor=lightyellow, width=1.5, height=0.8];'
in result
)
assert (
'"Tool1" [label="Tool1", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];'
in result
)
assert (
'"Tool2" [label="Tool2", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];'
in result
)
assert (
'"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, fillcolor=lightyellow, width=1.5, height=0.8];'
in result
)
def test_get_all_edges(mock_agent):
result = get_all_edges(mock_agent)
@ -53,16 +83,29 @@ def test_get_all_edges(mock_agent):
assert '"Agent1" -> "Handoff1";' in result
assert '"Handoff1" -> "__end__";' in result
def test_draw_graph(mock_agent):
graph = draw_graph(mock_agent)
assert isinstance(graph, graphviz.Source)
assert "digraph G" in graph.source
assert 'graph [splines=true];' in graph.source
assert "graph [splines=true];" in graph.source
assert 'node [fontname="Arial"];' in graph.source
assert 'edge [penwidth=1.5];' in graph.source
assert "edge [penwidth=1.5];" in graph.source
assert '"__start__" [shape=ellipse, style=filled, fillcolor=lightblue];' in graph.source
assert '"__end__" [shape=ellipse, style=filled, fillcolor=lightblue];' in graph.source
assert '"Agent1" [label="Agent1", shape=box, style=filled, fillcolor=lightyellow, width=1.5, height=0.8];' in graph.source
assert '"Tool1" [label="Tool1", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in graph.source
assert '"Tool2" [label="Tool2", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in graph.source
assert '"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, fillcolor=lightyellow, width=1.5, height=0.8];' in graph.source
assert (
'"Agent1" [label="Agent1", shape=box, style=filled, fillcolor=lightyellow, width=1.5, height=0.8];'
in graph.source
)
assert (
'"Tool1" [label="Tool1", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];'
in graph.source
)
assert (
'"Tool2" [label="Tool2", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];'
in graph.source
)
assert (
'"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, fillcolor=lightyellow, width=1.5, height=0.8];'
in graph.source
)