feat: Added qdrant support for BaseEngine and added MediaWiki RAG!

docker-compose.test.yml

@@ -34,6 +34,9 @@ services:
       - D_RETRIEVER_SEARCH=7
       - COHERE_API_KEY=some_credentials
       - OPENAI_API_KEY=some_credentials2
+      - QDRANT_HOST=qdrant
+      - QDRANT_PORT=6333
+      - QDRANT_API_KEY=
     volumes:
       - ./coverage:/project/coverage
     depends_on:
@@ -43,6 +46,8 @@ services:
         condition: service_healthy
       postgres:
         condition: service_healthy
+      qdrant-healthcheck:
+        condition: service_healthy
   neo4j:
     image: "neo4j:5.9.0"
     environment:
@@ -87,3 +92,25 @@ services:
       timeout: 30s
       retries: 2
       start_period: 40s
+  qdrant:
+    image: qdrant/qdrant:v1.9.2
+    restart: always
+    container_name: qdrant
+    ports:
+      - 6333:6333
+    expose:
+      - 6333
+    volumes:
+      - ./qdrant_data:/qdrant_data
+  qdrant-healthcheck:
+    restart: always
+    image: curlimages/curl:latest
+    entrypoint: ["/bin/sh", "-c", "--", "while true; do sleep 30; done;"]
+    depends_on:
+      - qdrant
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://qdrant:6333/readyz"]
+      interval: 10s
+      timeout: 2s
+      retries: 5
+

requirements.txt

@@ -15,7 +15,7 @@ neo4j>=5.14.1, <6.0.0
 coverage>=7.3.3, <8.0.0
 pytest>=7.4.3, <8.0.0
 python-dotenv==1.0.0
-tc-hivemind-backend==1.1.5
+tc-hivemind-backend==1.2.0
 llama-index-question-gen-guidance==0.1.2
 llama-index-vector-stores-postgres==0.1.2
 celery>=5.3.6, <6.0.0

subquery.py

@@ -11,6 +11,7 @@
     DEFAULT_GUIDANCE_SUB_QUESTION_PROMPT_TMPL,
     GDriveQueryEngine,
     GitHubQueryEngine,
+    MediaWikiQueryEngine,
     NotionQueryEngine,
     prepare_discord_engine_auto_filter,
 )
