diff --git a/ai_coach/ai_coach.py b/ai_coach/ai_coach.py
index 0799a3e..b65204f 100644
--- a/ai_coach/ai_coach.py
+++ b/ai_coach/ai_coach.py
@@ -1,4 +1,4 @@
-from xblockutils.studio_editable import StudioEditableXBlockMixin
+from xblock.utils.studio_editable import StudioEditableXBlockMixin
 from xblock.fields import Float, Integer, Scope, String
 from xblock.core import XBlock
 from xblock.completable import CompletableXBlockMixin
diff --git a/setup.py b/setup.py
index 04bb048..edc4fb7 100644
--- a/setup.py
+++ b/setup.py
@@ -4,7 +4,7 @@
 import os
 
 from pathlib import Path
-from setuptools import setup
+from setuptools import find_packages, setup
 
 
 this_directory = Path(__file__).parent
@@ -36,13 +36,13 @@ def package_data(pkg, roots):
     url='https://github.com/edly-io/ai-coach-xblock',
     license='AGPL v3',
     author='edly',
-    packages=[
-        'ai_coach',
-    ],
+    packages=find_packages(
+        include=['ai_coach', 'ai_coach.*'],
+        exclude=["*tests"],
+    ),
     install_requires=[
         'XBlock',
-        'openai==1.3.5',
-        'xblock-utils'
+        'openai==1.3.5'
     ],
     entry_points={
         'xblock.v1': [