fix(multimodal): smooth rough edges

Signed-off-by: Anupam Kumar <kyteinsky@gmail.com>

Author: kyteinsky

context_chat_backend/chain/ingest/doc_loader.py

@@ -20,7 +20,7 @@
 from pypdf.errors import FileNotDecryptedError as PdfFileNotDecryptedError
 from striprtf import striprtf
 
-from ...types import IndexingException, SourceItem, TaskProcException, TaskProcFatalException
+from ...types import IndexingException, SourceItem, TaskProcClientException, TaskProcException, TaskProcFatalException
 from .task_proc import (
 	OCR_TASK_TYPE,
 	SPEECH_TO_TEXT_TASK_TYPE,
@@ -35,7 +35,7 @@
 logger = logging.getLogger('ccb.doc_loader')
 PDF_IMAGES_BATCH_SIZE = 25
 PDF_IMAGES_MAX_SIZE = 100 * 1024 * 1024  # 100 MiB in bytes
-# todo: reduce the frequency of this check?
+# todo: cache it across processes?
 IS_OCR_AVAILABLE = is_task_type_available(OCR_TASK_TYPE)
 IS_STT_AVAILABLE = is_task_type_available(SPEECH_TO_TEXT_TASK_TYPE)
 
@@ -60,7 +60,7 @@ def _load_pdf(file: BytesIO, source: SourceItem) -> str:
 
 	for page in pdf_reader.pages:
 		text = page.extract_text().strip()
-		ocr_outputs = []
+		page_ocr_outputs = []
 
 		if IS_OCR_AVAILABLE:
 			for i in range(0, len(page.images), PDF_IMAGES_BATCH_SIZE):
@@ -87,12 +87,12 @@ def _load_pdf(file: BytesIO, source: SourceItem) -> str:
 				try:
 					ocr_tp_outputs = asyncio.run(do_ocr(source.userIds[0], file_ids_map.values()))  # pyright: ignore[reportArgumentType]
 					asyncio.run(delete_temp_files(file_ids_map.values()))  # pyright: ignore[reportArgumentType]
-				except TaskProcFatalException:
+				except TaskProcFatalException as e:
 					# task type is not present anymore, flip the task type indicator manually
 					# for this complete batched injest process
 					logger.warning(
 						'The OCR provider disappeared mid-operation, disabling OCR for this batch of'
-						f' documents including {source.reference}'
+						f' documents including {source.reference}: {e}'
 					)
 					IS_OCR_AVAILABLE = False
 					break
@@ -101,11 +101,11 @@ def _load_pdf(file: BytesIO, source: SourceItem) -> str:
 					continue
 
 				# for each image batch
-				ocr_outputs += ocr_tp_outputs
+				page_ocr_outputs += ocr_tp_outputs
 
-		# for each page
+		# for each page, append text then its OCR outputs to preserve document order
 		output.append(text)
-		output += ocr_outputs
+		output += page_ocr_outputs
 
 	return '\n\n'.join(output)
 
@@ -192,6 +192,7 @@ def decode_source(source: SourceItem) -> str:
 	IndexingException
 	'''
 
+	global IS_OCR_AVAILABLE
 	io_obj: BytesIO | None = None
 	try:
 		# .pot files are powerpoint templates but also plain text files,
@@ -200,16 +201,36 @@ def decode_source(source: SourceItem) -> str:
 			raise IndexingException('PowerPoint template files (.pot) are not supported')
 
 		try:
-			if IS_OCR_AVAILABLE and source.type.startswith('image/'):
-				return asyncio.run(do_ocr(source.userIds[0], [source.file_id]))[0]
-			if IS_STT_AVAILABLE and source.type.startswith('audio/'):
-				return asyncio.run(do_transcription(source.userIds[0], source.file_id))
-		except TaskProcException as e:
-			# todo: convert this to error obj return
-			# todo: short circuit all other ocr/transcription files when a fatal error arrives
-			# todo:  maybe with a global ttl, with a retryable tag
+			if source.type.startswith('image/'):
+				if IS_OCR_AVAILABLE:
+					return asyncio.run(do_ocr(source.userIds[0], [source.file_id]))[0]
+				raise IndexingException(
+					f'Image file ({source.reference}) cannot be processed since OCR task type is not present'
+				)
+			if source.type.startswith('audio/'):
+				if IS_STT_AVAILABLE:
+					return asyncio.run(do_transcription(source.userIds[0], source.file_id))
+				raise IndexingException(
+					f'Audio file ({source.reference}) cannot be processed since Speech-to-Text task type is not present'
+				)
+		except TaskProcFatalException as e:
+			logger.warning(
+				'The OCR provider disappeared mid-operation, disabling OCR for this batch of'
+				f' documents including {source.reference}: {e}'
+			)
+			IS_OCR_AVAILABLE = False
+			raise IndexingException(  # noqa: B904
+				f'Image file ({source.reference}) cannot be processed since OCR task type is not present'
+			)
+		except TaskProcClientException as e:
 			logger.warning(f'OCR task failed for source file ({source.reference}): {e}')
 			raise IndexingException(f'OCR task failed for source file ({source.reference}): {e}')  # noqa: B904
+		except TaskProcException as e:
+			logger.warning(f'OCR task failed for source file ({source.reference}), it will be retried: {e}')
+			raise IndexingException(  # noqa: B904
+				f'OCR task failed for source file ({source.reference}), it will be retried: {e}',
+				retryable=True,
+			)
 		except ValueError:
 			# should not happen
 			logger.warning(f'Unexpected ValueError for source file ({source.reference})')

context_chat_backend/chain/ingest/injest.py

@@ -15,9 +15,9 @@
 
 from ...dyn_loader import VectorDBLoader
 from ...mimetype_list import AUDIO_MIMETYPES, IMAGE_MIMETYPES, SUPPORTED_MIMETYPES
+from ...types import FILES_PROVIDER_ID, IndexingError, IndexingException, ReceivedFileItem, SourceItem, TConfig
 from ...vectordb.base import BaseVectorDB
 from ...vectordb.types import DbException, SafeDbException, UpdateAccessOp
-from ...types import FILES_PROVIDER_ID, IndexingError, IndexingException, ReceivedFileItem, SourceItem, TConfig
 from ..types import InDocument
 from .doc_loader import decode_source
 from .doc_splitter import get_splitter_for
@@ -311,7 +311,7 @@ def _sources_to_indocuments(
 def _increase_access_for_existing_sources(
 	vectordb: BaseVectorDB,
 	existing_sources: Mapping[int, SourceItem | ReceivedFileItem]
-) -> Mapping[int, IndexingError | None]:
+) -> dict[int, IndexingError | None]:
 	'''
 	update userIds for existing sources
 	allow the userIds as additional users, not as the only users
@@ -352,7 +352,7 @@ def _process_sources(
 	vectordb: BaseVectorDB,
 	config: TConfig,
 	sources: Mapping[int, SourceItem | ReceivedFileItem]
-) -> Mapping[int, IndexingError | None]:
+) -> dict[int, IndexingError | None]:
 	'''
 	Processes the sources and adds them to the vectordb.
 	Returns the list of source ids that were successfully added and those that need to be retried.

context_chat_backend/chain/ingest/task_proc.py

@@ -226,6 +226,8 @@ async def __get_task_result(user_id: str, task: Task) -> Any:
 			raise TaskProcException('Failed to parse TaskProcessing task result') from e
 
 	if task.status != 'STATUS_SUCCESSFUL':
+		# it's not possible to know if the document is unprocessable or the provider is having a temporary issue
+		# so we retry it
 		raise TaskProcException(
 			f'TaskProcessing task id {task.id} failed with status {task.status}'
 			f' after waiting {now_waiting_for} seconds',

context_chat_backend/types.py

@@ -150,6 +150,10 @@ class TaskProcFatalException(TaskProcException):
 	...
 
 
+class TaskProcClientException(TaskProcException):
+	...
+
+
 class AppRole(str, Enum):
 	NORMAL = 'normal'
 	INDEXING = 'indexing'