@@ -25,7 +26,7 @@ def query_multiple_source(
     notion: bool = False,
     telegram: bool = False,
     github: bool = False,
-    media_wiki: bool = False,
+    mediaWiki: bool = False,
 ) -> tuple[str, list[NodeWithScore]]:
     """
     query multiple platforms and get an answer from the multiple
@@ -72,7 +73,7 @@ def query_multiple_source(
     # discourse_query_engine: BaseQueryEngine
     gdrive_query_engine: BaseQueryEngine
     notion_query_engine: BaseQueryEngine
-    # media_wiki_query_engine: BaseQueryEngine
+    mediawiki_query_engine: BaseQueryEngine
     # telegram_query_engine: BaseQueryEngine
 
     # query engine perparation
@@ -143,8 +144,20 @@ def query_multiple_source(
                 metadata=tool_metadata,
             )
         )
-    if media_wiki:
-        raise NotImplementedError
+    if mediaWiki:
+        mediawiki_query_engine = MediaWikiQueryEngine(
+            community_id=community_id
+        ).prepare()
+        tool_metadata = ToolMetadata(
+            name="WikiPedia",
+            description="Hosts articles about any information on internet",
+        )
+        query_engine_tools.append(
+            QueryEngineTool(
+                query_engine=mediawiki_query_engine,
+                metadata=tool_metadata,
+            )
+        )
 
     embed_model = CohereEmbedding()
     llm = OpenAI("gpt-3.5-turbo")

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)

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)

tests/unit/test_mediawiki_query_engine.py

@@ -0,0 +1,15 @@
+from unittest import TestCase
+
+from llama_index.core.indices.vector_store.retrievers.retriever import VectorIndexRetriever
+from utils.query_engine import MediaWikiQueryEngine
+
+
+class TestMediaWikiQueryEngine(TestCase):
+    def setUp(self) -> None:
+        community_id = "sample_community"
+        self.notion_query_engine = MediaWikiQueryEngine(community_id)
+
+    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, VectorIndexRetriever)

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)

utils/query_engine/__init__.py

@@ -1,6 +1,7 @@
 # flake8: noqa
 from .gdrive import GDriveQueryEngine
 from .github import GitHubQueryEngine
+from .media_wiki import MediaWikiQueryEngine
 from .notion import NotionQueryEngine
 from .prepare_discord_query_engine import prepare_discord_engine_auto_filter
 from .subquery_gen_prompt import DEFAULT_GUIDANCE_SUB_QUESTION_PROMPT_TMPL

utils/query_engine/base_engine.py

@@ -1,10 +1,10 @@
-from bot.retrievers.custom_retriever import CustomVectorStoreRetriever
 from bot.retrievers.utils.load_hyperparams import load_hyperparams
 from llama_index.core import VectorStoreIndex, get_response_synthesizer
-from llama_index.core.base.embeddings.base import BaseEmbedding
+from llama_index.core.indices.vector_store.retrievers.retriever import (
+    VectorIndexRetriever,
+)
 from llama_index.core.query_engine import RetrieverQueryEngine
-from tc_hivemind_backend.embeddings import CohereEmbedding
-from tc_hivemind_backend.pg_vector_access import PGVectorAccess
+from tc_hivemind_backend.qdrant_vector_access import QDrantVectorAccess
 
 
 class BaseEngine:
@@ -24,16 +24,17 @@ def __init__(self, platform_name: str, community_id: str) -> None:
         """
         self.platform_name = platform_name
         self.community_id = community_id
-        self.dbname = f"community_{self.community_id}"
+        self.collection_name = f"{self.community_id}_{platform_name}"
 
     def prepare(self, testing=False):
-        index = self._setup_vector_store_index(
+        vector_store_index = self._setup_vector_store_index(
             testing=testing,
         )
         _, similarity_top_k, _ = load_hyperparams()
 
-        retriever = CustomVectorStoreRetriever(
-            index=index, similarity_top_k=similarity_top_k
+        retriever = VectorIndexRetriever(
+            index=vector_store_index,
+            similarity_top_k=similarity_top_k,
         )
         query_engine = RetrieverQueryEngine(
             retriever=retriever,
@@ -54,27 +55,13 @@ def _setup_vector_store_index(
         testing : bool
             for testing purposes
         **kwargs :
-            table_name : str
-                to override the default table_name
-            dbname : str
-                to override the default database name
+            collection_name : str
+                to override the default collection_name
         """
-        table_name = kwargs.get("table_name", self.platform_name)
-        dbname = kwargs.get("dbname", self.dbname)
-
-        embed_model: BaseEmbedding
-        if testing:
-            from llama_index.core import MockEmbedding
-
-            embed_model = MockEmbedding(embed_dim=1024)
-        else:
-            embed_model = CohereEmbedding()
-
-        pg_vector = PGVectorAccess(
-            table_name=table_name,
-            dbname=dbname,
+        collection_name = kwargs.get("collection_name", self.collection_name)
+        qdrant_vector = QDrantVectorAccess(
+            collection_name=collection_name,
             testing=testing,
-            embed_model=embed_model,
         )
-        index = pg_vector.load_index()
+        index = qdrant_vector.load_index()
         return index

utils/query_engine/media_wiki.py

@@ -0,0 +1,7 @@
+from utils.query_engine.base_engine import BaseEngine
+
+
+class MediaWikiQueryEngine(BaseEngine):
+    def __init__(self, community_id: str) -> None:
+        platform_name = "mediawiki"
+        super().__init__(platform_name, community_id)