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 if not all(doc.meta.get("__level") for doc in matched_leaf_documents): 

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

109 

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

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

112 

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

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

115 """ 

116 Run the AutoMergingRetriever. 

117 

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

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

120 

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

122 :returns: 

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

124 """ 

125 

126 AutoMergingRetriever._check_valid_documents(documents) 

127 

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

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

130 if len(parent_docs) != 1: 

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

132 

133 parent_doc = parent_docs[0] 

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

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

136 

137 return parent_doc 

138 

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

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

141 

142 for doc in docs_to_merge: 

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

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

145 else: 

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

147 

148 # Process each parent group 

149 merged_docs = [] 

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

151 parent_doc = _get_parent_doc(parent_doc_id) 

152 

153 # Calculate merge score 

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

155 if score > self.threshold: 

156 merged_docs.append(parent_doc) # Merge into parent 

157 else: 

158 docs_to_return.extend(child_docs) # Keep children separate 

159 

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

161 if not merged_docs: 

162 return merged_docs + docs_to_return 

163 

164 # Recursively try to merge the next level 

165 return _try_merge_level(merged_docs, docs_to_return) 

166 

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

168 

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

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

171 """ 

172 Asynchronously run the AutoMergingRetriever. 

173 

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

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

176 

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

178 :returns: 

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

180 """ 

181 

182 AutoMergingRetriever._check_valid_documents(documents) 

183 

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

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

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

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

188 ) 

189 if len(parent_docs) != 1: 

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

191 

192 parent_doc = parent_docs[0] 

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

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

195 

196 return parent_doc 

197 

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

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

200 

201 for doc in docs_to_merge: 

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

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

204 else: 

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

206 

207 # Process each parent group 

208 merged_docs = [] 

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

210 parent_doc = await _get_parent_doc(parent_doc_id) 

211 

212 # Calculate merge score 

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

214 if score > self.threshold: 

215 merged_docs.append(parent_doc) # Merge into parent 

216 else: 

217 docs_to_return.extend(child_docs) # Keep children separate 

218 

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

220 if not merged_docs: 

221 return merged_docs + docs_to_return 

222 

223 # Recursively try to merge the next level 

224 return await _try_merge_level(merged_docs, docs_to_return) 

225 

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

227 

228 def close(self) -> None: 

229 """ 

230 Release the synchronous resources of the underlying Document Store. 

231 """ 

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

233 self.document_store.close() 

234 

235 async def close_async(self) -> None: 

236 """ 

237 Release the asynchronous resources of the underlying Document Store. 

238 """ 

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

240 await self.document_store.close_async()