Coverage for haystack/document_stores/in_memory/document_store.py: 97%

431 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-21 13:53 +0000

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import asyncio 

6import json 

7import math 

8import re 

9import uuid 

10from collections import Counter 

11from collections.abc import Callable, Iterable 

12from concurrent.futures import ThreadPoolExecutor 

13from dataclasses import dataclass, replace 

14from pathlib import Path 

15from typing import Any, Literal 

16 

17import numpy as np 

18 

19from haystack import default_from_dict, default_to_dict, logging 

20from haystack.dataclasses import Document 

21from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError 

22from haystack.document_stores.types import DuplicatePolicy 

23from haystack.utils import expit 

24from haystack.utils.filters import document_matches_filter 

25 

26logger = logging.getLogger(__name__) 

27 

28# document scores are essentially unbounded and will be scaled to values between 0 and 1 if scale_score is set to 

29# True (default). Scaling uses the expit function (inverse of the logit function) after applying a scaling factor 

30# (e.g., BM25_SCALING_FACTOR for the bm25_retrieval method). 

31# Larger scaling factor decreases scaled scores. For example, an input of 10 is scaled to 0.99 with 

32# BM25_SCALING_FACTOR=2 but to 0.78 with BM25_SCALING_FACTOR=8 (default). The defaults were chosen empirically. 

33# Increase the default if most unscaled scores are larger than expected (>30) and otherwise would incorrectly all be 

34# mapped to scores ~1. 

35BM25_SCALING_FACTOR = 8 

36DOT_PRODUCT_SCALING_FACTOR = 100 

37 

38 

39@dataclass 

40class BM25DocumentStats: 

41 """ 

42 A dataclass for managing document statistics for BM25 retrieval. 

43 

44 :param freq_token: A Counter of token frequencies in the document. 

45 :param doc_len: Number of tokens in the document. 

46 """ 

47 

48 freq_token: dict[str, int] 

49 doc_len: int 

50 

51 

52# Process-global storage, keyed by index, for stores created with shared=True. Lets instances sharing an index 

53# operate on the same data. Non-shared stores keep their data instance-local instead (see `__init__`). 

54_STORAGES: dict[str, dict[str, Document]] = {} 

55_BM25_STATS_STORAGES: dict[str, dict[str, BM25DocumentStats]] = {} 

56_AVERAGE_DOC_LEN_STORAGES: dict[str, float] = {} 

57_FREQ_VOCAB_FOR_IDF_STORAGES: dict[str, Counter] = {} 

58 

59 

60class InMemoryDocumentStore: 

61 """ 

62 Stores data in-memory. It's ephemeral and cannot be saved to disk. 

63 """ 

64 

65 def __init__( 

66 self, 

67 bm25_tokenization_regex: str = r"(?u)\b\w+\b", 

68 bm25_algorithm: Literal["BM25Okapi", "BM25L", "BM25Plus"] = "BM25L", 

69 bm25_parameters: dict | None = None, 

70 embedding_similarity_function: Literal["dot_product", "cosine"] = "dot_product", 

71 index: str | None = None, 

72 shared: bool = True, 

73 async_executor: ThreadPoolExecutor | None = None, 

74 return_embedding: bool = True, 

75 *, 

76 strict_datetime_comparison: bool = False, 

77 ) -> None: 

78 """ 

79 Initializes the DocumentStore. 

80 

81 :param bm25_tokenization_regex: The regular expression used to tokenize the text for BM25 retrieval. 

82 :param bm25_algorithm: The BM25 algorithm to use. One of "BM25Okapi", "BM25L", or "BM25Plus". 

83 :param bm25_parameters: Parameters for BM25 implementation in a dictionary format. 

84 For example: `{'k1':1.5, 'b':0.75, 'epsilon':0.25}` 

85 You can learn more about these parameters by visiting https://github.com/dorianbrown/rank_bm25. 

86 :param embedding_similarity_function: The similarity function used to compare Documents embeddings. 

87 One of "dot_product" (default) or "cosine". To choose the most appropriate function, look for information 

88 about your embedding model. 

89 :param index: A specific index to store the documents. If not specified, a random UUID is used. 

90 When `shared` is True, instances using the same index share the same documents. 

91 :param shared: Whether the documents live in process-global storage shared across instances using the same 

92 index (True, the default), or are kept instance-local and freed when this instance is garbage collected 

93 (False). Shared storage persists for the lifetime of the process, so prefer `shared=False` for stores 

94 that are created frequently (for example per request) to avoid unbounded memory growth. 

95 :param async_executor: 

96 Optional ThreadPoolExecutor to use for async calls. If not provided, a single-threaded 

97 executor will be initialized and used. 

98 :param return_embedding: Whether to return the embedding of the retrieved Documents. Default is True. 

99 :param strict_datetime_comparison: 

100 If `True`, timezone-naive and timezone-aware datetimes never match each other in filters. 

101 If `False` (the default), the timezone from the aware datetime is copied to the naive one before 

102 comparing. 

103 """ 

104 self.bm25_tokenization_regex = bm25_tokenization_regex 

105 self.tokenizer = re.compile(bm25_tokenization_regex).findall 

106 

107 # Shared stores keep their data in the process-global dicts keyed by index, so instances sharing an index 

108 # operate on the same documents. Non-shared stores keep their data instance-local so it is freed with the 

109 # instance instead of accumulating in the globals. 

110 self._shared = shared 

111 self.index = index if index is not None else str(uuid.uuid4()) 

112 

113 if self._shared: 

114 if self.index not in _STORAGES: 

115 _STORAGES[self.index] = {} 

116 _BM25_STATS_STORAGES[self.index] = {} 

117 _AVERAGE_DOC_LEN_STORAGES[self.index] = 0.0 

118 _FREQ_VOCAB_FOR_IDF_STORAGES[self.index] = Counter() 

