Coverage for haystack/components/retrievers/auto_merging_retriever.py: 100%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from collections import defaultdict 

6from typing import Any 

7 

8from haystack import Document, component, default_from_dict, default_to_dict 

9from haystack.document_stores.types import DocumentStore 

10 

11 

12@component 

13class AutoMergingRetriever: 

14 """ 

15 A retriever which returns parent documents of the matched leaf nodes documents, based on a threshold setting. 

16 

17 The AutoMergingRetriever assumes you have a hierarchical tree structure of documents, where the leaf nodes 

18 are indexed in a document store. See the HierarchicalDocumentSplitter for more information on how to create 

19 such a structure. During retrieval, if the number of matched leaf documents below the same parent is 

20 higher than a defined threshold, the retriever will return the parent document instead of the individual leaf 

21 documents. 

22 

23 The rational is, given that a paragraph is split into multiple chunks represented as leaf documents, and if for 

24 a given query, multiple chunks are matched, the whole paragraph might be more informative than the individual 

25 chunks alone. 

26 

27 Currently the AutoMergingRetriever can only be used by the following DocumentStores: 

28 - [AstraDB](https://haystack.deepset.ai/integrations/astradb) 

29 - [ElasticSearch](https://haystack.deepset.ai/docs/latest/documentstore/elasticsearch) 

30 - [OpenSearch](https://haystack.deepset.ai/docs/latest/documentstore/opensearch) 

31 - [PGVector](https://haystack.deepset.ai/docs/latest/documentstore/pgvector) 

32 - [Qdrant](https://haystack.deepset.ai/docs/latest/documentstore/qdrant) 

33 

34 ```python 

35 from haystack import Document 

36 from haystack.components.preprocessors import HierarchicalDocumentSplitter 

37 from haystack.components.retrievers.auto_merging_retriever import AutoMergingRetriever 

38 from haystack.document_stores.in_memory import InMemoryDocumentStore 

39 

40 # create a hierarchical document structure with 3 levels, where the parent document has 3 children 

41 text = "The sun rose early in the morning. It cast a warm glow over the trees. Birds began to sing." 

42 original_document = Document(content=text) 

43 builder = HierarchicalDocumentSplitter(block_sizes={10, 3}, split_overlap=0, split_by="word") 

44 docs = builder.run([original_document])["documents"] 

45 

46 # store level-1 parent documents and initialize the retriever 

47 doc_store_parents = InMemoryDocumentStore() 

48 for doc in docs: 

49 if doc.meta["__children_ids"] and doc.meta["__level"] in [0,1]: # store the root document and level 1 documents 

50 doc_store_parents.write_documents([doc]) 

51 

52 retriever = AutoMergingRetriever(doc_store_parents, threshold=0.5) 

53 

54 # assume we retrieved 2 leaf docs from the same parent, the parent document should be returned, 

55 # since it has 3 children and the threshold=0.5, and we retrieved 2 children (2/3 > 0.66(6)) 

56 leaf_docs = [doc for doc in docs if not doc.meta["__children_ids"]] 

57 retrieved_docs = retriever.run(leaf_docs[4:6]) 

58 print(retrieved_docs["documents"]) 

59 # [Document(id=538..), 

60 # content: 'warm glow over the trees. Birds began to sing.', 

61 # meta: {'block_size': 10, 'parent_id': '835..', 'children_ids': ['c17...', '3ff...', '352...'], 'level': 1, 'source_id': '835...', 

62 # 'page_number': 1, 'split_id': 1, 'split_idx_start': 45})]} 

63 ``` 

64 """ # noqa: E501 

65 

66 def __init__(self, document_store: DocumentStore, threshold: float = 0.5) -> None: 

67 """ 

68 Initialize the AutoMergingRetriever. 

69 

70 :param document_store: DocumentStore from which to retrieve the parent documents 

71 :param threshold: Threshold to decide whether the parent instead of the individual documents is returned 

72 """ 

73 

74 if not 0 < threshold < 1: 

75 raise ValueError("The threshold parameter must be between 0 and 1.") 

76 

77 self.document_store = document_store 

78 self.threshold = threshold 

79 

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

81 """ 

82 Serializes the component to a dictionary. 

83 

84 :returns: 

85 Dictionary with serialized data. 

86 """ 

87 return default_to_dict(self, document_store=self.document_store, threshold=self.threshold) 

88 

89 @classmethod 

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

91 """ 

92 Deserializes the component from a dictionary. 

93 

94 :param data: 

95 Dictionary with serialized data. 

96 :returns: 

97 An instance of the component. 

98 """ 

99 return default_from_dict(cls, data) 

100 

101 @staticmethod 

102 def _check_valid_documents(matched_leaf_documents: list[Document]) -> None: 

103 # check if the matched leaf documents have the required meta fields 

104 if not all(doc.meta.get("__parent_id") for doc in matched_leaf_documents): 

105 raise ValueError("The matched leaf documents do not have the required meta field '__parent_id'") 

106 

107 # HierarchicalDocumentSplitter uses __level=0 and __block_size=0 for the root, so these 

108 # fields must be checked for presence rather than truthiness. 

109 if not all("__level" in doc.meta for doc in matched_leaf_documents): 

110 raise ValueError("The matched leaf documents do not have the required meta field '__level'") 

111 

112 if not all("__block_size" in doc.meta for doc in matched_leaf_documents): 

113 raise ValueError("The matched leaf documents do not have the required meta field '__block_size'") 

