From f7c594da083477e5eefbbf8cfe78cfda052bdaae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mart=C3=ADn=20Bravo?=
 <123977407+MartinEBravo@users.noreply.github.com>
Date: Tue, 18 Mar 2025 09:53:48 +0100
Subject: [PATCH] feat: add visualization functions for agent graphs

---
 src/agents/{ => extensions}/visualizations.py | 2 +-
 tests/test_visualizations.py                  | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)
 rename src/agents/{ => extensions}/visualizations.py (98%)

diff --git a/src/agents/visualizations.py b/src/agents/extensions/visualizations.py
similarity index 98%
rename from src/agents/visualizations.py
rename to src/agents/extensions/visualizations.py
index 42019d4..c020708 100644
--- a/src/agents/visualizations.py
+++ b/src/agents/extensions/visualizations.py
@@ -1,6 +1,6 @@
 import graphviz
 
-from src.agents.agent import Agent
+from agents import Agent
 
 
 def get_main_graph(agent: Agent) -> str:
diff --git a/tests/test_visualizations.py b/tests/test_visualizations.py
index d813978..ac02cee 100644
--- a/tests/test_visualizations.py
+++ b/tests/test_visualizations.py
@@ -3,8 +3,8 @@ from unittest.mock import Mock
 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
+from agents import Agent
+from agents.extensions.visualizations import draw_graph, get_all_edges, get_all_nodes, get_main_graph
 
 
 @pytest.fixture