119 else: 

120 self._local_storage: dict[str, Document] = {} 

121 self._local_bm25_attr: dict[str, BM25DocumentStats] = {} 

122 self._local_avg_doc_len: float = 0.0 

123 self._local_freq_vocab_for_idf: Counter = Counter() 

124 

125 self.bm25_algorithm = bm25_algorithm 

126 self.bm25_algorithm_inst = self._dispatch_bm25() 

127 self.bm25_parameters = bm25_parameters or {} 

128 self.embedding_similarity_function = embedding_similarity_function 

129 

130 # keep track of whether we own the executor if we created it we must also clean it up 

131 self._owns_executor = async_executor is None 

132 self.executor = ( 

133 ThreadPoolExecutor(thread_name_prefix=f"async-inmemory-docstore-executor-{id(self)}", max_workers=1) 

134 if async_executor is None 

135 else async_executor 

136 ) 

137 self.return_embedding = return_embedding 

138 self.strict_datetime_comparison = strict_datetime_comparison 

139 

140 def __del__(self) -> None: 

141 """ 

142 Cleanup when the instance is being destroyed. 

143 """ 

144 if hasattr(self, "_owns_executor") and self._owns_executor and hasattr(self, "executor"): 

145 self.executor.shutdown(wait=True) 

146 

147 def shutdown(self) -> None: 

148 """ 

149 Explicitly shutdown the executor if we own it. 

150 """ 

151 if self._owns_executor: 

152 self.executor.shutdown(wait=True) 

153 

154 @property 

155 def storage(self) -> dict[str, Document]: 

156 """ 

157 Utility property that returns the storage used by this instance of InMemoryDocumentStore. 

158 """ 

159 return _STORAGES[self.index] if self._shared else self._local_storage 

160 

161 @property 

162 def _bm25_attr(self) -> dict[str, BM25DocumentStats]: 

163 return _BM25_STATS_STORAGES[self.index] if self._shared else self._local_bm25_attr 

164 

165 @property 

166 def _avg_doc_len(self) -> float: 

167 return _AVERAGE_DOC_LEN_STORAGES[self.index] if self._shared else self._local_avg_doc_len 

168 

169 @_avg_doc_len.setter 

170 def _avg_doc_len(self, value: float) -> None: 

171 if self._shared: 

172 _AVERAGE_DOC_LEN_STORAGES[self.index] = value 

173 else: 

174 self._local_avg_doc_len = value 

175 

176 @property 

177 def _freq_vocab_for_idf(self) -> Counter: 

178 return _FREQ_VOCAB_FOR_IDF_STORAGES[self.index] if self._shared else self._local_freq_vocab_for_idf 

179 

180 def _dispatch_bm25(self) -> "Callable[[str, list[Document]], list[tuple[Document, float]]]": 

181 """ 

182 Select the correct BM25 algorithm based on user specification. 

183 

184 :returns: 

185 The BM25 algorithm method. 

186 """ 

187 table = {"BM25Okapi": self._score_bm25okapi, "BM25L": self._score_bm25l, "BM25Plus": self._score_bm25plus} 

188 

189 if self.bm25_algorithm not in table: 

190 raise ValueError(f"BM25 algorithm '{self.bm25_algorithm}' is not supported.") 

191 return table[self.bm25_algorithm] 

192 

193 def _tokenize_bm25(self, text: str) -> list[str]: 

194 """ 

195 Tokenize text using the BM25 tokenization regex. 

196 

197 Here we explicitly create a tokenization method to encapsulate 

198 all pre-processing logic used to create BM25 tokens, such as 

199 lowercasing. This helps track the exact tokenization process 

200 used for BM25 scoring at any given time. 

201 

202 :param text: 

203 The text to tokenize. 

204 :returns: 

205 A list of tokens. 

206 """ 

207 text = text.lower() 

208 return self.tokenizer(text) 

209 

210 def _score_bm25l(self, query: str, documents: list[Document]) -> list[tuple[Document, float]]: 

211 """ 

212 Calculate BM25L scores for the given query and filtered documents. 

213 

214 :param query: 

215 The query string. 

216 :param documents: 

217 The list of documents to score, should be produced by 

218 the filter_documents method; may be an empty list. 

219 :returns: 

220 A list of tuples, each containing a Document and its BM25L score. 

221 """ 

222 k = self.bm25_parameters.get("k1", 1.5) 

223 b = self.bm25_parameters.get("b", 0.75) 

224 delta = self.bm25_parameters.get("delta", 0.5) 

225 

226 def _compute_idf(tokens: list[str]) -> dict[str, float]: 

227 """Per-token IDF computation for all tokens.""" 

228 idf = {} 

229 n_corpus = len(self._bm25_attr) 

230 for tok in tokens: 

231 n = self._freq_vocab_for_idf.get(tok, 0) 

232 idf[tok] = math.log((n_corpus + 1.0) / (n + 0.5)) * int(n != 0) 

233 return idf 

234 

235 def _compute_tf(token: str, freq: dict[str, int], doc_len: int) -> float: 

236 """Per-token BM25L computation.""" 

237 freq_term = freq.get(token, 0.0) 

238 ctd = freq_term / (1 - b + b * doc_len / self._avg_doc_len) 

239 return (1.0 + k) * (ctd + delta) / (k + ctd + delta) 

240 

241 idf = _compute_idf(self._tokenize_bm25(query)) 

242 bm25_attr = {doc.id: self._bm25_attr[doc.id] for doc in documents} 

243 

244 ret = [] 

245 for doc in documents: 

246 doc_stats = bm25_attr[doc.id] 

247 freq = doc_stats.freq_token 

248 doc_len = doc_stats.doc_len 

249 

250 score = 0.0 

251 for tok in idf.keys(): 

252 score += idf[tok] * _compute_tf(tok, freq, doc_len) 

