diff --git a/tests/unit/test_gdrive_query_engine.py b/tests/unit/test_gdrive_query_engine.py
index 706d211..7ab4e65 100644
--- a/tests/unit/test_gdrive_query_engine.py
+++ b/tests/unit/test_gdrive_query_engine.py
@@ -1,6 +1,6 @@
 from unittest import TestCase
 
-from bot.retrievers.custom_retriever import CustomVectorStoreRetriever
+from llama_index.core.indices.vector_store.retrievers.retriever import VectorIndexRetriever
 from utils.query_engine import GDriveQueryEngine
 
 
@@ -12,4 +12,4 @@ def setUp(self) -> None:
     def test_prepare_engine(self):
         gdrive_query_engine = self.gdrive_query_engine.prepare(testing=True)
         print(gdrive_query_engine.__dict__)
-        self.assertIsInstance(gdrive_query_engine.retriever, CustomVectorStoreRetriever)
+        self.assertIsInstance(gdrive_query_engine.retriever, VectorIndexRetriever)
diff --git a/tests/unit/test_github_query_engine.py b/tests/unit/test_github_query_engine.py
index 1361398..b3b657f 100644
--- a/tests/unit/test_github_query_engine.py
+++ b/tests/unit/test_github_query_engine.py
@@ -1,6 +1,6 @@
 from unittest import TestCase
 
-from bot.retrievers.custom_retriever import CustomVectorStoreRetriever
+from llama_index.core.indices.vector_store.retrievers.retriever import VectorIndexRetriever
 from utils.query_engine import GitHubQueryEngine
 
 
@@ -12,4 +12,4 @@ def setUp(self) -> None:
     def test_prepare_engine(self):
         github_query_engine = self.github_query_engine.prepare(testing=True)
         print(github_query_engine.__dict__)
-        self.assertIsInstance(github_query_engine.retriever, CustomVectorStoreRetriever)
+        self.assertIsInstance(github_query_engine.retriever, VectorIndexRetriever)
diff --git a/tests/unit/test_mediawiki_query_engine.py b/tests/unit/test_mediawiki_query_engine.py
index a909bc2..303aa63 100644
--- a/tests/unit/test_mediawiki_query_engine.py
+++ b/tests/unit/test_mediawiki_query_engine.py
@@ -1,6 +1,6 @@
 from unittest import TestCase
 
-from bot.retrievers.custom_retriever import CustomVectorStoreRetriever
+from llama_index.core.indices.vector_store.retrievers.retriever import VectorIndexRetriever
 from utils.query_engine import MediaWikiQueryEngine
 
 
@@ -12,4 +12,4 @@ def setUp(self) -> None:
     def test_prepare_engine(self):
         notion_query_engine = self.notion_query_engine.prepare(testing=True)
         print(notion_query_engine.__dict__)
-        self.assertIsInstance(notion_query_engine.retriever, CustomVectorStoreRetriever)
+        self.assertIsInstance(notion_query_engine.retriever, VectorIndexRetriever)
diff --git a/tests/unit/test_notion_query_engine.py b/tests/unit/test_notion_query_engine.py
index caaf903..57f2b96 100644
--- a/tests/unit/test_notion_query_engine.py
+++ b/tests/unit/test_notion_query_engine.py
@@ -1,6 +1,6 @@
 from unittest import TestCase
 
-from bot.retrievers.custom_retriever import CustomVectorStoreRetriever
+from llama_index.core.indices.vector_store.retrievers.retriever import VectorIndexRetriever
 from utils.query_engine import NotionQueryEngine
 
 
@@ -12,4 +12,4 @@ def setUp(self) -> None:
     def test_prepare_engine(self):
         notion_query_engine = self.notion_query_engine.prepare(testing=True)
         print(notion_query_engine.__dict__)
-        self.assertIsInstance(notion_query_engine.retriever, CustomVectorStoreRetriever)
+        self.assertIsInstance(notion_query_engine.retriever, VectorIndexRetriever)