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