253 ret.append((doc, score)) 

254 

255 return ret 

256 

257 def _score_bm25okapi(self, query: str, documents: list[Document]) -> list[tuple[Document, float]]: 

258 """ 

259 Calculate BM25Okapi scores for the given query and filtered documents. 

260 

261 :param query: 

262 The query string. 

263 :param documents: 

264 The list of documents to score, should be produced by 

265 the filter_documents method; may be an empty list. 

266 :returns: 

267 A list of tuples, each containing a Document and its BM25Okapi score. 

268 """ 

269 k = self.bm25_parameters.get("k1", 1.5) 

270 b = self.bm25_parameters.get("b", 0.75) 

271 epsilon = self.bm25_parameters.get("epsilon", 0.25) 

272 

273 def _compute_idf(tokens: list[str]) -> dict[str, float]: 

274 """Per-token IDF computation for all tokens.""" 

275 sum_idf = 0.0 

276 neg_idf_tokens = [] 

277 

278 # Although this is a global statistic, we compute it here 

279 # to make the computation more self-contained. And the 

280 # complexity is O(vocab_size), which is acceptable. 

281 idf = {} 

282 for tok, n in self._freq_vocab_for_idf.items(): 

283 idf[tok] = math.log((len(self._bm25_attr) - n + 0.5) / (n + 0.5)) 

284 sum_idf += idf[tok] 

285 if idf[tok] < 0: 

286 neg_idf_tokens.append(tok) 

287 

288 eps = epsilon * sum_idf / len(self._freq_vocab_for_idf) 

289 for tok in neg_idf_tokens: 

290 idf[tok] = eps 

291 return {tok: idf.get(tok, 0.0) for tok in tokens} 

292 

293 def _compute_tf(token: str, freq: dict[str, int], doc_len: int) -> float: 

294 """Per-token BM25Okapi computation.""" 

295 freq_term = freq.get(token, 0.0) 

296 freq_norm = freq_term + k * (1 - b + b * doc_len / self._avg_doc_len) 

297 return freq_term * (1.0 + k) / freq_norm 

298 

299 idf = _compute_idf(self._tokenize_bm25(query)) 

300 bm25_attr = {doc.id: self._bm25_attr[doc.id] for doc in documents} 

301 

302 ret = [] 

303 for doc in documents: 

304 doc_stats = bm25_attr[doc.id] 

305 freq = doc_stats.freq_token 

306 doc_len = doc_stats.doc_len 

307 

308 score = 0.0 

309 for tok in idf.keys(): 

310 score += idf[tok] * _compute_tf(tok, freq, doc_len) 

311 ret.append((doc, score)) 

312 

313 return ret 

314 

315 def _score_bm25plus(self, query: str, documents: list[Document]) -> list[tuple[Document, float]]: 

316 """ 

317 Calculate BM25+ scores for the given query and filtered documents. 

318 

319 This implementation follows the document on BM25 Wikipedia page, 

320 which add 1 (smoothing factor) to document frequency when computing IDF. 

321 

322 :param query: 

323 The query string. 

324 :param documents: 

325 The list of documents to score, should be produced by 

326 the filter_documents method; may be an empty list. 

327 :returns: 

328 A list of tuples, each containing a Document and its BM25+ score. 

329 """ 

330 k = self.bm25_parameters.get("k1", 1.5) 

331 b = self.bm25_parameters.get("b", 0.75) 

332 delta = self.bm25_parameters.get("delta", 1.0) 

333 

334 def _compute_idf(tokens: list[str]) -> dict[str, float]: 

335 """Per-token IDF computation.""" 

336 idf = {} 

337 n_corpus = len(self._bm25_attr) 

338 for tok in tokens: 

339 n = self._freq_vocab_for_idf.get(tok, 0) 

340 idf[tok] = math.log(1 + (n_corpus - n + 0.5) / (n + 0.5)) * int(n != 0) 

341 return idf 

342 

343 def _compute_tf(token: str, freq: dict[str, int], doc_len: float) -> float: 

344 """Per-token normalized term frequency.""" 

345 freq_term = freq.get(token, 0.0) 

346 freq_damp = k * (1 - b + b * doc_len / self._avg_doc_len) 

347 return freq_term * (1.0 + k) / (freq_term + freq_damp) + delta 

348 

349 idf = _compute_idf(self._tokenize_bm25(query)) 

350 bm25_attr = {doc.id: self._bm25_attr[doc.id] for doc in documents} 

351 

352 ret = [] 

353 for doc in documents: 

354 doc_stats = bm25_attr[doc.id] 

355 freq = doc_stats.freq_token 

356 doc_len = doc_stats.doc_len 

357 

358 score = 0.0 

359 for tok in idf.keys(): 

360 score += idf[tok] * _compute_tf(tok, freq, doc_len) 

361 ret.append((doc, score)) 

362 

363 return ret 

364 

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

366 """ 

367 Serializes the component to a dictionary. 

368 

369 :returns: 

370 Dictionary with serialized data. 

371 """ 

372 return default_to_dict( 

373 self, 

374 bm25_tokenization_regex=self.bm25_tokenization_regex, 

375 bm25_algorithm=self.bm25_algorithm, 

376 bm25_parameters=self.bm25_parameters, 

377 embedding_similarity_function=self.embedding_similarity_function, 

378 index=self.index, 

379 shared=self._shared, 

380 return_embedding=self.return_embedding, 

381 strict_datetime_comparison=self.strict_datetime_comparison, 

382 ) 

383 

384 @classmethod 

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

386 """ 

387 Deserializes the component from a dictionary. 

388 

389 :param data: 

390 The dictionary to deserialize from. 

391 :returns: 

392 The deserialized component. 

393 """ 

394 return default_from_dict(cls, data) 

395 

396 def save_to_disk(self, path: str) -> None: 

