Persist AckQueue for indexing

Author: Roey7

app/data_sources/bookstack.py

@@ -6,7 +6,7 @@
 from data_source_api.base_data_source import BaseDataSource, ConfigField, HTMLInputType
 from data_source_api.exception import InvalidDataSourceConfig
 from data_source_api.utils import parse_with_workers
-from indexing_queue import IndexingQueue
+from index_queue import IndexQueue
 from parsers.html import html_to_text
 from pydantic import BaseModel
 from requests import Session, HTTPError
@@ -194,10 +194,10 @@ def _parse_documents_worker(self, raw_docs: List[Dict]):
                                              type=DocumentType.DOCUMENT))
             if len(parsed_docs) >= 50:
                 total_fed += len(parsed_docs)
-                IndexingQueue.get().feed(docs=parsed_docs)
+                IndexQueue.get_instance().put(docs=parsed_docs)
                 parsed_docs = []
 
-        IndexingQueue.get().feed(docs=parsed_docs)
+        IndexQueue.get_instance().put(docs=parsed_docs)
         total_fed += len(parsed_docs)
         if total_fed > 0:
             logging.info(f"Worker fed {total_fed} documents")

app/data_sources/confluence.py

@@ -8,7 +8,7 @@
 from data_source_api.base_data_source import BaseDataSource, ConfigField, HTMLInputType
 from data_source_api.exception import InvalidDataSourceConfig
 from data_source_api.utils import parse_with_workers
-from indexing_queue import IndexingQueue
+from index_queue import IndexQueue
 from parsers.html import html_to_text
 from pydantic import BaseModel
 
@@ -115,10 +115,10 @@ def _parse_documents_worker(self, raw_docs: List[Dict]):
                                              type=DocumentType.DOCUMENT))
             if len(parsed_docs) >= 50:
                 total_fed += len(parsed_docs)
-                IndexingQueue.get().feed(docs=parsed_docs)
+                IndexQueue.get_instance().put(docs=parsed_docs)
                 parsed_docs = []
 
-        IndexingQueue.get().feed(docs=parsed_docs)
+        IndexQueue.get_instance().put(docs=parsed_docs)
         total_fed += len(parsed_docs)
         if total_fed > 0:
             logging.info(f'Worker fed {total_fed} documents')

app/data_sources/google_drive.py

@@ -17,7 +17,7 @@
 from data_source_api.base_data_source import BaseDataSource, ConfigField, HTMLInputType
 from data_source_api.basic_document import BasicDocument, DocumentType, FileType
 from data_source_api.exception import InvalidDataSourceConfig, KnownException
-from indexing_queue import IndexingQueue
+from index_queue import IndexQueue
 from parsers.html import html_to_text
 from parsers.pptx import pptx_to_text
 from parsers.docx import docx_to_html
@@ -188,7 +188,7 @@ def _index_files_from_drive(self, drive) -> List[dict]:
                 file_type=FileType.from_mime_type(mime_type=file['mimeType'])
             ))
 
-        IndexingQueue.get().feed(documents)
+        IndexQueue.get_instance().put(documents)
 
     def _get_all_drives(self) -> List[dict]:
         return [{'name': 'My Drive', 'id': None}] \

app/data_sources/mattermost.py

@@ -11,7 +11,7 @@
 from data_source_api.basic_document import BasicDocument, DocumentType
 from data_source_api.exception import InvalidDataSourceConfig
 from data_source_api.utils import parse_with_workers
-from indexing_queue import IndexingQueue
+from index_queue import IndexQueue
 
 logger = logging.getLogger(__name__)
 
@@ -150,7 +150,7 @@ def _feed_channel(self, channel: MattermostChannel):
                         parsed_posts.append(last_message)
                         if len(parsed_posts) >= MattermostDataSource.FEED_BATCH_SIZE:
                             total_fed += len(parsed_posts)