114 

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

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

117 """ 

118 Run the AutoMergingRetriever. 

119 

120 Recursively groups documents by their parents and merges them if they meet the threshold, 

121 continuing up the hierarchy until no more merges are possible. 

122 

123 :param documents: List of leaf documents that were matched by a retriever 

124 :returns: 

125 List of documents (could be a mix of different hierarchy levels) 

126 """ 

127 

128 AutoMergingRetriever._check_valid_documents(documents) 

129 

130 def _get_parent_doc(parent_id: str) -> Document: 

131 parent_docs = self.document_store.filter_documents({"field": "id", "operator": "==", "value": parent_id}) 

132 if len(parent_docs) != 1: 

133 raise ValueError(f"Expected 1 parent document with id {parent_id}, found {len(parent_docs)}") 

134 

135 parent_doc = parent_docs[0] 

136 if not parent_doc.meta.get("__children_ids"): 

137 raise ValueError(f"Parent document with id {parent_id} does not have any children.") 

138 

139 return parent_doc 

140 

141 def _try_merge_level(docs_to_merge: list[Document], docs_to_return: list[Document]) -> list[Document]: 

142 parent_doc_id_to_child_docs: dict[str, list[Document]] = defaultdict(list) # to group documents by parent 

143 

144 for doc in docs_to_merge: 

145 if doc.meta.get("__parent_id"): # only docs that have parents 

146 parent_doc_id_to_child_docs[doc.meta["__parent_id"]].append(doc) 

147 else: 

148 docs_to_return.append(doc) # keep docs that have no parents 

149 

150 # Process each parent group 

151 merged_docs = [] 

152 for parent_doc_id, child_docs in parent_doc_id_to_child_docs.items(): 

153 parent_doc = _get_parent_doc(parent_doc_id) 

154 

155 # Calculate merge score 

156 score = len(child_docs) / len(parent_doc.meta["__children_ids"]) 

157 if score > self.threshold: 

158 merged_docs.append(parent_doc) # Merge into parent 

159 else: 

160 docs_to_return.extend(child_docs) # Keep children separate 

161 

162 # if no new merges were made, we're done 

163 if not merged_docs: 

164 return merged_docs + docs_to_return 

165 

166 # Recursively try to merge the next level 

167 return _try_merge_level(merged_docs, docs_to_return) 

168 

169 return {"documents": _try_merge_level(documents, [])} 

170 

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

172 async def run_async(self, documents: list[Document]) -> dict[str, list[Document]]: 

173 """ 

174 Asynchronously run the AutoMergingRetriever. 

175 

176 Recursively groups documents by their parents and merges them if they meet the threshold, 

177 continuing up the hierarchy until no more merges are possible. 

178 

179 :param documents: List of leaf documents that were matched by a retriever 

180 :returns: 

181 List of documents (could be a mix of different hierarchy levels) 

182 """ 

183 

184 AutoMergingRetriever._check_valid_documents(documents) 

185 

186 async def _get_parent_doc(parent_id: str) -> Document: 

187 # 'ignore' since filter_documents_async is not defined in the Protocol but exists in the implementations 

188 parent_docs = await self.document_store.filter_documents_async( # type: ignore[attr-defined] 

189 {"field": "id", "operator": "==", "value": parent_id} 

190 ) 

191 if len(parent_docs) != 1: 

192 raise ValueError(f"Expected 1 parent document with id {parent_id}, found {len(parent_docs)}") 

193 

194 parent_doc = parent_docs[0] 

195 if not parent_doc.meta.get("__children_ids"): 

196 raise ValueError(f"Parent document with id {parent_id} does not have any children.") 

197 

198 return parent_doc 

199 

200 async def _try_merge_level(docs_to_merge: list[Document], docs_to_return: list[Document]) -> list[Document]: 

201 parent_doc_id_to_child_docs: dict[str, list[Document]] = defaultdict(list) # to group documents by parent 

202 

203 for doc in docs_to_merge: 

204 if doc.meta.get("__parent_id"): # only docs that have parents 

205 parent_doc_id_to_child_docs[doc.meta["__parent_id"]].append(doc) 

206 else: 

207 docs_to_return.append(doc) # keep docs that have no parents 

208 

209 # Process each parent group 

210 merged_docs = [] 

211 for parent_doc_id, child_docs in parent_doc_id_to_child_docs.items(): 

212 parent_doc = await _get_parent_doc(parent_doc_id) 

213 

214 # Calculate merge score 

215 score = len(child_docs) / len(parent_doc.meta["__children_ids"]) 

216 if score > self.threshold: 

217 merged_docs.append(parent_doc) # Merge into parent 

218 else: 

219 docs_to_return.extend(child_docs) # Keep children separate 

220 

221 # if no new merges were made, we're done 

222 if not merged_docs: 

223 return merged_docs + docs_to_return 

224 

225 # Recursively try to merge the next level 

226 return await _try_merge_level(merged_docs, docs_to_return) 

227 

228 return {"documents": await _try_merge_level(documents, [])} 

229 

230 def close(self) -> None: 

231 """ 

232 Release the synchronous resources of the underlying Document Store. 

233 """ 

234 if hasattr(self.document_store, "close"): 

235 self.document_store.close() 

236 

237 async def close_async(self) -> None: 

238 """ 

239 Release the asynchronous resources of the underlying Document Store. 

240 """ 

241 if hasattr(self.document_store, "close_async"): 

242 await self.document_store.close_async()