fix: picklable ingest errors wrapper (#448)

Author: pawel-chmielak-deepsense

packages/ragbits-document-search/CHANGELOG.md

@@ -2,9 +2,8 @@
 
 ## Unreleased
 
-### Changed
-
-- Add support for Git source to fetch files from Git repositories
+- Introduce picklable ingest error wrapper (#448)
+- Add support for Git source to fetch files from Git repositories (#439)
 
 ## 0.10.2 (2025-03-21)
 

packages/ragbits-document-search/src/ragbits/document_search/documents/exceptions.py

@@ -1,7 +1,3 @@
-import inspect
-from typing import Any
-
-
 class SourceError(Exception):
     """
     Class for all exceptions raised by the document source.
@@ -11,19 +7,6 @@ def __init__(self, message: str) -> None:
         super().__init__(message)
         self.message = message
 
-    def __reduce__(self) -> tuple[type["SourceError"], tuple[Any, ...]]:
-        # This __reduce__ method is written in a way that it automatically handles any subclass of SourceError.
-        # It requires the subclass to have an initializer that store the arguments in the instance's state,
-        # under the same name.
-        init_params = inspect.signature(self.__class__.__init__).parameters
-
-        args = [
-            self.__getattribute__(param_name)
-            for param_name in list(init_params.keys())[1:]  # Skip 'self'
-        ]
-
-        return self.__class__, tuple(args)
-
 
 class SourceConnectionError(SourceError):
     """

packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/dummy.py

@@ -67,3 +67,25 @@ async def process(self, document_meta: DocumentMeta) -> Sequence[Element | Inter
                 document_meta=document_meta,
             )
         ]
+
+
+class DummyFailureProvider(BaseProvider):
+    """
+    This is a simple provider that always raises an exception.
+    """
+
+    SUPPORTED_DOCUMENT_TYPES = {DocumentType.TXT, DocumentType.MD, DocumentType.JPG, DocumentType.PNG}
+
+    async def process(self, document_meta: DocumentMeta) -> Sequence[Element | IntermediateElement]:
+        """
+        Process the text document.
+
+        Args:
+            document_meta: The document to process.
+
+        Raises:
+            RuntimeError: This is a dummy exception.
+        """
+        self.validate_document_type(document_meta.document_type)
+
+        raise RuntimeError("This is a dummy exception")

packages/ragbits-document-search/src/ragbits/document_search/ingestion/strategies/base.py

@@ -1,5 +1,6 @@
 import asyncio
 import random
+import traceback
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Awaitable, Callable, Iterable, Sequence
@@ -20,6 +21,31 @@
 _CallReturnT = TypeVar("_CallReturnT")
 
 
+@dataclass
+class IngestError:
+    """
+    Represents an error that occurred during the document ingest execution
+    """
+
+    type: type[Exception]
+    message: str
+    stacktrace: str
+
+    @classmethod
+    def from_exception(cls, exc: Exception) -> "IngestError":
+        """
+        Create an IngestError from an exception.
+
+        Args:
+            exc: The exception to create the IngestError from.
+
+        Returns:
+            The IngestError instance.
+        """
+        stacktrace = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__))
+        return cls(type=type(exc), message=str(exc), stacktrace=stacktrace)
+
+
 @dataclass
 class IngestDocumentResult:
     """
@@ -28,7 +54,7 @@ class IngestDocumentResult:
 
     document_uri: str
     num_elements: int = 0
-    error: BaseException | None = None
+    error: IngestError | None = None
 
 
 @dataclass

packages/ragbits-document-search/src/ragbits/document_search/ingestion/strategies/batched.py

@@ -11,6 +11,7 @@
 from ragbits.document_search.ingestion.intermediate_handlers.base import BaseIntermediateHandler
 from ragbits.document_search.ingestion.strategies.base import (
     IngestDocumentResult,
+    IngestError,
     IngestExecutionResult,
     IngestStrategy,
 )
@@ -19,7 +20,7 @@
 @dataclass
 class IngestTaskResult:
     """
