fix: update retriever type assertion!

We were using the default `VectorIndexRetriever` for the new pipelines and it was wrong to assert those with CustomVectorStoreRetriever.
TogetherCrew · May 20, 2024 · b1f15f4 · b1f15f4
1 parent 9a9247d
commit b1f15f4
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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)