397 """ 

398 Write the database and its data to disk as a JSON file. 

399 

400 :param path: The path to the JSON file. 

401 """ 

402 data: dict[str, Any] = self.to_dict() 

403 data["documents"] = [doc.to_dict(flatten=False) for doc in self.storage.values()] 

404 with open(path, "w") as f: 

405 json.dump(data, f) 

406 

407 @classmethod 

408 def load_from_disk(cls, path: str) -> "InMemoryDocumentStore": 

409 """ 

410 Load the database and its data from disk as a JSON file. 

411 

412 :param path: The path to the JSON file. 

413 :returns: The loaded InMemoryDocumentStore. 

414 """ 

415 if Path(path).exists(): 

416 try: 

417 with open(path) as f: 

418 data = json.load(f) 

419 except Exception as e: 

420 raise DocumentStoreError(f"Error loading InMemoryDocumentStore from disk. error: {e}") from e 

421 

422 documents = data.pop("documents") 

423 cls_object = default_from_dict(cls, data) 

424 cls_object.write_documents( 

425 documents=[Document.from_dict(doc) for doc in documents], policy=DuplicatePolicy.OVERWRITE 

426 ) 

427 return cls_object 

428 

429 raise FileNotFoundError(f"File {path} not found.") 

430 

431 def count_documents(self) -> int: 

432 """ 

433 Returns the number of documents present in the DocumentStore. 

434 """ 

435 return len(self.storage.keys()) 

436 

437 def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Document]: 

438 """ 

439 Returns the documents that match the filters provided. 

440 

441 :param filters: The filters to apply. For a detailed specification of the filters, refer to the 

442 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering). 

443 :returns: A list of Documents that match the given filters. 

444 """ 

445 if filters: 

446 InMemoryDocumentStore._validate_filters(filters) 

447 docs = [ 

448 doc 

449 for doc in self.storage.values() 

450 if document_matches_filter( 

451 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison 

452 ) 

453 ] 

454 else: 

455 docs = list(self.storage.values()) 

456 

457 if not self.return_embedding: 

458 docs = [replace(doc, embedding=None) for doc in docs] 

459 

460 return docs 

461 

462 def write_documents(self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE) -> int: 

463 """ 

464 Refer to the DocumentStore.write_documents() protocol documentation. 

465 

466 If `policy` is set to `DuplicatePolicy.NONE` defaults to `DuplicatePolicy.FAIL`. 

467 """ 

468 if ( 

469 not isinstance(documents, Iterable) 

470 or isinstance(documents, str) 

471 or any(not isinstance(doc, Document) for doc in documents) 

472 ): 

473 raise ValueError("Please provide a list of Documents.") 

474 

475 if policy == DuplicatePolicy.NONE: 

476 policy = DuplicatePolicy.FAIL 

477 

478 written_documents = len(documents) 

479 for document in documents: 

480 if policy != DuplicatePolicy.OVERWRITE and document.id in self.storage.keys(): 

481 if policy == DuplicatePolicy.FAIL: 

482 raise DuplicateDocumentError(f"ID '{document.id}' already exists.") 

483 if policy == DuplicatePolicy.SKIP: 

484 logger.warning("ID '{document_id}' already exists", document_id=document.id) 

485 written_documents -= 1 

486 continue 

487 

488 # Since the statistics are updated in an incremental manner, 

489 # we need to explicitly remove the existing document to revert 

490 # the statistics before updating them with the new document. 

491 if document.id in self.storage.keys(): 

492 self.delete_documents([document.id]) 

493 

494 tokens = [] 

495 if document.content is not None: 

496 tokens = self._tokenize_bm25(document.content) 

497 

498 self.storage[document.id] = document 

499 

500 self._bm25_attr[document.id] = BM25DocumentStats(Counter(tokens), len(tokens)) 

501 self._freq_vocab_for_idf.update(set(tokens)) 

502 # Update avg doc len based on the new document and the previous average 

503 n_docs = len(self._bm25_attr) 

504 self._avg_doc_len = (len(tokens) + self._avg_doc_len * (n_docs - 1)) / n_docs 

505 return written_documents 

506 

507 def delete_documents(self, document_ids: list[str]) -> None: 

508 """ 

509 Deletes all documents with matching document_ids from the DocumentStore. 

510 

511 :param document_ids: The document_ids to delete. 

512 """ 

513 for doc_id in document_ids: 

514 if doc_id not in self.storage.keys(): 

515 continue 

516 del self.storage[doc_id] 

517 

518 # Update statistics accordingly 

519 doc_stats = self._bm25_attr.pop(doc_id) 

520 freq = doc_stats.freq_token 

521 doc_len = doc_stats.doc_len 

522 

523 self._freq_vocab_for_idf.subtract(Counter(freq.keys())) 

524 for token in freq: 

525 if self._freq_vocab_for_idf[token] <= 0: 

526 del self._freq_vocab_for_idf[token] 

527 try: 

528 self._avg_doc_len = (self._avg_doc_len * (len(self._bm25_attr) + 1) - doc_len) / len(self._bm25_attr) 

529 except ZeroDivisionError: 

530 self._avg_doc_len = 0 

531 

532 @staticmethod 

533 def _validate_filters(filters: dict[str, Any] | None) -> None: 

534 if filters and "operator" not in filters and "conditions" not in filters: 

535 raise ValueError( 

536 "Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details." 

537 ) 

538 

539 def delete_all_documents(self) -> None: 

540 """ 

541 Deletes all documents in the document store. 

542 """ 

543 if self._shared: 

544 _STORAGES[self.index] = {} 

545 _BM25_STATS_STORAGES[self.index] = {} 

546 _AVERAGE_DOC_LEN_STORAGES[self.index] = 0.0 

547 _FREQ_VOCAB_FOR_IDF_STORAGES[self.index] = Counter() 

