chore: import upstream snapshot with attribution

haystack/components/preprocessors/hierarchical_document_splitter.py

@@ -0,0 +1,156 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import replace
+from typing import Any, Literal
+
+from haystack import Document, component, default_from_dict, default_to_dict
+from haystack.components.preprocessors import DocumentSplitter
+
+
+@component
+class HierarchicalDocumentSplitter:
+    """
+    Splits a documents into different block sizes building a hierarchical tree structure of blocks of different sizes.
+
+    The root node of the tree is the original document, the leaf nodes are the smallest blocks. The blocks in between
+    are connected such that the smaller blocks are children of the parent-larger blocks.
+
+    ## Usage example
+    ```python
+    from haystack import Document
+    from haystack.components.preprocessors import HierarchicalDocumentSplitter
+
+    doc = Document(content="This is a simple test document")
+    splitter = HierarchicalDocumentSplitter(block_sizes={3, 2}, split_overlap=0, split_by="word")
+    splitter.run([doc])
+    # >> {'documents': [Document(id=3f7..., content: 'This is a simple test document', meta: {'block_size': 0, 'parent_id': None, 'children_ids': ['5ff..', '8dc..'], 'level': 0}),
+    # >> Document(id=5ff.., content: 'This is a ', meta: {'block_size': 3, 'parent_id': '3f7..', 'children_ids': ['f19..', '52c..'], 'level': 1, 'source_id': '3f7..', 'page_number': 1, 'split_id': 0, 'split_idx_start': 0}),
+    # >> Document(id=8dc.., content: 'simple test document', meta: {'block_size': 3, 'parent_id': '3f7..', 'children_ids': ['39d..', 'e23..'], 'level': 1, 'source_id': '3f7..', 'page_number': 1, 'split_id': 1, 'split_idx_start': 10}),
+    # >> Document(id=f19.., content: 'This is ', meta: {'block_size': 2, 'parent_id': '5ff..', 'children_ids': [], 'level': 2, 'source_id': '5ff..', 'page_number': 1, 'split_id': 0, 'split_idx_start': 0}),
+    # >> Document(id=52c.., content: 'a ', meta: {'block_size': 2, 'parent_id': '5ff..', 'children_ids': [], 'level': 2, 'source_id': '5ff..', 'page_number': 1, 'split_id': 1, 'split_idx_start': 8}),
+    # >> Document(id=39d.., content: 'simple test ', meta: {'block_size': 2, 'parent_id': '8dc..', 'children_ids': [], 'level': 2, 'source_id': '8dc..', 'page_number': 1, 'split_id': 0, 'split_idx_start': 0}),
+    # >> Document(id=e23.., content: 'document', meta: {'block_size': 2, 'parent_id': '8dc..', 'children_ids': [], 'level': 2, 'source_id': '8dc..', 'page_number': 1, 'split_id': 1, 'split_idx_start': 12})]}
+    ```
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        block_sizes: set[int],
+        split_overlap: int = 0,
+        split_by: Literal["word", "sentence", "page", "passage"] = "word",
+    ) -> None:
+        """
+        Initialize HierarchicalDocumentSplitter.
+
+        :param block_sizes: Set of block sizes to split the document into. The blocks are split in descending order.
+        :param split_overlap: The number of overlapping units for each split.
+        :param split_by: The unit for splitting your documents.
+        :raises ValueError: If `block_sizes` is empty, if `split_overlap` is negative, or if `split_overlap` is
+            greater than or equal to the smallest value in `block_sizes`.
+        """
+
+        if not block_sizes:
+            raise ValueError("block_sizes must not be empty. Provide at least one block size.")
+
+        if split_overlap < 0:
+            raise ValueError("split_overlap must be greater than or equal to 0.")
+
+        smallest_block_size = min(block_sizes)
+        if split_overlap >= smallest_block_size:
+            raise ValueError(
+                f"split_overlap ({split_overlap}) must be less than the smallest value in block_sizes "
+                f"({smallest_block_size}). Reduce split_overlap or increase the smallest block size."
+            )
+
+        self.block_sizes = sorted(set(block_sizes), reverse=True)
+        self.splitters: dict[int, DocumentSplitter] = {}
+        self.split_overlap = split_overlap
+        self.split_by = split_by
+        self._build_block_sizes()
+
+    @component.output_types(documents=list[Document])
+    def run(self, documents: list[Document]) -> dict[str, list[Document]]:
+        """
+        Builds a hierarchical document structure for each document in a list of documents.
+
+        :param documents: List of Documents to split into hierarchical blocks.
+        :returns: List of HierarchicalDocument
+        """
+        hierarchical_docs = []
+        for doc in documents:
+            hierarchical_docs.extend(self.build_hierarchy_from_doc(doc))
+        return {"documents": hierarchical_docs}
+
+    def _build_block_sizes(self) -> None:
+        for block_size in self.block_sizes:
+            self.splitters[block_size] = DocumentSplitter(
+                split_length=block_size, split_overlap=self.split_overlap, split_by=self.split_by
+            )
+
+    @staticmethod
+    def _add_meta_data(document: Document) -> Document:
+        new_meta = {**document.meta, "__block_size": 0, "__parent_id": None, "__children_ids": [], "__level": 0}
+        return replace(document, meta=new_meta)
+
+    def build_hierarchy_from_doc(self, document: Document) -> list[Document]:
+        """
+        Build a hierarchical tree document structure from a single document.
+
+        Given a document, this function splits the document into hierarchical blocks of different sizes represented
+        as HierarchicalDocument objects.
+
+        :param document: Document to split into hierarchical blocks.
+        :returns:
+            List of HierarchicalDocument
+        """
+
+        root = self._add_meta_data(document)
+        current_level_nodes = [root]
+        all_docs = []
+
+        for block in self.block_sizes:
+            next_level_nodes = []
+            for doc in current_level_nodes:
+                splitted_docs = self.splitters[block].run([doc])
+                child_docs = splitted_docs["documents"]
+                # if it's only one document skip
+                if len(child_docs) == 1:
+                    next_level_nodes.append(doc)
+                    continue
+                for child_doc in child_docs:
+                    child_doc = self._add_meta_data(child_doc)
+                    child_doc.meta["__level"] = doc.meta["__level"] + 1
+                    child_doc.meta["__block_size"] = block
+                    child_doc.meta["__parent_id"] = doc.id
+                    all_docs.append(child_doc)
+                    doc.meta["__children_ids"].append(child_doc.id)
+                    next_level_nodes.append(child_doc)
+            current_level_nodes = next_level_nodes
+
+        return [root] + all_docs
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Returns a dictionary representation of the component.
+
+        :returns:
+                Serialized dictionary representation of the component.
+        """
+        return default_to_dict(
+            self, block_sizes=self.block_sizes, split_overlap=self.split_overlap, split_by=self.split_by
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "HierarchicalDocumentSplitter":
+        """
+        Deserialize this component from a dictionary.
+
+        :param data:
+            The dictionary to deserialize and create the component.
+
+        :returns:
+            The deserialized component.
+        """
+        return default_from_dict(cls, data)