Final fixes

Querent-ai · Apr 18, 2024 · f7b8f74 · f7b8f74
1 parent 48d9abe
commit f7b8f74
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 103 deletions.
diff --git a/querent/common/types/querent_event.py b/querent/common/types/querent_event.py
@@ -1,8 +1,7 @@
 from typing import Any
-from enum import Enum
 
 
-class EventType(Enum):
+class EventType:
     Graph = "Graph"
     Vector = "Vector"
     Terminate="Terminate"

diff --git a/querent/core/base_engine.py b/querent/core/base_engine.py
@@ -180,7 +180,7 @@ async def _listen_for_state_changes(self):
                 if new_state.payload == "Terminate":
                     break
                 new_state = {
-                    "event_type": str(new_state.event_type.value),
+                    "event_type": str(new_state.event_type),
                     "timestamp": new_state.timestamp,
                     "payload": new_state.payload,
                     "file": new_state.file,
@@ -221,7 +221,7 @@ async def _inner_worker():
                 none_counter = 0
                 while not self.termination_event.is_set():
                     retries = 0
-                    await asyncio.sleep(30)
+                    await asyncio.sleep(5)
                     data = await self.input_queue.get()
                     try:
                         if isinstance(data, IngestedMessages):

diff --git a/querent/ingestors/pdfs/pdf_ingestor_v1.py b/querent/ingestors/pdfs/pdf_ingestor_v1.py
@@ -65,20 +65,19 @@ async def ingest(
                         data=None,
                         error=None,
                         doc_source=chunk_bytes.doc_source,
-                        is_token_stream=True,
                     )
                 collected_bytes += chunk_bytes.data
         except Exception as e:
             # at the queue level, we can sample out the error
-            yield IngestedTokens(file=current_file, data=None, error=f"Exception: {e}", doc_source=chunk_bytes.doc_source, is_token_stream=True)
+            yield IngestedTokens(file=current_file, data=None, error=f"Exception: {e}", doc_source=chunk_bytes.doc_source)
         finally:
             # process the last file
             try:
                 async for page_text in self.extract_and_process_pdf(
                     CollectedBytes(file=current_file, data=collected_bytes), chunk_bytes.doc_source
                 ):
                     yield page_text
-                yield IngestedTokens(file=current_file, data=None, error=None, doc_source=chunk_bytes.doc_source, is_token_stream=True)
+                yield IngestedTokens(file=current_file, data=None, error=None, doc_source=chunk_bytes.doc_source)
             except Exception as exc:
                 yield IngestedTokens(
                     file=current_file,
@@ -109,7 +108,6 @@ async def extract_and_process_pdf(
                     data=processed_text,
                     error=collected_bytes.error,
                     doc_source=doc_source,
-                    is_token_stream=True,
                 )
                 # async for image_result in self.extract_images_and_ocr(
                 #     page,

diff --git a/querent/storage/s3-data-management.py b/querent/storage/s3-data-management.py