-                            IndexingQueue.get().feed(docs=parsed_posts)
+                            IndexQueue.get_instance().put(docs=parsed_posts)
                             parsed_posts = []
 
                 author_image_url = f"{self._get_mattermost_url()}/api/v4/users/{post['user_id']}/image?_=0"
@@ -175,7 +175,7 @@ def _feed_channel(self, channel: MattermostChannel):
                 break
             page += 1
 
-        IndexingQueue.get().feed(docs=parsed_posts)
+        IndexQueue.get_instance().put(docs=parsed_posts)
         total_fed += len(parsed_posts)
 
         if len(parsed_posts) > 0:

app/data_sources/rocketchat.py

@@ -9,7 +9,7 @@
 from data_source_api.base_data_source import BaseDataSource, ConfigField, HTMLInputType
 from data_source_api.basic_document import DocumentType, BasicDocument
 from data_source_api.exception import InvalidDataSourceConfig
-from indexing_queue import IndexingQueue
+from index_queue import IndexQueue
 
 
 @dataclass
@@ -191,7 +191,7 @@ def _feed_new_documents(self) -> None:
                 documents.append(last_msg)
 
         logging.info(f"Total messages : {len(documents)}")
-        IndexingQueue.get().feed(docs=documents)
+        IndexQueue.get_instance().put(docs=documents)
 
 
 if __name__ == "__main__":

app/data_sources/slack.py

@@ -10,7 +10,7 @@
 from data_source_api.base_data_source import BaseDataSource, ConfigField, HTMLInputType
 from data_source_api.basic_document import DocumentType, BasicDocument
 from data_source_api.utils import parse_with_workers
-from indexing_queue import IndexingQueue
+from index_queue import IndexQueue
 
 logger = logging.getLogger(__name__)
 
@@ -125,7 +125,7 @@ def _feed_conversation(self, conv):
                     documents.append(last_msg)
                     if len(documents) == SlackDataSource.FEED_BATCH_SIZE:
                         total_fed += SlackDataSource.FEED_BATCH_SIZE
-                        IndexingQueue.get().feed(docs=documents)
+                        IndexQueue.get_instance().put(docs=documents)
                         documents = []
 
             timestamp = message['ts']
@@ -141,7 +141,7 @@ def _feed_conversation(self, conv):
         if last_msg is not None:
             documents.append(last_msg)
 
-        IndexingQueue.get().feed(docs=documents)
+        IndexQueue.get_instance().put(docs=documents)
         total_fed += len(documents)
         if total_fed > 0:
             logger.info(f'Slack worker fed {total_fed} documents')
@@ -151,7 +151,7 @@ def _fetch_conversation_messages(self, conv):
         cursor = None
         has_more = True
         last_index_unix = self._last_index_time.timestamp()