548 else: 

549 self._local_storage = {} 

550 self._local_bm25_attr = {} 

551 self._local_avg_doc_len = 0.0 

552 self._local_freq_vocab_for_idf = Counter() 

553 

554 def update_by_filter(self, filters: dict[str, Any], meta: dict[str, Any]) -> int: 

555 """ 

556 Updates the metadata of all documents that match the provided filters. 

557 

558 :param filters: The filters to apply to select documents for updating. 

559 For filter syntax, see filter_documents. 

560 :param meta: The metadata fields to update. These will be merged with existing metadata. 

561 :returns: The number of documents updated. 

562 :raises ValueError: if filters have invalid syntax. 

563 """ 

564 InMemoryDocumentStore._validate_filters(filters) 

565 matching = [ 

566 doc 

567 for doc in self.storage.values() 

568 if document_matches_filter( 

569 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison 

570 ) 

571 ] 

572 for doc in matching: 

573 doc.meta.update(meta) 

574 self.storage[doc.id] = doc 

575 return len(matching) 

576 

577 def delete_by_filter(self, filters: dict[str, Any]) -> int: 

578 """ 

579 Deletes all documents that match the provided filters. 

580 

581 :param filters: The filters to apply to select documents for deletion. 

582 For filter syntax, see filter_documents. 

583 :returns: The number of documents deleted. 

584 :raises ValueError: if filters have invalid syntax. 

585 """ 

586 InMemoryDocumentStore._validate_filters(filters) 

587 matching = [ 

588 doc 

589 for doc in self.storage.values() 

590 if document_matches_filter( 

591 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison 

592 ) 

593 ] 

594 doc_ids = [doc.id for doc in matching] 

595 self.delete_documents(doc_ids) 

596 return len(doc_ids) 

597 

598 def count_documents_by_filter(self, filters: dict[str, Any]) -> int: 

599 """ 

600 Returns the number of documents that match the provided filters. 

601 

602 :param filters: The filters to apply. 

603 For a detailed specification of the filters, refer to the 

604 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering). 

605 :returns: The number of documents that match the filters. 

606 """ 

607 if filters: 

608 InMemoryDocumentStore._validate_filters(filters) 

609 return sum( 

610 1 

611 for doc in self.storage.values() 

612 if document_matches_filter( 

613 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison 

614 ) 

615 ) 

616 return len(self.storage) 

617 

618 def count_unique_metadata_by_filter(self, filters: dict[str, Any], metadata_fields: list[str]) -> dict[str, int]: 

619 """ 

620 Returns the number of unique values for each specified metadata field from documents matching the filters. 

621 

622 :param filters: The filters to apply. 

623 For a detailed specification of the filters, refer to the 

624 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering). 

625 :param metadata_fields: List of field names to count unique values for. 

626 Field names can include or omit the "meta." prefix. 

627 :returns: A dictionary mapping each metadata field name (without "meta." prefix) 

628 to the count of its unique values among the filtered documents. 

629 """ 

630 if filters: 

631 InMemoryDocumentStore._validate_filters(filters) 

632 docs = [ 

633 doc 

634 for doc in self.storage.values() 

635 if document_matches_filter( 

636 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison 

637 ) 

638 ] 

639 else: 

640 docs = list(self.storage.values()) 

641 

642 result: dict[str, int] = {} 

643 for field in metadata_fields: 

644 key = field.removeprefix("meta.") if field.startswith("meta.") else field 

645 values = {doc.meta.get(key) for doc in docs if key in doc.meta and doc.meta[key] is not None} 

646 result[key] = len(values) 

647 return result 

648 

649 def get_metadata_fields_info(self) -> dict[str, dict[str, str]]: 

650 """ 

651 Returns information about the metadata fields present in the stored documents. 

652 

653 Types are inferred from the stored values (keyword, int, float, boolean). 

654 

655 :returns: A dictionary mapping each metadata field name to a dict with a "type" key. 

656 """ 

657 type_map: dict[str, str] = {} 

658 for doc in self.storage.values(): 

659 for key, value in doc.meta.items(): 

660 if value is None: 

661 continue 

662 if isinstance(value, bool): 

663 type_map[key] = "boolean" 

664 elif isinstance(value, int): 

665 type_map[key] = "int" 

666 elif isinstance(value, float): 

667 type_map[key] = "float" 

668 else: 

669 type_map[key] = "keyword" 

670 return {k: {"type": v} for k, v in type_map.items()} 

671 

672 def get_metadata_field_min_max(self, metadata_field: str) -> dict[str, Any]: 

673 """ 

674 Returns the minimum and maximum values for the given metadata field across all documents. 

675 

676 :param metadata_field: The metadata field name. Can include or omit the "meta." prefix. 

677 :returns: A dictionary with "min" and "max" keys. Returns `{"min": None, "max": None}` 

678 if the field is missing or has no values. 

679 """ 

680 key = metadata_field.removeprefix("meta.") if metadata_field.startswith("meta.") else metadata_field 

681 values = [ 

682 doc.meta[key] 

683 for doc in self.storage.values() 

684 if key in doc.meta and doc.meta[key] is not None and isinstance(doc.meta[key], (int, float, str)) 

685 ] 

686 if not values: 

687 return {"min": None, "max": None} 

688 try: 

689 return {"min": min(values), "max": max(values)} 

690 except TypeError: 

691 return {"min": None, "max": None} 

692 

693 def get_metadata_field_unique_values( 

694 self, 

695 metadata_field: str, 

696 search_term: str | None = None, 

697 from_: int = 0, 

698 size: int = 10, 

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

700 ) -> tuple[list[Any], int]: 

