Coverage for haystack/components/preprocessors/hierarchical_document_splitter.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from copy import deepcopy 

6from dataclasses import replace 

7from typing import Any, Literal 

8 

9from haystack import Document, component, default_from_dict, default_to_dict 

10from haystack.components.preprocessors import DocumentSplitter 

11 

12 

13@component 

14class HierarchicalDocumentSplitter: 

15 """ 

16 Splits a documents into different block sizes building a hierarchical tree structure of blocks of different sizes. 

17 

18 The root node of the tree is the original document, the leaf nodes are the smallest blocks. The blocks in between 

19 are connected such that the smaller blocks are children of the parent-larger blocks. 

20 

21 ## Usage example 

22 ```python 

23 from haystack import Document 

24 from haystack.components.preprocessors import HierarchicalDocumentSplitter 

25 

26 doc = Document(content="This is a simple test document") 

27 splitter = HierarchicalDocumentSplitter(block_sizes={3, 2}, split_overlap=0, split_by="word") 

28 splitter.run([doc]) 

29 # >> {'documents': [Document(id=3f7..., content: 'This is a simple test document', meta: {'block_size': 0, 'parent_id': None, 'children_ids': ['5ff..', '8dc..'], 'level': 0}), 

30 # >> 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}), 

31 # >> 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}), 

32 # >> 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}), 

33 # >> 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}), 

34 # >> 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}), 

35 # >> 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})]} 

36 ``` 

37 """ # noqa: E501 

38 

39 def __init__( 

40 self, 

41 block_sizes: set[int], 

42 split_overlap: int = 0, 

43 split_by: Literal["word", "sentence", "page", "passage"] = "word", 

44 ) -> None: 

45 """ 

46 Initialize HierarchicalDocumentSplitter. 

47 

48 :param block_sizes: Set of block sizes to split the document into. The blocks are split in descending order. 

49 :param split_overlap: The number of overlapping units for each split. 

50 :param split_by: The unit for splitting your documents. 

51 :raises ValueError: If `block_sizes` is empty, if `split_overlap` is negative, or if `split_overlap` is 

52 greater than or equal to the smallest value in `block_sizes`. 

53 """ 

54 

55 if not block_sizes: 

56 raise ValueError("block_sizes must not be empty. Provide at least one block size.") 

57 

58 if split_overlap < 0: 

59 raise ValueError("split_overlap must be greater than or equal to 0.") 

60 

61 smallest_block_size = min(block_sizes) 

62 if split_overlap >= smallest_block_size: 

63 raise ValueError( 

64 f"split_overlap ({split_overlap}) must be less than the smallest value in block_sizes " 

65 f"({smallest_block_size}). Reduce split_overlap or increase the smallest block size." 

66 ) 

67 

68 self.block_sizes = sorted(set(block_sizes), reverse=True) 

69 self.splitters: dict[int, DocumentSplitter] = {} 

70 self.split_overlap = split_overlap 

71 self.split_by = split_by 

72 self._build_block_sizes() 

73 

74 @component.output_types(documents=list[Document]) 

75 def run(self, documents: list[Document]) -> dict[str, list[Document]]: 

76 """ 

77 Builds a hierarchical document structure for each document in a list of documents. 

78 

79 :param documents: List of Documents to split into hierarchical blocks. 

80 :returns: List of HierarchicalDocument 

81 """ 

82 hierarchical_docs = [] 

83 for doc in documents: 

84 hierarchical_docs.extend(self.build_hierarchy_from_doc(doc)) 

85 return {"documents": hierarchical_docs} 

86 

87 def _build_block_sizes(self) -> None: 

88 for block_size in self.block_sizes: 

89 self.splitters[block_size] = DocumentSplitter( 

90 split_length=block_size, split_overlap=self.split_overlap, split_by=self.split_by 

91 ) 

92 

93 @staticmethod 

94 def _add_meta_data(document: Document) -> Document: 

95 new_meta = {**document.meta, "__block_size": 0, "__parent_id": None, "__children_ids": [], "__level": 0} 

96 return replace(document, meta=new_meta) 

97 

98 def build_hierarchy_from_doc(self, document: Document) -> list[Document]: 

99 """ 

100 Build a hierarchical tree document structure from a single document. 

101 

102 Given a document, this function splits the document into hierarchical blocks of different sizes represented 

103 as HierarchicalDocument objects. 

104 

105 :param document: Document to split into hierarchical blocks. 

106 :returns: 

107 List of HierarchicalDocument 

108 """ 

109 

110 # the root is the only node built from the caller's Document, so it is the only one that needs detaching 

111 root = self._add_meta_data(replace(document, meta=deepcopy(document.meta))) 

112 current_level_nodes = [root] 

113 all_docs = [] 

114 

115 for block in self.block_sizes: 

116 next_level_nodes = [] 

117 for doc in current_level_nodes: 

118 splitted_docs = self.splitters[block].run([doc]) 

119 child_docs = splitted_docs["documents"] 

120 # if it's only one document skip 

121 if len(child_docs) == 1: 

122 next_level_nodes.append(doc) 

123 continue 

124 for child_doc in child_docs: 

125 child_doc = self._add_meta_data(child_doc) 

126 child_doc.meta["__level"] = doc.meta["__level"] + 1 

127 child_doc.meta["__block_size"] = block 

128 child_doc.meta["__parent_id"] = doc.id 

129 all_docs.append(child_doc) 

130 doc.meta["__children_ids"].append(child_doc.id) 

131 next_level_nodes.append(child_doc) 

132 current_level_nodes = next_level_nodes 

133 

134 return [root] + all_docs 

135 

136 def to_dict(self) -> dict[str, Any]: 

137 """ 

138 Returns a dictionary representation of the component. 

139 

140 :returns: 

141 Serialized dictionary representation of the component. 

142 """ 

143 return default_to_dict( 

144 self, block_sizes=self.block_sizes, split_overlap=self.split_overlap, split_by=self.split_by 

145 ) 

146 

147 @classmethod 

148 def from_dict(cls, data: dict[str, Any]) -> "HierarchicalDocumentSplitter": 

149 """ 

150 Deserialize this component from a dictionary. 

151 

152 :param data: 

153 The dictionary to deserialize and create the component. 

154 

155 :returns: 

156 The deserialized component. 

157 """ 

158 return default_from_dict(cls, data)