-        logger.info(f'Fetching messages for conversation {conv.name} since {last_index_unix}')
+        logger.info(f'Fetching messages for conversation {conv.name}')
 
         while has_more:
             response = self._slack.conversations_history(channel=conv.id, oldest=str(last_index_unix),

app/index_queue.py

@@ -0,0 +1,56 @@
+import threading
+from dataclasses import dataclass
+from typing import List
+
+from persistqueue import SQLiteAckQueue
+
+from data_source_api.basic_document import BasicDocument
+from paths import SQLITE_TASKS_PATH
+
+
+@dataclass
+class IndexQueueItem:
+    queue_item_id: int
+    doc: BasicDocument
+
+
+class IndexQueue(SQLiteAckQueue):
+    __instance = None
+    __lock = threading.Lock()
+
+    @classmethod
+    def get_instance(cls):
+        with cls.__lock:
+            if cls.__instance is None:
+                cls.__instance = cls()
+        return cls.__instance
+
+    def __init__(self):
+        if IndexQueue.__instance is not None:
+            raise RuntimeError("Queue is a singleton, use .get() to get the instance")
+
+        self.condition = threading.Condition()
+        super().__init__(path=SQLITE_TASKS_PATH, multithreading=True, name="index")
+
+    def put_single(self, doc: BasicDocument):
+        self.put([doc])
+
+    def put(self, docs: List[BasicDocument]):
+        with self.condition:
+            for doc in docs:
+                super().put(doc)
+
+            self.condition.notify_all()
+
+    def consume_all(self, max_docs=5000, timeout=1) -> List[IndexQueueItem]:
+        with self.condition:
+            self.condition.wait(timeout=timeout)
+
+            queue_items = []
+            count = 0
+            while not super().empty() and count < max_docs:
+                raw_items = super().get(raw=True)
+                queue_items.append(IndexQueueItem(queue_item_id=raw_items['pqid'], doc=raw_items['data']))
+                count += 1
+
+            return queue_items

app/indexing/background_indexer.py

@@ -1,10 +1,14 @@
 import logging
 import threading
+from typing import List
 
-from indexing_queue import IndexingQueue
+from index_queue import IndexQueue
 from indexing.index_documents import Indexer
 
 
+logger = logging.getLogger()
+
+
 class BackgroundIndexer:
     _thread = None
     _stop_event = threading.Event()
@@ -31,17 +35,26 @@ def stop(cls):
 
     @staticmethod
     def run():
-        logger = logging.getLogger()
-        docs_queue_instance = IndexingQueue.get()
+        docs_queue_instance = IndexQueue.get_instance()
         logger.info(f'Background indexer started...')
 
         while not BackgroundIndexer._stop_event.is_set():
-            docs_chunk = docs_queue_instance.consume_all()
-            if not docs_chunk:
+            queue_items = docs_queue_instance.consume_all()
+            if not queue_items:
                 continue
 
-            BackgroundIndexer._currently_indexing_count = len(docs_chunk)
-            logger.info(f'Got chunk of {len(docs_chunk)} documents')
-            Indexer.index_documents(docs_chunk)
-            logger.info(f'Finished indexing chunk of {len(docs_chunk)} documents')
-            BackgroundIndexer._currently_indexing_count = 0
+            BackgroundIndexer._currently_indexing_count = len(queue_items)
+            logger.info(f'Got chunk of {len(queue_items)} documents')
+
+            docs = [doc.doc for doc in queue_items]
+            Indexer.index_documents(docs)
+            BackgroundIndexer._ack_chunk(docs_queue_instance, [doc.queue_item_id for doc in queue_items])
+
+    @staticmethod
+    def _ack_chunk(queue: IndexQueue, ids: List[int]):
+        logger.info(f'Finished indexing chunk of {len(ids)} documents')
+        for item_id in ids:
+            queue.ack(id=item_id)
+
+        logger.info(f'Acked {len(ids)} documents.')
+        BackgroundIndexer._currently_indexing_count = 0

app/indexing_queue.py

@@ -20,7 +20,7 @@
 from indexing.background_indexer import BackgroundIndexer
 from indexing.bm25_index import Bm25Index
 from indexing.faiss_index import FaissIndex
-from indexing_queue import IndexingQueue
+from index_queue import IndexQueue
 from paths import UI_PATH
 from schemas import DataSource
 from schemas.data_source_type import DataSourceType
@@ -138,7 +138,7 @@ class Status:
         docs_left_to_index: int
 
     return Status(docs_in_indexing=BackgroundIndexer.get_currently_indexing(),
-                  docs_left_to_index=IndexingQueue.get().get_how_many_left())
+                  docs_left_to_index=IndexQueue.get_instance().qsize())
 
 
 @app.post("/clear-index")
@@ -157,6 +157,6 @@ async def clear_index():
     logger.warning(f"Failed to mount UI (you probably need to build it): {e}")
 
 
-# if __name__ == '__main__':
-#     import uvicorn
-#     uvicorn.run("main:app", host="localhost", port=8000)
+if __name__ == '__main__':
+    import uvicorn
+    uvicorn.run("main:app", host="localhost", port=8000)