Coverage for haystack/utils/misc.py: 95%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import json 

6import mimetypes 

7import tempfile 

8from collections import defaultdict 

9from dataclasses import replace 

10from math import inf 

11from pathlib import Path 

12from typing import TYPE_CHECKING, Any, Literal, overload 

13 

14from numpy import exp, ndarray 

15 

16from haystack import logging 

17 

18if TYPE_CHECKING: 

19 from haystack.dataclasses import Document 

20 

21CUSTOM_MIMETYPES = { 

22 # we add markdown because it is not added by the mimetypes module 

23 # see https://github.com/python/cpython/pull/17995 

24 ".md": "text/markdown", 

25 ".markdown": "text/markdown", 

26 # we add msg because it is not added by the mimetypes module 

27 ".msg": "application/vnd.ms-outlook", 

28} 

29 

30logger = logging.getLogger(__name__) 

31 

32 

33def expand_page_range(page_range: list[str | int]) -> list[int]: 

34 """ 

35 Takes a list of page numbers and ranges and expands them into a list of page numbers. 

36 

37 For example, given a page_range=['1-3', '5', '8', '10-12'] the function will return [1, 2, 3, 5, 8, 10, 11, 12] 

38 

39 :param page_range: List of page numbers and ranges 

40 :returns: 

41 An expanded list of page integers 

42 :raises ValueError: 

43 If any element is not a valid integer or a range string in the format `'start-end'`. 

44 

45 """ 

46 expanded_page_range = [] 

47 

48 for page in page_range: 

49 if isinstance(page, int): 

50 # check if it's a range wrongly passed as an integer expression 

51 if "-" in str(page): 

52 msg = "range must be a string in the format 'start-end'" 

53 raise ValueError(f"Invalid page range: {page} - {msg}") 

54 expanded_page_range.append(page) 

55 

56 elif isinstance(page, str) and page.isdigit(): 

57 expanded_page_range.append(int(page)) 

58 

59 elif isinstance(page, str) and "-" in page: 

60 parts = page.split("-", maxsplit=1) 

61 if not parts[0].isdigit() or not parts[1].isdigit(): 

62 msg = "range must be a string in the format 'start-end'" 

63 raise ValueError(f"Invalid page range: {page} - {msg}") 

64 start, end = parts 

65 expanded_page_range.extend(range(int(start), int(end) + 1)) 

66 

67 else: 

68 msg = "range must be a string in the format 'start-end' or an integer" 

69 raise ValueError(f"Invalid page range: {page} - {msg}") 

70 

71 if not expanded_page_range: 

72 raise ValueError("No valid page numbers or ranges found in the input list") 

73 

74 return expanded_page_range 

75 

76 

77@overload 

78def expit(x: float) -> float: ... 

79@overload 

80def expit(x: ndarray[Any, Any]) -> ndarray[Any, Any]: ... 

81def expit(x: float | ndarray[Any, Any]) -> float | ndarray[Any, Any]: 

82 """ 

83 Compute logistic sigmoid function. Maps input values to a range between 0 and 1 

84 

85 :param x: input value. Can be a scalar or a numpy array. 

86 """ 

87 return 1 / (1 + exp(-x)) 

88 

89 

90def _guess_mime_type(path: Path) -> str | None: 

91 """ 

92 Guess the MIME type of the provided file path. 

93 

94 :param path: The file path to get the MIME type for. 

95 

96 :returns: The MIME type of the provided file path, or `None` if the MIME type cannot be determined. 

97 """ 

98 extension = path.suffix.lower() 

99 mime_type = mimetypes.guess_type(path.as_posix())[0] 

100 # lookup custom mappings if the mime type is not found 

101 return CUSTOM_MIMETYPES.get(extension, mime_type) 

102 

103 

104def _get_output_dir(out_dir: str) -> str: 

105 """ 

106 Find or create a writable directory for saving status files. 

107 

108 Tries in the following order: 

109 

110 1. ~/.haystack/{out_dir} 

111 2. {tempdir}/haystack/{out_dir} 

112 3. ./.haystack/{out_dir} 

113 

114 :raises RuntimeError: If no directory could be created. 

115 :returns: 

116 The path to the created directory. 

117 """ 

118 

119 candidates = [ 

120 Path.home() / ".haystack" / out_dir, 

121 Path(tempfile.gettempdir()) / "haystack" / out_dir, 

122 Path.cwd() / ".haystack" / out_dir, 

123 ] 

124 

125 for candidate in candidates: 

126 try: 

127 candidate.mkdir(parents=True, exist_ok=True) 

128 return str(candidate) 

129 except Exception: 

130 continue 

131 

132 raise RuntimeError( 

133 f"Could not create a writable directory for output files in any of the following locations: {candidates}" 

134 ) 

135 

136 

137def _deduplicate_documents(documents: list["Document"]) -> list["Document"]: 

138 """ 

139 Deduplicate a list of documents by their id keeping the duplicate with the highest score if a score is present. 

140 

141 :param documents: List of documents to deduplicate. 

142 :returns: List of deduplicated documents. 

143 """ 

