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
« 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
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
17import numpy as np
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
26logger = logging.getLogger(__name__)
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
39@dataclass
40class BM25DocumentStats:
41 """
42 A dataclass for managing document statistics for BM25 retrieval.
44 :param freq_token: A Counter of token frequencies in the document.
45 :param doc_len: Number of tokens in the document.
46 """
48 freq_token: dict[str, int]
49 doc_len: int
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] = {}
60class InMemoryDocumentStore:
61 """
62 Stores data in-memory. It's ephemeral and cannot be saved to disk.
63 """
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.
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
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())
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()
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
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
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)
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)
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
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
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
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
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
180 def _dispatch_bm25(self) -> "Callable[[str, list[Document]], list[tuple[Document, float]]]":
181 """
182 Select the correct BM25 algorithm based on user specification.
184 :returns:
185 The BM25 algorithm method.
186 """
187 table = {"BM25Okapi": self._score_bm25okapi, "BM25L": self._score_bm25l, "BM25Plus": self._score_bm25plus}
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]
193 def _tokenize_bm25(self, text: str) -> list[str]:
194 """
195 Tokenize text using the BM25 tokenization regex.
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.
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)
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.
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)
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
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)
241 idf = _compute_idf(self._tokenize_bm25(query))
242 bm25_attr = {doc.id: self._bm25_attr[doc.id] for doc in documents}
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
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))
255 return ret
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.
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)
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 = []
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)
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}
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
299 idf = _compute_idf(self._tokenize_bm25(query))
300 bm25_attr = {doc.id: self._bm25_attr[doc.id] for doc in documents}
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
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))
313 return ret
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.
319 This implementation follows the document on BM25 Wikipedia page,
320 which add 1 (smoothing factor) to document frequency when computing IDF.
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)
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
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
349 idf = _compute_idf(self._tokenize_bm25(query))
350 bm25_attr = {doc.id: self._bm25_attr[doc.id] for doc in documents}
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
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))
363 return ret
365 def to_dict(self) -> dict[str, Any]:
366 """
367 Serializes the component to a dictionary.
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 )
384 @classmethod
385 def from_dict(cls, data: dict[str, Any]) -> "InMemoryDocumentStore":
386 """
387 Deserializes the component from a dictionary.
389 :param data:
390 The dictionary to deserialize from.
391 :returns:
392 The deserialized component.
393 """
394 return default_from_dict(cls, data)
396 def save_to_disk(self, path: str) -> None:
397 """
398 Write the database and its data to disk as a JSON file.
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)
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.
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
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
429 raise FileNotFoundError(f"File {path} not found.")
431 def count_documents(self) -> int:
432 """
433 Returns the number of documents present in the DocumentStore.
434 """
435 return len(self.storage.keys())
437 def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Document]:
438 """
439 Returns the documents that match the filters provided.
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())
457 if not self.return_embedding:
458 docs = [replace(doc, embedding=None) for doc in docs]
460 return docs
462 def write_documents(self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE) -> int:
463 """
464 Refer to the DocumentStore.write_documents() protocol documentation.
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.")
475 if policy == DuplicatePolicy.NONE:
476 policy = DuplicatePolicy.FAIL
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
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])
494 tokens = []
495 if document.content is not None:
496 tokens = self._tokenize_bm25(document.content)
498 self.storage[document.id] = document
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
507 def delete_documents(self, document_ids: list[str]) -> None:
508 """
509 Deletes all documents with matching document_ids from the DocumentStore.
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]
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
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
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 )
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()
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.
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)
577 def delete_by_filter(self, filters: dict[str, Any]) -> int:
578 """
579 Deletes all documents that match the provided filters.
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)
598 def count_documents_by_filter(self, filters: dict[str, Any]) -> int:
599 """
600 Returns the number of documents that match the provided filters.
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)
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.
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())
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
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.
653 Types are inferred from the stored values (keyword, int, float, boolean).
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()}
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.
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}
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.
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)
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()}
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)
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.
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")
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
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 []
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)
766 results = sorted(scored_documents, key=lambda x: x[1], reverse=True)[:top_k]
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
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)
779 if not negatives_are_valid and score <= 0.0:
780 continue
782 doc_fields = doc.to_dict()
783 doc_fields["score"] = score
785 if not self.return_embedding and "embedding" in doc_fields:
786 doc_fields.pop("embedding")
788 return_document = Document.from_dict(doc_fields)
790 return_documents.append(return_document)
792 return return_documents
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.
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.")
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())
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 )
842 scores = self._compute_query_embedding_similarity_scores(
843 embedding=query_embedding, documents=documents_with_embeddings, scale_score=scale_score
844 )
846 resolved_return_embedding = self.return_embedding if return_embedding is None else return_embedding
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))
859 return top_documents
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.
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 """
873 query_embedding = np.array(embedding)
874 if query_embedding.ndim == 1:
875 query_embedding = np.expand_dims(a=query_embedding, axis=0)
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)
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)
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
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]
913 return scores
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())
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.
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 )
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.
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 )
945 async def delete_documents_async(self, document_ids: list[str]) -> None:
946 """
947 Deletes all documents with matching document_ids from the DocumentStore.
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 )
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.
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 )
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.
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 )
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.
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 )
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.
1004 Types are inferred from the stored values (keyword, int, float, boolean).
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)
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.
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 )
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.
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 )
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)
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.
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 )
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.
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 )