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

81
tests/test_visualizations.py
View file

@ -1,8 +1,11 @@
import pytest
from unittest.mock import Mock 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 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 @pytest.fixture
def mock_agent(): def mock_agent():
@ -23,25 +26,52 @@ def mock_agent():
return agent return agent
def test_get_main_graph(mock_agent): def test_get_main_graph(mock_agent):
result = get_main_graph(mock_agent) result = get_main_graph(mock_agent)
assert "digraph G" in result assert "digraph G" in result
assert 'graph [splines=true];' in result assert "graph [splines=true];" in result
assert 'node [fontname="Arial"];' 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 '"__start__" [shape=ellipse, style=filled, fillcolor=lightblue];' in result
assert '"__end__" [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 (
assert '"Tool1" [label="Tool1", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in result '"Agent1" [label="Agent1", shape=box, style=filled, fillcolor=lightyellow, width=1.5, height=0.8];'
assert '"Tool2" [label="Tool2", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in result in result
assert '"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, 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): def test_get_all_nodes(mock_agent):
result = 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 (
assert '"Tool1" [label="Tool1", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in result '"Agent1" [label="Agent1", shape=box, style=filled, fillcolor=lightyellow, width=1.5, height=0.8];'
assert '"Tool2" [label="Tool2", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in result in result
assert '"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, 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): def test_get_all_edges(mock_agent):
result = 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 '"Agent1" -> "Handoff1";' in result
assert '"Handoff1" -> "__end__";' in result assert '"Handoff1" -> "__end__";' 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)
assert "digraph G" in graph.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 '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 '"__start__" [shape=ellipse, style=filled, fillcolor=lightblue];' in graph.source
assert '"__end__" [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 (
assert '"Tool1" [label="Tool1", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in graph.source '"Agent1" [label="Agent1", shape=box, style=filled, fillcolor=lightyellow, width=1.5, height=0.8];'
assert '"Tool2" [label="Tool2", shape=ellipse, style=filled, fillcolor=lightgreen, width=0.5, height=0.3];' in graph.source 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 (
'"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
)