701 """ 

702 Returns unique values for a metadata field, optionally filtered by a search term, with pagination. 

703 

704 :param metadata_field: The metadata field name. Can include or omit the "meta." prefix. 

705 :param search_term: Optional search term to filter values, matched as a case-insensitive substring 

706 against the metadata field's value. 

707 :param from_: The offset to start returning values from (for pagination). 

708 :param size: The maximum number of unique values to return. 

709 :param filters: Optional filters to restrict the documents considered. 

710 :returns: A tuple of (paginated list of unique values, total count of unique values). 

711 """ 

712 key = metadata_field.removeprefix("meta.") if metadata_field.startswith("meta.") else metadata_field 

713 unique_values: dict[tuple[str, str], Any] = {} 

714 for doc in self.filter_documents(filters=filters): 

715 value = doc.meta.get(key) 

716 if value is not None: 

717 unique_values.setdefault((type(value).__name__, str(value)), value) 

718 

719 if search_term: 

720 search_term_lower = search_term.lower() 

721 unique_values = {k: v for k, v in unique_values.items() if search_term_lower in k[1].lower()} 

722 

723 sorted_keys = sorted(unique_values, key=lambda k: (k[1], k[0])) 

724 paginated_keys = sorted_keys[from_ : from_ + size] 

725 return [unique_values[k] for k in paginated_keys], len(sorted_keys) 

726 

727 def bm25_retrieval( 

728 self, query: str, filters: dict[str, Any] | None = None, top_k: int = 10, scale_score: bool = False 

729 ) -> list[Document]: 

730 """ 

731 Retrieves documents that are most relevant to the query using BM25 algorithm. 

732 

733 :param query: The query string. 

734 :param filters: A dictionary with filters to narrow down the search space. 

735 :param top_k: The number of top documents to retrieve. Default is 10. 

736 :param scale_score: Whether to scale the scores of the retrieved documents. Default is False. 

737 :returns: A list of the top_k documents most relevant to the query. 

738 """ 

739 if not query: 

740 raise ValueError("Query should be a non-empty string") 

741 

742 content_type_filter = {"field": "content", "operator": "!=", "value": None} 

743 if filters: 

744 if "operator" not in filters: 

745 raise ValueError( 

746 "Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details." 

747 ) 

748 filters = {"operator": "AND", "conditions": [content_type_filter, filters]} 

749 else: 

750 filters = content_type_filter 

751 

752 all_documents = self.filter_documents(filters=filters) 

753 if len(all_documents) == 0: 

754 logger.info("No documents found for BM25 retrieval. Returning empty list.") 

755 return [] 

756 

757 # A tokenless corpus (every stored document has empty content) has no vocabulary and an 

758 # average document length of zero, which would make all three BM25 algorithms divide by 

759 # zero during scoring. Score every candidate as 0.0 instead; the non-positive-score 

760 # handling below then keeps them for BM25Okapi (unscaled) and drops them otherwise. 

761 if self._avg_doc_len == 0: 

762 scored_documents = [(doc, 0.0) for doc in all_documents] 

763 else: 

764 scored_documents = self.bm25_algorithm_inst(query, all_documents) 

765 

766 results = sorted(scored_documents, key=lambda x: x[1], reverse=True)[:top_k] 

767 

768 # BM25Okapi can return meaningful negative values, so they should not be filtered out when scale_score is False. 

769 # It's the only algorithm supported by rank_bm25 at the time of writing (2024) that can return negative scores. 

770 # see https://github.com/deepset-ai/haystack/pull/6889 for more context. 

771 negatives_are_valid = self.bm25_algorithm == "BM25Okapi" and not scale_score 

772 

773 # Create documents with the BM25 score to return them 

774 return_documents = [] 

775 for doc, score in results: 

776 if scale_score: 

777 score = expit(score / BM25_SCALING_FACTOR) 

778 

779 if not negatives_are_valid and score <= 0.0: 

780 continue 

781 

782 doc_fields = doc.to_dict() 

783 doc_fields["score"] = score 

784 

785 if not self.return_embedding and "embedding" in doc_fields: 

786 doc_fields.pop("embedding") 

787 

788 return_document = Document.from_dict(doc_fields) 

789 

790 return_documents.append(return_document) 

791 

792 return return_documents 

793 

794 def embedding_retrieval( 

795 self, 

796 query_embedding: list[float], 

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

798 top_k: int = 10, 

799 scale_score: bool = False, 

800 return_embedding: bool | None = False, 

801 ) -> list[Document]: 

802 """ 

803 Retrieves documents that are most similar to the query embedding using a vector similarity metric. 

804 

805 :param query_embedding: Embedding of the query. 

806 :param filters: A dictionary with filters to narrow down the search space. 

807 :param top_k: The number of top documents to retrieve. Default is 10. 

808 :param scale_score: Whether to scale the scores of the retrieved Documents. Default is False. 

809 :param return_embedding: Whether to return the embedding of the retrieved Documents. 

810 If not provided, the value of the `return_embedding` parameter set at component 

811 initialization will be used. Default is False. 

812 :returns: A list of the top_k documents most relevant to the query. 

813 :raises ValueError: if filters have invalid syntax. 

814 """ 

815 if len(query_embedding) == 0 or not isinstance(query_embedding[0], float): 

816 raise ValueError("query_embedding should be a non-empty list of floats.") 

817 

818 if filters: 

819 InMemoryDocumentStore._validate_filters(filters) 

820 all_documents = [ 

821 doc 

822 for doc in self.storage.values() 

823 if document_matches_filter( 

824 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison 

825 ) 

826 ] 

827 else: 

828 all_documents = list(self.storage.values()) 

829 

830 documents_with_embeddings = [doc for doc in all_documents if doc.embedding is not None] 

831 if len(documents_with_embeddings) == 0: 

832 logger.warning( 

833 "No Documents found with embeddings. Returning empty list. " 

834 "To generate embeddings, use a DocumentEmbedder." 

835 ) 

