Move IngestionPipeline docstore document insertion after transformations (#19849)

* Move IngestionPipeline docstore document insertion after transformations

* Fix type check

---------

Co-authored-by: Florian Bacho <[email protected]>

Author: Florian-BACHO

llama-index-core/llama_index/core/ingestion/pipeline.py

@@ -380,7 +380,6 @@ def _prepare_inputs(
     def _handle_duplicates(
         self,
         nodes: Sequence[BaseNode],
-        store_doc_text: bool = True,
     ) -> Sequence[BaseNode]:
         """Handle docstore duplicates by checking all hashes."""
         assert self.docstore is not None
@@ -394,14 +393,11 @@ def _handle_duplicates(
                 nodes_to_run.append(node)
                 current_hashes.append(node.hash)
 
-        self.docstore.add_documents(nodes_to_run, store_text=store_doc_text)
-
         return nodes_to_run
 
     def _handle_upserts(
         self,
         nodes: Sequence[BaseNode],
-        store_doc_text: bool = True,
     ) -> Sequence[BaseNode]:
         """Handle docstore upserts by checking hashes and ids."""
         assert self.docstore is not None
@@ -437,11 +433,7 @@ def _handle_upserts(
                 if self.vector_store is not None:
                     self.vector_store.delete(ref_doc_id)
 
-        nodes_to_run = list(deduped_nodes_to_run.values())
-        self.docstore.set_document_hashes({n.id_: n.hash for n in nodes_to_run})
-        self.docstore.add_documents(nodes_to_run, store_text=store_doc_text)
-
-        return nodes_to_run
+        return list(deduped_nodes_to_run.values())
 
     @staticmethod
     def _node_batcher(
@@ -452,6 +444,23 @@ def _node_batcher(
         for i in range(0, len(nodes), batch_size):
             yield nodes[i : i + batch_size]
 
+    def _update_docstore(
+        self, nodes: Sequence[BaseNode], store_doc_text: bool = True
+    ) -> None:
+        """Update the document store with the given nodes."""
+        assert self.docstore is not None
+
+        if self.docstore_strategy in (
+            DocstoreStrategy.UPSERTS,
+            DocstoreStrategy.UPSERTS_AND_DELETE,
+        ):
+            self.docstore.set_document_hashes({n.id_: n.hash for n in nodes})
+            self.docstore.add_documents(nodes, store_text=store_doc_text)
+        elif self.docstore_strategy == DocstoreStrategy.DUPLICATES_ONLY:
+            self.docstore.add_documents(nodes, store_text=store_doc_text)
+        else:
+            raise ValueError(f"Invalid docstore strategy: {self.docstore_strategy}")
+
     @dispatcher.span
     def run(
         self,
@@ -493,13 +502,9 @@ def run(
                 DocstoreStrategy.UPSERTS,
                 DocstoreStrategy.UPSERTS_AND_DELETE,
             ):
-                nodes_to_run = self._handle_upserts(
-                    input_nodes, store_doc_text=store_doc_text
-                )
+                nodes_to_run = self._handle_upserts(input_nodes)
             elif self.docstore_strategy == DocstoreStrategy.DUPLICATES_ONLY:
-                nodes_to_run = self._handle_duplicates(
-                    input_nodes, store_doc_text=store_doc_text
-                )
+                nodes_to_run = self._handle_duplicates(input_nodes)
             else:
                 raise ValueError(f"Invalid docstore strategy: {self.docstore_strategy}")
         elif self.docstore is not None and self.vector_store is None:
@@ -515,10 +520,7 @@ def run(
                     "Switching to duplicates_only strategy."
                 )
                 self.docstore_strategy = DocstoreStrategy.DUPLICATES_ONLY
-            nodes_to_run = self._handle_duplicates(
-                input_nodes, store_doc_text=store_doc_text
-            )
-
+            nodes_to_run = self._handle_duplicates(input_nodes)
         else:
             nodes_to_run = input_nodes
 
@@ -564,6 +566,9 @@ def run(
             if nodes_with_embeddings:
                 self.vector_store.add(nodes_with_embeddings)
 
+        if self.docstore is not None:
+            self._update_docstore(nodes_to_run, store_doc_text=store_doc_text)
+
         return nodes
 
     # ------ async methods ------

llama-index-core/tests/ingestion/test_pipeline.py

@@ -1,12 +1,14 @@
 from multiprocessing import cpu_count
+from typing import Sequence, Any
 
+import pytest
 from llama_index.core.embeddings.mock_embed_model import MockEmbedding
 from llama_index.core.extractors import KeywordExtractor
 from llama_index.core.ingestion.pipeline import IngestionPipeline
 from llama_index.core.llms.mock import MockLLM
 from llama_index.core.node_parser import SentenceSplitter, MarkdownElementNodeParser
 from llama_index.core.readers import ReaderConfig, StringIterableReader
-from llama_index.core.schema import Document
+from llama_index.core.schema import Document, TransformComponent, BaseNode
 from llama_index.core.storage.docstore import SimpleDocumentStore
 
 
@@ -260,3 +262,27 @@ def test_pipeline_parallel() -> None:
     assert len(nodes) == 20
     assert pipeline.docstore is not None
     assert len(pipeline.docstore.docs) == 2
+
+
+def test_pipeline_with_transform_error() -> None:
+    class RaisingTransform(TransformComponent):
+        def __call__(
+            self, nodes: Sequence[BaseNode], **kwargs: Any
+        ) -> Sequence[BaseNode]:
+            raise RuntimeError
+
+    document1 = Document.example()
+    document1.id_ = "1"
+
+    pipeline = IngestionPipeline(
+        transformations=[
+            SentenceSplitter(chunk_size=25, chunk_overlap=0),
+            RaisingTransform(),
+        ],
+        docstore=SimpleDocumentStore(),
+    )
+
+    with pytest.raises(RuntimeError):
+        pipeline.run(documents=[document1])
+
+    assert pipeline.docstore.get_node("1", raise_error=False) is None