-    Represents the result of the document batch ingest tast.
+    Represents the result of the document batch ingest task.
     """
 
     document_uri: str
@@ -145,18 +146,29 @@ async def _parse_batch(
             ],
             return_exceptions=True,
         )
-        return [
-            IngestDocumentResult(
-                document_uri=uri,
-                error=response,
-            )
-            if isinstance(response, BaseException)
-            else IngestTaskResult(
-                document_uri=uri,
-                elements=response,
-            )
-            for uri, response in zip(uris, responses, strict=True)
-        ]
+
+        results: list[IngestTaskResult | IngestDocumentResult] = []
+        for uri, response in zip(uris, responses, strict=True):
+            if isinstance(response, BaseException):
+                if isinstance(response, Exception):
+                    results.append(
+                        IngestDocumentResult(
+                            document_uri=uri,
+                            error=IngestError.from_exception(response),
+                        )
+                    )
+                # Handle only standard exceptions, not BaseExceptions like SystemExit, KeyboardInterrupt, etc.
+                else:
+                    raise response
+            else:
+                results.append(
+                    IngestTaskResult(
+                        document_uri=uri,
+                        elements=response,
+                    )
+                )
+
+        return results
 
     async def _enrich_batch(
         self,
@@ -184,18 +196,29 @@ async def _enrich_batch(
             ],
             return_exceptions=True,
         )
-        return [
-            IngestDocumentResult(
-                document_uri=result.document_uri,
-                error=response,
-            )
-            if isinstance(response, BaseException)
-            else IngestTaskResult(
-                document_uri=result.document_uri,
-                elements=[element for element in result.elements if isinstance(element, Element)] + response,
-            )
-            for result, response in zip(batch, responses, strict=True)
-        ]
+
+        results: list[IngestTaskResult | IngestDocumentResult] = []
+        for result, response in zip(batch, responses, strict=True):
+            if isinstance(response, BaseException):
+                if isinstance(response, Exception):
+                    results.append(
+                        IngestDocumentResult(
+                            document_uri=result.document_uri,
+                            error=IngestError.from_exception(response),
+                        )
+                    )
+                # Handle only standard exceptions, not BaseExceptions like SystemExit, KeyboardInterrupt, etc.
+                else:
+                    raise response
+            else:
+                results.append(
+                    IngestTaskResult(
+                        document_uri=result.document_uri,
+                        elements=[element for element in result.elements if isinstance(element, Element)] + response,
+                    )
+                )
+
+        return results
 
     async def _index_batch(
         self,
@@ -230,7 +253,7 @@ async def _index_batch(
             return [
                 IngestDocumentResult(
                     document_uri=result.document_uri,
-                    error=exc,
+                    error=IngestError.from_exception(exc),
                 )
                 for result in batch
             ]

packages/ragbits-document-search/src/ragbits/document_search/ingestion/strategies/sequential.py

@@ -8,6 +8,7 @@
 from ragbits.document_search.ingestion.intermediate_handlers.base import BaseIntermediateHandler
 from ragbits.document_search.ingestion.strategies.base import (
     IngestDocumentResult,
+    IngestError,
     IngestExecutionResult,
     IngestStrategy,
 )
@@ -70,7 +71,7 @@ async def __call__(
                 results.failed.append(
                     IngestDocumentResult(
                         document_uri=document_uri,
-                        error=exc,
+                        error=IngestError.from_exception(exc),
                     )
                 )
             else:

packages/ragbits-document-search/tests/unit/test_ingest_strategies.py

@@ -4,7 +4,7 @@
 from ragbits.core.vector_stores.in_memory import InMemoryVectorStore
 from ragbits.document_search.documents.document import DocumentMeta, DocumentType
 from ragbits.document_search.ingestion.document_processor import DocumentProcessorRouter
-from ragbits.document_search.ingestion.providers.dummy import DummyProvider
+from ragbits.document_search.ingestion.providers.dummy import DummyFailureProvider, DummyProvider
 from ragbits.document_search.ingestion.strategies.base import IngestStrategy
 from ragbits.document_search.ingestion.strategies.batched import BatchedIngestStrategy
 from ragbits.document_search.ingestion.strategies.ray import RayDistributedIngestStrategy
@@ -48,3 +48,27 @@ async def test_ingest_strategy_call(ingest_strategy: IngestStrategy, documents:
 
     assert len(results.successful) == len(documents)
     assert len(results.failed) == 0
+
+
+async def test_ingest_errors_are_returned(ingest_strategy: IngestStrategy, documents: list[DocumentMeta]) -> None:
+    vector_store = InMemoryVectorStore(embedder=NoopEmbedder())
+    parser_router = DocumentProcessorRouter.from_config({DocumentType.TXT: DummyFailureProvider()})
+
+    results = await ingest_strategy(
+        documents=documents,
+        vector_store=vector_store,
+        parser_router=parser_router,
+        enricher_router={},
+    )
+
+    assert len(results.successful) == 0
+    assert len(results.failed) == len(documents)
+
+    expected_error_message = "This is a dummy exception"
+    expected_error_type = RuntimeError
+    for result in results.failed:
+        assert result.error is not None
+        assert result.error.type == expected_error_type
+        assert result.error.message == expected_error_message
+        assert result.error.stacktrace.startswith("Traceback")
+        assert expected_error_message in result.error.stacktrace

packages/ragbits-document-search/tests/unit/test_source_exceptions.py

@@ -1,5 +1,3 @@
-import pickle
-
 from ragbits.document_search.documents.exceptions import (
     SourceConnectionError,
     SourceError,
@@ -14,50 +12,19 @@ def test_source_error_init():
     assert str(error) == "Test error message"
 
 
-def test_source_error_pickle():
-    original = SourceError("Test pickle message")
-    pickled = pickle.dumps(original)
-    unpickled = pickle.loads(pickled)  # noqa: S301
-
-    assert isinstance(unpickled, SourceError)
-    assert unpickled.message == "Test pickle message"
-    assert str(unpickled) == "Test pickle message"
-
-
 def test_source_connection_error_init():
     error = SourceConnectionError()
     assert error.message == "Connection error."
     assert str(error) == "Connection error."
 
 
-def test_source_connection_error_pickle():
-    original = SourceConnectionError()
-    pickled = pickle.dumps(original)
-    unpickled = pickle.loads(pickled)  # noqa: S301
-
-    assert isinstance(unpickled, SourceConnectionError)
-    assert unpickled.message == "Connection error."
-    assert str(unpickled) == "Connection error."
-
-
 def test_source_not_found_error_init():
     error = SourceNotFoundError("test-source-id")
     assert error.source_id == "test-source-id"
     assert error.message == "Source with ID test-source-id not found."
     assert str(error) == "Source with ID test-source-id not found."
 
 
-def test_source_not_found_error_pickle():
-    original = SourceNotFoundError("test-source-id")
-    pickled = pickle.dumps(original)
-    unpickled = pickle.loads(pickled)  # noqa: S301
-
-    assert isinstance(unpickled, SourceNotFoundError)
-    assert unpickled.source_id == "test-source-id"
-    assert unpickled.message == "Source with ID test-source-id not found."
-    assert str(unpickled) == "Source with ID test-source-id not found."
-
-
 def test_web_download_error_init():
     url = "https://example.com/file.pdf"
     code = 404
@@ -67,17 +34,3 @@ def test_web_download_error_init():
     assert error.code == code
     assert error.message == f"Download of {url} failed with code {code}."
     assert str(error) == f"Download of {url} failed with code {code}."
-
-
-def test_web_download_error_pickle():
-    url = "https://example.com/file.pdf"
-    code = 404
-    original = WebDownloadError(url, code)
-    pickled = pickle.dumps(original)
-    unpickled = pickle.loads(pickled)  # noqa: S301
-
-    assert isinstance(unpickled, WebDownloadError)
-    assert unpickled.url == url
-    assert unpickled.code == code
-    assert unpickled.message == f"Download of {url} failed with code {code}."
-    assert str(unpickled) == f"Download of {url} failed with code {code}."