Coverage for haystack/components/retrievers/in_memory/bm25_retriever.py: 100%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from typing import Any 

6 

7from haystack import Document, component, default_from_dict, default_to_dict 

8from haystack.document_stores.in_memory import InMemoryDocumentStore 

9from haystack.document_stores.types import FilterPolicy, apply_filter_policy 

10 

11 

12@component 

13class InMemoryBM25Retriever: 

14 """ 

15 Retrieves documents that are most similar to the query using keyword-based algorithm. 

16 

17 Use this retriever with the InMemoryDocumentStore. 

18 

19 ### Usage example 

20 

21 ```python 

22 from haystack import Document 

23 from haystack.components.retrievers.in_memory import InMemoryBM25Retriever 

24 from haystack.document_stores.in_memory import InMemoryDocumentStore 

25 

26 docs = [ 

27 Document(content="Python is a popular programming language"), 

28 Document(content="python ist eine beliebte Programmiersprache"), 

29 ] 

30 

31 doc_store = InMemoryDocumentStore() 

32 doc_store.write_documents(docs) 

33 retriever = InMemoryBM25Retriever(doc_store) 

34 

35 result = retriever.run(query="Programmiersprache") 

36 

37 print(result["documents"]) 

38 ``` 

39 """ 

40 

41 def __init__( 

42 self, 

43 document_store: InMemoryDocumentStore, 

44 filters: dict[str, Any] | None = None, 

45 top_k: int = 10, 

46 scale_score: bool = False, 

47 filter_policy: FilterPolicy = FilterPolicy.REPLACE, 

48 ) -> None: 

49 """ 

50 Create the InMemoryBM25Retriever component. 

51 

52 :param document_store: 

53 An instance of InMemoryDocumentStore where the retriever should search for relevant documents. 

54 :param filters: 

55 A dictionary with filters to narrow down the retriever's search space in the document store. 

56 :param top_k: 

57 The maximum number of documents to retrieve. 

58 :param scale_score: 

59 When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant. 

60 When `False`, uses raw similarity scores. 

61 :param filter_policy: The filter policy to apply during retrieval. 

62 Filter policy determines how filters are applied when retrieving documents. You can choose: 

63 - `REPLACE` (default): Overrides the initialization filters with the filters specified at runtime. 

64 Use this policy to dynamically change filtering for specific queries. 

65 - `MERGE`: Combines runtime filters with initialization filters to narrow down the search. 

66 :raises TypeError: If the document_store is not an instance of InMemoryDocumentStore. 

67 :raises ValueError: 

68 If the specified `top_k` is not > 0. 

69 """ 

70 if not isinstance(document_store, InMemoryDocumentStore): 

71 raise TypeError("document_store must be an instance of InMemoryDocumentStore") 

72 

73 self.document_store = document_store 

74 

75 if top_k <= 0: 

76 raise ValueError(f"top_k must be greater than 0. Currently, the top_k is {top_k}") 

77 

78 self.filters = filters 

79 self.top_k = top_k 

80 self.scale_score = scale_score 

81 self.filter_policy = filter_policy 

82 

83 def _get_telemetry_data(self) -> dict[str, Any]: 

84 """ 

85 Data that is sent to Posthog for usage analytics. 

86 """ 

87 return {"document_store": type(self.document_store).__name__} 

88 

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

90 """ 

91 Serializes the component to a dictionary. 

92 

93 :returns: 

94 Dictionary with serialized data. 

95 """ 

96 return default_to_dict( 

97 self, 

98 document_store=self.document_store, 

99 filters=self.filters, 

100 top_k=self.top_k, 

101 scale_score=self.scale_score, 

102 filter_policy=self.filter_policy.value, 

103 ) 

104 

105 @classmethod 

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

107 """ 

108 Deserializes the component from a dictionary. 

109 

110 :param data: 

111 The dictionary to deserialize from. 

112 :returns: 

113 The deserialized component. 

114 """ 

115 init_params = data.get("init_parameters", {}) 

116 if "filter_policy" in init_params: 

117 init_params["filter_policy"] = FilterPolicy.from_str(init_params["filter_policy"]) 

118 return default_from_dict(cls, data) 

119 

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

121 def run( 

122 self, 

123 query: str, 

124 filters: dict[str, Any] | None = None, 

125 top_k: int | None = None, 

126 scale_score: bool | None = None, 

127 ) -> dict[str, list[Document]]: 

128 """ 

129 Run the InMemoryBM25Retriever on the given input data. 

130 

131 :param query: 

132 The query string for the Retriever. 

133 :param filters: 

134 A dictionary with filters to narrow down the search space when retrieving documents. 

135 :param top_k: 

136 The maximum number of documents to return. 

137 :param scale_score: 

138 When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant. 

139 When `False`, uses raw similarity scores. 

140 :returns: 

141 The retrieved documents. 

142 

143 :raises ValueError: 

144 If the specified DocumentStore is not found or is not a InMemoryDocumentStore instance. 

145 """ 

146 filters = apply_filter_policy(self.filter_policy, self.filters, filters) 

147 if top_k is None: 

148 top_k = self.top_k 

149 if scale_score is None: 

150 scale_score = self.scale_score 

151 

152 docs = self.document_store.bm25_retrieval(query=query, filters=filters, top_k=top_k, scale_score=scale_score) 

153 return {"documents": docs} 

154 

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

156 async def run_async( 

157 self, 

158 query: str, 

159 filters: dict[str, Any] | None = None, 

160 top_k: int | None = None, 

161 scale_score: bool | None = None, 

162 ) -> dict[str, list[Document]]: 

163 """ 

164 Run the InMemoryBM25Retriever on the given input data. 

165 

166 :param query: 

167 The query string for the Retriever. 

168 :param filters: 

169 A dictionary with filters to narrow down the search space when retrieving documents. 

170 :param top_k: 

171 The maximum number of documents to return. 

172 :param scale_score: 

173 When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant. 

174 When `False`, uses raw similarity scores. 

175 :returns: 

176 The retrieved documents. 

177 

178 :raises ValueError: 

179 If the specified DocumentStore is not found or is not a InMemoryDocumentStore instance. 

180 """ 

181 filters = apply_filter_policy(self.filter_policy, self.filters, filters) 

182 if top_k is None: 

183 top_k = self.top_k 

184 if scale_score is None: 

185 scale_score = self.scale_score 

186 

187 docs = await self.document_store.bm25_retrieval_async( 

188 query=query, filters=filters, top_k=top_k, scale_score=scale_score 

189 ) 

190 return {"documents": docs}