836 return [] 

837 if len(documents_with_embeddings) < len(all_documents): 

838 logger.info( 

839 "Skipping some Documents that don't have an embedding. To generate embeddings, use a DocumentEmbedder." 

840 ) 

841 

842 scores = self._compute_query_embedding_similarity_scores( 

843 embedding=query_embedding, documents=documents_with_embeddings, scale_score=scale_score 

844 ) 

845 

846 resolved_return_embedding = self.return_embedding if return_embedding is None else return_embedding 

847 

848 # create Documents with the similarity score for the top k results 

849 top_documents = [] 

850 for doc, score in sorted(zip(documents_with_embeddings, scores, strict=True), key=lambda x: x[1], reverse=True)[ 

851 :top_k 

852 ]: 

853 doc_fields = doc.to_dict() 

854 doc_fields["score"] = score 

855 if resolved_return_embedding is False: 

856 doc_fields["embedding"] = None 

857 top_documents.append(Document.from_dict(doc_fields)) 

858 

859 return top_documents 

860 

861 def _compute_query_embedding_similarity_scores( 

862 self, embedding: list[float], documents: list[Document], scale_score: bool = False 

863 ) -> list[float]: 

864 """ 

865 Computes the similarity scores between the query embedding and the embeddings of the documents. 

866 

867 :param embedding: Embedding of the query. 

868 :param documents: A list of Documents. 

869 :param scale_score: Whether to scale the scores of the Documents. Default is False. 

870 :returns: A list of scores. 

871 """ 

872 

873 query_embedding = np.array(embedding) 

874 if query_embedding.ndim == 1: 

875 query_embedding = np.expand_dims(a=query_embedding, axis=0) 

876 

877 try: 

878 document_embeddings = np.array([doc.embedding for doc in documents]) 

879 except ValueError as e: 

880 if "inhomogeneous shape" in str(e): 

881 raise DocumentStoreError( 

882 "The embedding size of all Documents should be the same. " 

883 "Please make sure that the Documents have been embedded with the same model." 

884 ) from e 

885 raise e 

886 if document_embeddings.ndim == 1: 

887 document_embeddings = np.expand_dims(a=document_embeddings, axis=0) 

888 

889 if self.embedding_similarity_function == "cosine": 

890 # cosine similarity is a normed dot product; guard against zero-norm vectors 

891 # (e.g. a zero embedding) which would otherwise divide by zero and yield NaN scores. 

892 query_norm = np.linalg.norm(x=query_embedding, axis=1, keepdims=True) 

893 document_norms = np.linalg.norm(x=document_embeddings, axis=1, keepdims=True) 

894 query_embedding /= np.where(query_norm == 0.0, 1.0, query_norm) 

895 document_embeddings /= np.where(document_norms == 0.0, 1.0, document_norms) 

896 

897 try: 

898 scores = np.dot(a=query_embedding, b=document_embeddings.T)[0].tolist() 

899 except ValueError as e: 

900 if "shapes" in str(e) and "not aligned" in str(e): 

901 raise DocumentStoreError( 

902 "The embedding size of the query should be the same as the embedding size of the Documents. " 

903 "Please make sure that the query has been embedded with the same model as the Documents." 

904 ) from e 

905 raise e 

906 

907 if scale_score: 

908 if self.embedding_similarity_function == "dot_product": 

909 scores = [expit(float(score / DOT_PRODUCT_SCALING_FACTOR)) for score in scores] 

910 elif self.embedding_similarity_function == "cosine": 

911 scores = [(score + 1) / 2 for score in scores] 

912 

913 return scores 

914 

915 async def count_documents_async(self) -> int: 

916 """ 

917 Returns the number of documents present in the DocumentStore. 

918 """ 

919 return len(self.storage.keys()) 

920 

921 async def filter_documents_async(self, filters: dict[str, Any] | None = None) -> list[Document]: 

922 """ 

923 Returns the documents that match the filters provided. 

924 

925 :param filters: The filters to apply. For a detailed specification of the filters, refer to the 

926 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering). 

927 :returns: A list of Documents that match the given filters. 

928 """ 

929 return await asyncio.get_running_loop().run_in_executor( 

930 self.executor, lambda: self.filter_documents(filters=filters) 

931 ) 

932 

933 async def write_documents_async( 

934 self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE 

935 ) -> int: 

936 """ 

937 Refer to the DocumentStore.write_documents() protocol documentation. 

938 

939 If `policy` is set to `DuplicatePolicy.NONE` defaults to `DuplicatePolicy.FAIL`. 

940 """ 

941 return await asyncio.get_running_loop().run_in_executor( 

942 self.executor, lambda: self.write_documents(documents=documents, policy=policy) 

943 ) 

944 

945 async def delete_documents_async(self, document_ids: list[str]) -> None: 

946 """ 

947 Deletes all documents with matching document_ids from the DocumentStore. 

948 

949 :param document_ids: The document_ids to delete. 

950 """ 

951 await asyncio.get_running_loop().run_in_executor( 

952 self.executor, lambda: self.delete_documents(document_ids=document_ids) 

953 ) 

954 

955 async def update_by_filter_async(self, filters: dict[str, Any], meta: dict[str, Any]) -> int: 

956 """ 

957 Updates the metadata of all documents that match the provided filters. 

958 

959 :param filters: The filters to apply to select documents for updating. 

960 For filter syntax, see filter_documents. 

961 :param meta: The metadata fields to update. These will be merged with existing metadata. 

962 :returns: The number of documents updated. 

963 """ 

964 return await asyncio.get_running_loop().run_in_executor( 

965 self.executor, lambda: self.update_by_filter(filters=filters, meta=meta) 

966 ) 

967 

968 async def count_documents_by_filter_async(self, filters: dict[str, Any]) -> int: 