144 # Keep for each Document id the one with the highest score 

145 highest_scoring_docs: dict[str, "Document"] = {} 

146 for doc in documents: 

147 score = doc.score if doc.score is not None else -inf 

148 best = highest_scoring_docs.get(doc.id) 

149 

150 if best is None or score > (best.score if best.score is not None else -inf): 

151 highest_scoring_docs[doc.id] = doc 

152 

153 return list(highest_scoring_docs.values()) 

154 

155 

156def _reciprocal_rank_fusion( 

157 document_lists: list[list["Document"]], weights: list[float] | None = None 

158) -> list["Document"]: 

159 """ 

160 Merge multiple ranked lists of Documents using Reciprocal Rank Fusion, deduplicating across lists. 

161 

162 See the original paper: https://plg.uwaterloo.ca/~gvcormac/cormacksigir09-rrf.pdf 

163 

164 The constant k is set to 61 (60 was suggested by the original paper, plus 1 as python lists are 0-based and the 

165 paper used 1-based ranking). 

166 

167 :param document_lists: A list of ranked document lists to fuse. 

168 :param weights: Optional per-list weights. Defaults to equal weights. 

169 :returns: 

170 Deduplicated list of documents with updated RRF scores. 

171 """ 

172 if not document_lists: 

173 return [] 

174 

175 k = 61 

176 scores_map: dict = defaultdict(int) 

177 documents_map: dict = {} 

178 resolved_weights = weights if weights else [1 / len(document_lists)] * len(document_lists) 

179 

180 for documents, weight in zip(document_lists, resolved_weights, strict=True): 

181 for rank, doc in enumerate(documents): 

182 scores_map[doc.id] += (weight * len(document_lists)) / (k + rank) 

183 documents_map[doc.id] = doc 

184 

185 for _id in scores_map: 

186 scores_map[_id] /= len(document_lists) / k 

187 

188 return [replace(doc, score=scores_map[doc.id]) for doc in documents_map.values()] 

189 

190 

191@overload 

192def _parse_dict_from_json( 

193 text: str, expected_keys: list[str] | None = ..., raise_on_failure: Literal[True] = ... 

194) -> dict[str, Any]: ... 

195@overload 

196def _parse_dict_from_json( 

197 text: str, expected_keys: list[str] | None = ..., raise_on_failure: Literal[False] = ... 

198) -> dict[str, Any] | None: ... 

199@overload 

200def _parse_dict_from_json( 

201 text: str, expected_keys: list[str] | None = ..., raise_on_failure: bool = ... 

202) -> dict[str, Any] | None: ... 

203def _parse_dict_from_json( 

204 text: str, expected_keys: list[str] | None = None, raise_on_failure: bool = True 

205) -> dict[str, Any] | None: 

206 """ 

207 Parses a JSON string containing a dictionary. 

208 

209 :param text: The string to parse. 

210 :param expected_keys: A list of keys that must be present in the parsed dictionary. 

211 :param raise_on_failure: If True, raises an exception on failure. If False, logs a warning and returns None. 

212 

213 :return: The parsed dictionary, or None if parsing fails and raise_on_failure is False. 

214 :raises json.JSONDecodeError: If the text is not valid JSON and raise_on_failure is True. 

215 :raises ValueError: If the parsed object is not a dictionary or has missing expected keys, 

216 and `raise_on_failure` is True. 

217 """ 

218 cleaned_text = text.strip() 

219 

220 try: 

221 parsed_json = json.loads(cleaned_text) 

222 except json.JSONDecodeError as e: 

223 if raise_on_failure: 

224 raise e 

225 logger.warning("Failed to parse JSON from text: {text}. Error: {error}", text=text, error=e) 

226 return None 

227 

228 if not isinstance(parsed_json, dict): 

229 if raise_on_failure: 

230 raise ValueError(f"Expected a JSON object containing a dictionary but got {type(parsed_json).__name__}") 

231 logger.warning( 

232 "Expected a JSON object containing a dictionary but got {type}. Returning None", 

233 type=type(parsed_json).__name__, 

234 ) 

235 return None 

236 

237 if not expected_keys: 

238 return parsed_json 

239 

240 missing_keys = [key for key in expected_keys if key not in parsed_json] 

241 if missing_keys: 

242 if raise_on_failure: 

243 raise ValueError(f"Missing expected keys in JSON: {missing_keys}. Got keys: {list(parsed_json.keys())}") 

244 logger.warning( 

245 "Missing expected keys in JSON: {missing_keys}. Got keys: {keys}", 

246 missing_keys=missing_keys, 

247 keys=list(parsed_json.keys()), 

248 ) 

249 return None 

250 

251 return parsed_json 

252 

253 

254def _normalize_metadata_field_name(metadata_field: str) -> str: 

255 """ 

256 Normalizes a metadata field name by removing the "meta." prefix if present. 

257 """ 

258 return metadata_field[5:] if metadata_field.startswith("meta.") else metadata_field