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 = int(parts[0]), int(parts[1]) 

65 if start > end: 

66 msg = "start must be less than or equal to end" 

67 raise ValueError(f"Invalid page range: '{parts[0]}-{parts[1]}' - {msg}") 

68 expanded_page_range.extend(range(start, end + 1)) 

69 

70 else: 

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

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

73 

74 if not expanded_page_range: 

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

76 

77 return expanded_page_range 

78 

79 

80@overload 

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

82@overload 

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

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

85 """ 

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

87 

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

89 """ 

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

91 

92 

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

94 """ 

95 Guess the MIME type of the provided file path. 

96 

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

98 

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

100 """ 

101 extension = path.suffix.lower() 

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

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

104 return CUSTOM_MIMETYPES.get(extension, mime_type) 

105 

106 

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

108 """ 

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

110 

111 Tries in the following order: 

112 

113 1. ~/.haystack/{out_dir} 

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

115 3. ./.haystack/{out_dir} 

116 

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

118 :returns: 

119 The path to the created directory. 

120 """ 

121 

122 candidates = [ 

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

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

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

126 ] 

127 

128 for candidate in candidates: 

129 try: 

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

131 return str(candidate) 

132 except Exception: 

133 continue 

134 

135 raise RuntimeError( 

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

137 ) 

138 

139 

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

141 """ 

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

143 

144 :param documents: List of documents to deduplicate. 

145 :returns: List of deduplicated documents. 

146 """ 

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

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

149 for doc in documents: 

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

151 best = highest_scoring_docs.get(doc.id) 

152 

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

154 highest_scoring_docs[doc.id] = doc 

155 

156 return list(highest_scoring_docs.values()) 

157 

158 

159def _reciprocal_rank_fusion( 

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

161) -> list["Document"]: 

162 """ 

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

164 

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

166 

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

168 paper used 1-based ranking). 

169 

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

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

172 :returns: 

173 Deduplicated list of documents with updated RRF scores. 

174 """ 

175 if not document_lists: 

176 return [] 

177 

178 k = 61 

179 scores_map: dict = defaultdict(int) 

180 documents_map: dict = {} 

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

182 

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

184 for rank, doc in enumerate(documents): 

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

186 documents_map[doc.id] = doc 

187 

188 for _id in scores_map: 

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

190 

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

192 

193 

194@overload 

195def _parse_dict_from_json( 

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

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

198@overload 

199def _parse_dict_from_json( 

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

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

202@overload 

203def _parse_dict_from_json( 

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

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

206def _parse_dict_from_json( 

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

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

209 """ 

210 Parses a JSON string containing a dictionary. 

211 

212 :param text: The string to parse. 

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

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

215 

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

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

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

219 and `raise_on_failure` is True. 

220 """ 

221 cleaned_text = text.strip() 

222 

223 try: 

224 parsed_json = json.loads(cleaned_text) 

225 except json.JSONDecodeError as e: 

226 if raise_on_failure: 

227 raise e 

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

229 return None 

230 

231 if not isinstance(parsed_json, dict): 

232 if raise_on_failure: 

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

234 logger.warning( 

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

236 type=type(parsed_json).__name__, 

237 ) 

238 return None 

239 

240 if not expected_keys: 

241 return parsed_json 

242 

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

244 if missing_keys: 

245 if raise_on_failure: 

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

247 logger.warning( 

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

249 missing_keys=missing_keys, 

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

251 ) 

252 return None 

253 

254 return parsed_json 

255 

256 

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

258 """ 

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

260 """ 

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