969 """ 

970 Returns the number of documents that match the provided filters. 

971 

972 :param filters: The filters to apply. 

973 For a detailed specification of the filters, refer to the 

974 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering). 

975 :returns: The number of documents that match the filters. 

976 """ 

977 return await asyncio.get_running_loop().run_in_executor( 

978 self.executor, lambda: self.count_documents_by_filter(filters=filters) 

979 ) 

980 

981 async def count_unique_metadata_by_filter_async( 

982 self, filters: dict[str, Any], metadata_fields: list[str] 

983 ) -> dict[str, int]: 

984 """ 

985 Returns the number of unique values for each specified metadata field from documents matching the filters. 

986 

987 :param filters: The filters to apply. 

988 For a detailed specification of the filters, refer to the 

989 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering). 

990 :param metadata_fields: List of field names to count unique values for. 

991 Field names can include or omit the "meta." prefix. 

992 :returns: A dictionary mapping each metadata field name (without "meta." prefix) 

993 to the count of its unique values among the filtered documents. 

994 """ 

995 return await asyncio.get_running_loop().run_in_executor( 

996 self.executor, 

997 lambda: self.count_unique_metadata_by_filter(filters=filters, metadata_fields=metadata_fields), 

998 ) 

999 

1000 async def get_metadata_fields_info_async(self) -> dict[str, dict[str, str]]: 

1001 """ 

1002 Returns information about the metadata fields present in the stored documents. 

1003 

1004 Types are inferred from the stored values (keyword, int, float, boolean). 

1005 

1006 :returns: A dictionary mapping each metadata field name to a dict with a "type" key. 

1007 """ 

1008 return await asyncio.get_running_loop().run_in_executor(self.executor, self.get_metadata_fields_info) 

1009 

1010 async def get_metadata_field_min_max_async(self, metadata_field: str) -> dict[str, Any]: 

1011 """ 

1012 Returns the minimum and maximum values for the given metadata field across all documents. 

1013 

1014 :param metadata_field: The metadata field name. Can include or omit the "meta." prefix. 

1015 :returns: A dictionary with "min" and "max" keys. Returns `{"min": None, "max": None}` 

1016 if the field is missing or has no values. 

1017 """ 

1018 return await asyncio.get_running_loop().run_in_executor( 

1019 self.executor, lambda: self.get_metadata_field_min_max(metadata_field=metadata_field) 

1020 ) 

1021 

1022 async def get_metadata_field_unique_values_async( 

1023 self, 

1024 metadata_field: str, 

1025 search_term: str | None = None, 

1026 from_: int = 0, 

1027 size: int = 10, 

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

1029 ) -> tuple[list[Any], int]: 

1030 """ 

1031 Returns unique values for a metadata field, optionally filtered by a search term, with pagination. 

1032 

1033 :param metadata_field: The metadata field name. Can include or omit the "meta." prefix. 

1034 :param search_term: Optional search term to filter values, matched as a case-insensitive substring 

1035 against the metadata field's value. 

1036 :param from_: The offset to start returning values from (for pagination). 

1037 :param size: The maximum number of unique values to return. 

1038 :param filters: Optional filters to restrict the documents considered. 

1039 :returns: A tuple of (paginated list of unique values, total count of unique values). 

1040 """ 

1041 return await asyncio.get_running_loop().run_in_executor( 

1042 self.executor, 

1043 lambda: self.get_metadata_field_unique_values( 

1044 metadata_field=metadata_field, search_term=search_term, from_=from_, size=size, filters=filters 

1045 ), 

1046 ) 

1047 

1048 async def delete_all_documents_async(self) -> None: 

1049 """ 

1050 Deletes all documents in the document store. 

1051 """ 

1052 await asyncio.get_running_loop().run_in_executor(self.executor, self.delete_all_documents) 

1053 

1054 async def bm25_retrieval_async( 

1055 self, query: str, filters: dict[str, Any] | None = None, top_k: int = 10, scale_score: bool = False 

1056 ) -> list[Document]: 

1057 """ 

1058 Retrieves documents that are most relevant to the query using BM25 algorithm. 

1059 

1060 :param query: The query string. 

1061 :param filters: A dictionary with filters to narrow down the search space. 

1062 :param top_k: The number of top documents to retrieve. Default is 10. 

1063 :param scale_score: Whether to scale the scores of the retrieved documents. Default is False. 

1064 :returns: A list of the top_k documents most relevant to the query. 

1065 """ 

1066 return await asyncio.get_running_loop().run_in_executor( 

1067 self.executor, 

1068 lambda: self.bm25_retrieval(query=query, filters=filters, top_k=top_k, scale_score=scale_score), 

1069 ) 

1070 

1071 async def embedding_retrieval_async( 

1072 self, 

1073 query_embedding: list[float], 

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

1075 top_k: int = 10, 

1076 scale_score: bool = False, 

1077 return_embedding: bool = False, 

1078 ) -> list[Document]: 

1079 """ 

1080 Retrieves documents that are most similar to the query embedding using a vector similarity metric. 

1081 

1082 :param query_embedding: Embedding of the query. 

1083 :param filters: A dictionary with filters to narrow down the search space. 

1084 :param top_k: The number of top documents to retrieve. Default is 10. 

1085 :param scale_score: Whether to scale the scores of the retrieved Documents. Default is False. 

1086 :param return_embedding: Whether to return the embedding of the retrieved Documents. Default is False. 

1087 :returns: A list of the top_k documents most relevant to the query. 

1088 """ 

1089 return await asyncio.get_running_loop().run_in_executor( 

1090 self.executor, 

1091 lambda: self.embedding_retrieval( 

1092 query_embedding=query_embedding, 

1093 filters=filters, 

1094 top_k=top_k, 

1095 scale_score=scale_score, 

1096 return_embedding=return_embedding, 

1097 ), 

1098 )