Coverage for haystack/document_stores/in_memory/document_store.py: 96%
446 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +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 unicodedata
10import uuid
11from collections import Counter
12from collections.abc import Callable, Iterable
13from concurrent.futures import ThreadPoolExecutor
14from dataclasses import dataclass, replace
15from pathlib import Path
16from typing import Any, Literal
18import numpy as np
20from haystack import default_from_dict, default_to_dict, logging
21from haystack.dataclasses import Document
22from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
23from haystack.document_stores.types import DuplicatePolicy
24from haystack.utils import expit
25from haystack.utils.filters import document_matches_filter
27logger = logging.getLogger(__name__)
29# document scores are essentially unbounded and will be scaled to values between 0 and 1 if scale_score is set to
30# True (default). Scaling uses the expit function (inverse of the logit function) after applying a scaling factor
31# (e.g., BM25_SCALING_FACTOR for the bm25_retrieval method).
32# Larger scaling factor decreases scaled scores. For example, an input of 10 is scaled to 0.99 with
33# BM25_SCALING_FACTOR=2 but to 0.78 with BM25_SCALING_FACTOR=8 (default). The defaults were chosen empirically.
34# Increase the default if most unscaled scores are larger than expected (>30) and otherwise would incorrectly all be
35# mapped to scores ~1.
36BM25_SCALING_FACTOR = 8
37DOT_PRODUCT_SCALING_FACTOR = 100
39# Unicode ranges for scripts that are written without spaces between words (Chinese, Japanese, Korean). These
40# characters are tokenized one per token so that bare-term queries for individual characters or short words can
41# match, while every other script keeps the usual word-based behaviour. The kana ranges intentionally exclude
42# punctuation such as the katakana middle dot (・) so it does not become a token on its own.
43_CJK_CHAR_CLASS = (
44 r"\u1100-\u11ff" # Hangul Jamo
45 r"\ua960-\ua97f" # Hangul Jamo Extended-A
46 r"\ud7b0-\ud7ff" # Hangul Jamo Extended-B
47 r"\u3130-\u318f" # Hangul Compatibility Jamo
48 r"\uac00-\ud7af" # Hangul Syllables
49 r"\u3041-\u3096" # Hiragana letters (excludes combining marks and kana iteration marks)
50 r"\u30a1-\u30fa\u30fc" # Katakana letters + prolonged sound mark, excludes ・ middle dot
51 r"\u31f0-\u31ff" # Katakana Phonetic Extensions
52 r"\u3400-\u4dbf" # CJK Unified Ideographs Extension A
53 r"\u4e00-\u9fff" # CJK Unified Ideographs
54 r"\uf900-\ufaff" # CJK Compatibility Ideographs
55 r"\uff66-\uff9f" # Halfwidth Katakana
56 r"\uffa0-\uffdc" # Halfwidth Hangul Jamo
57)
58_DEFAULT_BM25_TOKENIZATION_REGEX = rf"[^\W{_CJK_CHAR_CLASS}]+|[{_CJK_CHAR_CLASS}]"
61def _make_metadata_value_hashable(value: Any) -> Any:
62 """Convert nested metadata values into values that can be used for deduplication."""
63 if isinstance(value, list):
64 return ("list", tuple(_make_metadata_value_hashable(item) for item in value))
65 if isinstance(value, tuple):
66 return ("tuple", tuple(_make_metadata_value_hashable(item) for item in value))
67 if isinstance(value, dict):
68 return (
69 "dict",
70 frozenset(
71 (_make_metadata_value_hashable(key), _make_metadata_value_hashable(item)) for key, item in value.items()
72 ),
73 )
74 return value
77@dataclass
78class BM25DocumentStats:
79 """
80 A dataclass for managing document statistics for BM25 retrieval.
82 :param freq_token: A Counter of token frequencies in the document.
83 :param doc_len: Number of tokens in the document.
84 """
86 freq_token: dict[str, int]
87 doc_len: int
90# Process-global storage, keyed by index, for stores created with shared=True. Lets instances sharing an index
91# operate on the same data. Non-shared stores keep their data instance-local instead (see `__init__`).
92_STORAGES: dict[str, dict[str, Document]] = {}
93_BM25_STATS_STORAGES: dict[str, dict[str, BM25DocumentStats]] = {}
94_AVERAGE_DOC_LEN_STORAGES: dict[str, float] = {}
95_FREQ_VOCAB_FOR_IDF_STORAGES: dict[str, Counter] = {}
98class InMemoryDocumentStore:
99 """
100 Stores data in-memory. It's ephemeral and cannot be saved to disk.
101 """
103 def __init__(
104 self,
105 bm25_tokenization_regex: str = _DEFAULT_BM25_TOKENIZATION_REGEX,
106 bm25_algorithm: Literal["BM25Okapi", "BM25L", "BM25Plus"] = "BM25L",
107 bm25_parameters: dict | None = None,
108 embedding_similarity_function: Literal["dot_product", "cosine"] = "dot_product",
109 index: str | None = None,
110 shared: bool = True,
111 async_executor: ThreadPoolExecutor | None = None,
112 return_embedding: bool = True,
113 *,
114 strict_datetime_comparison: bool = False,
115 ) -> None:
116 """
117 Initializes the DocumentStore.
119 :param bm25_tokenization_regex:
120 The regular expression used to tokenize the text for BM25 retrieval. The default groups word
121 characters into word tokens and splits Chinese, Japanese and Korean text into one token per character.
122 Text is lowercased and NFC-normalized before tokenization.
123 :param bm25_algorithm: The BM25 algorithm to use. One of "BM25Okapi", "BM25L", or "BM25Plus".
124 :param bm25_parameters: Parameters for BM25 implementation in a dictionary format.
125 For example: `{'k1':1.5, 'b':0.75, 'epsilon':0.25}`
126 You can learn more about these parameters by visiting https://github.com/dorianbrown/rank_bm25.
127 :param embedding_similarity_function: The similarity function used to compare Documents embeddings.
128 One of "dot_product" (default) or "cosine". To choose the most appropriate function, look for information
129 about your embedding model.
130 :param index: A specific index to store the documents. If not specified, a random UUID is used.
131 When `shared` is True, instances using the same index share the same documents.
132 :param shared: Whether the documents live in process-global storage shared across instances using the same
133 index (True, the default), or are kept instance-local and freed when this instance is garbage collected
134 (False). Shared storage persists for the lifetime of the process, so prefer `shared=False` for stores
135 that are created frequently (for example per request) to avoid unbounded memory growth.
136 :param async_executor:
137 Optional ThreadPoolExecutor to use for async calls. If not provided, a single-threaded
138 executor will be initialized and used.
139 :param return_embedding: Whether to return the embedding of the retrieved Documents. Default is True.
140 :param strict_datetime_comparison:
141 If `True`, timezone-naive and timezone-aware datetimes never match each other in filters.
142 If `False` (the default), the timezone from the aware datetime is copied to the naive one before
143 comparing.
144 """
145 self.bm25_tokenization_regex = bm25_tokenization_regex
146 self.tokenizer = re.compile(bm25_tokenization_regex).findall
148 # Shared stores keep their data in the process-global dicts keyed by index, so instances sharing an index
149 # operate on the same documents. Non-shared stores keep their data instance-local so it is freed with the
150 # instance instead of accumulating in the globals.
151 self._shared = shared
152 self.index = index if index is not None else str(uuid.uuid4())
154 if self._shared:
155 if self.index not in _STORAGES:
156 _STORAGES[self.index] = {}
157 _BM25_STATS_STORAGES[self.index] = {}
158 _AVERAGE_DOC_LEN_STORAGES[self.index] = 0.0
159 _FREQ_VOCAB_FOR_IDF_STORAGES[self.index] = Counter()
160 else:
161 self._local_storage: dict[str, Document] = {}
162 self._local_bm25_attr: dict[str, BM25DocumentStats] = {}
163 self._local_avg_doc_len: float = 0.0
164 self._local_freq_vocab_for_idf: Counter = Counter()
166 self.bm25_algorithm = bm25_algorithm
167 self.bm25_algorithm_inst = self._dispatch_bm25()
168 self.bm25_parameters = bm25_parameters or {}
169 self.embedding_similarity_function = embedding_similarity_function
171 # keep track of whether we own the executor if we created it we must also clean it up
172 self._owns_executor = async_executor is None
173 self.executor = (
174 ThreadPoolExecutor(thread_name_prefix=f"async-inmemory-docstore-executor-{id(self)}", max_workers=1)
175 if async_executor is None
176 else async_executor
177 )
178 self.return_embedding = return_embedding
179 self.strict_datetime_comparison = strict_datetime_comparison
181 def __del__(self) -> None:
182 """
183 Cleanup when the instance is being destroyed.
184 """
185 if hasattr(self, "_owns_executor") and self._owns_executor and hasattr(self, "executor"):
186 self.executor.shutdown(wait=True)
188 def shutdown(self) -> None:
189 """
190 Explicitly shutdown the executor if we own it.
191 """
192 if self._owns_executor:
193 self.executor.shutdown(wait=True)
195 @property
196 def storage(self) -> dict[str, Document]:
197 """
198 Utility property that returns the storage used by this instance of InMemoryDocumentStore.
199 """
200 return _STORAGES[self.index] if self._shared else self._local_storage
202 @property
203 def _bm25_attr(self) -> dict[str, BM25DocumentStats]:
204 return _BM25_STATS_STORAGES[self.index] if self._shared else self._local_bm25_attr
206 @property
207 def _avg_doc_len(self) -> float:
208 return _AVERAGE_DOC_LEN_STORAGES[self.index] if self._shared else self._local_avg_doc_len
210 @_avg_doc_len.setter
211 def _avg_doc_len(self, value: float) -> None:
212 if self._shared:
213 _AVERAGE_DOC_LEN_STORAGES[self.index] = value
214 else:
215 self._local_avg_doc_len = value
217 @property
218 def _freq_vocab_for_idf(self) -> Counter:
219 return _FREQ_VOCAB_FOR_IDF_STORAGES[self.index] if self._shared else self._local_freq_vocab_for_idf
221 def _dispatch_bm25(self) -> "Callable[[str, list[Document]], list[tuple[Document, float]]]":
222 """
223 Select the correct BM25 algorithm based on user specification.
225 :returns:
226 The BM25 algorithm method.
227 """
228 table = {"BM25Okapi": self._score_bm25okapi, "BM25L": self._score_bm25l, "BM25Plus": self._score_bm25plus}
230 if self.bm25_algorithm not in table:
231 raise ValueError(f"BM25 algorithm '{self.bm25_algorithm}' is not supported.")
232 return table[self.bm25_algorithm]
234 def _tokenize_bm25(self, text: str) -> list[str]:
235 """
236 Tokenize text using the BM25 tokenization regex.
238 Here we explicitly create a tokenization method to encapsulate
239 all pre-processing logic used to create BM25 tokens, such as
240 lowercasing and Unicode NFC normalization. This helps track the exact tokenization process used for
241 BM25 scoring at any given time.
243 :param text:
244 The text to tokenize.
245 :returns:
246 A list of tokens.
247 """
248 text = unicodedata.normalize("NFC", text.lower())
249 return self.tokenizer(text)
251 def _score_bm25l(self, query: str, documents: list[Document]) -> list[tuple[Document, float]]:
252 """
253 Calculate BM25L scores for the given query and filtered documents.
255 :param query:
256 The query string.
257 :param documents:
258 The list of documents to score, should be produced by
259 the filter_documents method; may be an empty list.
260 :returns:
261 A list of tuples, each containing a Document and its BM25L score.
262 """
263 k = self.bm25_parameters.get("k1", 1.5)
264 b = self.bm25_parameters.get("b", 0.75)
265 delta = self.bm25_parameters.get("delta", 0.5)
267 def _compute_idf(tokens: list[str]) -> dict[str, float]:
268 """Per-token IDF computation for all tokens."""
269 idf = {}
270 n_corpus = len(self._bm25_attr)
271 for tok in tokens:
272 n = self._freq_vocab_for_idf.get(tok, 0)
273 idf[tok] = math.log((n_corpus + 1.0) / (n + 0.5)) * int(n != 0)
274 return idf
276 def _compute_tf(token: str, freq: dict[str, int], doc_len: int) -> float:
277 """Per-token BM25L computation."""
278 freq_term = freq.get(token, 0.0)
279 # The lower bound (delta) only applies to terms that occur in the document;
280 # a missing term contributes nothing.
281 if freq_term == 0:
282 return 0.0
283 ctd = freq_term / (1 - b + b * doc_len / self._avg_doc_len)
284 return (1.0 + k) * (ctd + delta) / (k + ctd + delta)
286 idf = _compute_idf(self._tokenize_bm25(query))
287 bm25_attr = {doc.id: self._bm25_attr[doc.id] for doc in documents}
289 ret = []
290 for doc in documents:
291 doc_stats = bm25_attr[doc.id]
292 freq = doc_stats.freq_token
293 doc_len = doc_stats.doc_len
295 score = 0.0
296 for tok in idf.keys():
297 score += idf[tok] * _compute_tf(tok, freq, doc_len)
298 ret.append((doc, score))
300 return ret
302 def _score_bm25okapi(self, query: str, documents: list[Document]) -> list[tuple[Document, float]]:
303 """
304 Calculate BM25Okapi scores for the given query and filtered documents.
306 :param query:
307 The query string.
308 :param documents:
309 The list of documents to score, should be produced by
310 the filter_documents method; may be an empty list.
311 :returns:
312 A list of tuples, each containing a Document and its BM25Okapi score.
313 """
314 k = self.bm25_parameters.get("k1", 1.5)
315 b = self.bm25_parameters.get("b", 0.75)
316 epsilon = self.bm25_parameters.get("epsilon", 0.25)
318 def _compute_idf(tokens: list[str]) -> dict[str, float]:
319 """Per-token IDF computation for all tokens."""
320 sum_idf = 0.0
321 neg_idf_tokens = []
323 # Although this is a global statistic, we compute it here
324 # to make the computation more self-contained. And the
325 # complexity is O(vocab_size), which is acceptable.
326 idf = {}
327 for tok, n in self._freq_vocab_for_idf.items():
328 idf[tok] = math.log((len(self._bm25_attr) - n + 0.5) / (n + 0.5))
329 sum_idf += idf[tok]
330 if idf[tok] < 0:
331 neg_idf_tokens.append(tok)
333 eps = epsilon * sum_idf / len(self._freq_vocab_for_idf)
334 for tok in neg_idf_tokens:
335 idf[tok] = eps
336 return {tok: idf.get(tok, 0.0) for tok in tokens}
338 def _compute_tf(token: str, freq: dict[str, int], doc_len: int) -> float:
339 """Per-token BM25Okapi computation."""
340 freq_term = freq.get(token, 0.0)
341 freq_norm = freq_term + k * (1 - b + b * doc_len / self._avg_doc_len)
342 return freq_term * (1.0 + k) / freq_norm
344 idf = _compute_idf(self._tokenize_bm25(query))
345 bm25_attr = {doc.id: self._bm25_attr[doc.id] for doc in documents}
347 ret = []
348 for doc in documents:
349 doc_stats = bm25_attr[doc.id]
350 freq = doc_stats.freq_token
351 doc_len = doc_stats.doc_len
353 score = 0.0
354 for tok in idf.keys():
355 score += idf[tok] * _compute_tf(tok, freq, doc_len)
356 ret.append((doc, score))
358 return ret
360 def _score_bm25plus(self, query: str, documents: list[Document]) -> list[tuple[Document, float]]:
361 """
362 Calculate BM25+ scores for the given query and filtered documents.
364 This implementation follows the document on BM25 Wikipedia page,
365 which add 1 (smoothing factor) to document frequency when computing IDF.
367 :param query:
368 The query string.
369 :param documents:
370 The list of documents to score, should be produced by
371 the filter_documents method; may be an empty list.
372 :returns:
373 A list of tuples, each containing a Document and its BM25+ score.
374 """
375 k = self.bm25_parameters.get("k1", 1.5)
376 b = self.bm25_parameters.get("b", 0.75)
377 delta = self.bm25_parameters.get("delta", 1.0)
379 def _compute_idf(tokens: list[str]) -> dict[str, float]:
380 """Per-token IDF computation."""
381 idf = {}
382 n_corpus = len(self._bm25_attr)
383 for tok in tokens:
384 n = self._freq_vocab_for_idf.get(tok, 0)
385 idf[tok] = math.log(1 + (n_corpus - n + 0.5) / (n + 0.5)) * int(n != 0)
386 return idf
388 def _compute_tf(token: str, freq: dict[str, int], doc_len: float) -> float:
389 """Per-token normalized term frequency."""
390 freq_term = freq.get(token, 0.0)
391 # The lower bound (delta) only applies to terms that occur in the document;
392 # a missing term contributes nothing.
393 if freq_term == 0:
394 return 0.0
395 freq_damp = k * (1 - b + b * doc_len / self._avg_doc_len)
396 return freq_term * (1.0 + k) / (freq_term + freq_damp) + delta
398 idf = _compute_idf(self._tokenize_bm25(query))
399 bm25_attr = {doc.id: self._bm25_attr[doc.id] for doc in documents}
401 ret = []
402 for doc in documents:
403 doc_stats = bm25_attr[doc.id]
404 freq = doc_stats.freq_token
405 doc_len = doc_stats.doc_len
407 score = 0.0
408 for tok in idf.keys():
409 score += idf[tok] * _compute_tf(tok, freq, doc_len)
410 ret.append((doc, score))
412 return ret
414 def to_dict(self) -> dict[str, Any]:
415 """
416 Serializes the component to a dictionary.
418 :returns:
419 Dictionary with serialized data.
420 """
421 return default_to_dict(
422 self,
423 bm25_tokenization_regex=self.bm25_tokenization_regex,
424 bm25_algorithm=self.bm25_algorithm,
425 bm25_parameters=self.bm25_parameters,
426 embedding_similarity_function=self.embedding_similarity_function,
427 index=self.index,
428 shared=self._shared,
429 return_embedding=self.return_embedding,
430 strict_datetime_comparison=self.strict_datetime_comparison,
431 )
433 @classmethod
434 def from_dict(cls, data: dict[str, Any]) -> "InMemoryDocumentStore":
435 """
436 Deserializes the component from a dictionary.
438 :param data:
439 The dictionary to deserialize from.
440 :returns:
441 The deserialized component.
442 """
443 return default_from_dict(cls, data)
445 def save_to_disk(self, path: str) -> None:
446 """
447 Write the database and its data to disk as a JSON file.
449 :param path: The path to the JSON file.
450 """
451 data: dict[str, Any] = self.to_dict()
452 data["documents"] = [doc.to_dict(flatten=False) for doc in self.storage.values()]
453 with open(path, "w") as f:
454 json.dump(data, f)
456 @classmethod
457 def load_from_disk(cls, path: str) -> "InMemoryDocumentStore":
458 """
459 Load the database and its data from disk as a JSON file.
461 :param path: The path to the JSON file.
462 :returns: The loaded InMemoryDocumentStore.
463 """
464 if Path(path).exists():
465 try:
466 with open(path) as f:
467 data = json.load(f)
468 except Exception as e:
469 raise DocumentStoreError(f"Error loading InMemoryDocumentStore from disk. error: {e}") from e
471 documents = data.pop("documents")
472 cls_object = default_from_dict(cls, data)
473 cls_object.write_documents(
474 documents=[Document.from_dict(doc) for doc in documents], policy=DuplicatePolicy.OVERWRITE
475 )
476 return cls_object
478 raise FileNotFoundError(f"File {path} not found.")
480 def count_documents(self) -> int:
481 """
482 Returns the number of documents present in the DocumentStore.
483 """
484 return len(self.storage.keys())
486 def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Document]:
487 """
488 Returns the documents that match the filters provided.
490 :param filters: The filters to apply. For a detailed specification of the filters, refer to the
491 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering).
492 :returns: A list of Documents that match the given filters.
493 """
494 if filters:
495 InMemoryDocumentStore._validate_filters(filters)
496 docs = [
497 doc
498 for doc in self.storage.values()
499 if document_matches_filter(
500 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison
501 )
502 ]
503 else:
504 docs = list(self.storage.values())
506 if not self.return_embedding:
507 docs = [replace(doc, embedding=None) for doc in docs]
509 return docs
511 def write_documents(self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE) -> int:
512 """
513 Refer to the DocumentStore.write_documents() protocol documentation.
515 If `policy` is set to `DuplicatePolicy.NONE` defaults to `DuplicatePolicy.FAIL`.
516 """
517 if (
518 not isinstance(documents, Iterable)
519 or isinstance(documents, str)
520 or any(not isinstance(doc, Document) for doc in documents)
521 ):
522 raise ValueError("Please provide a list of Documents.")
524 if policy == DuplicatePolicy.NONE:
525 policy = DuplicatePolicy.FAIL
527 written_documents = len(documents)
528 for document in documents:
529 if policy != DuplicatePolicy.OVERWRITE and document.id in self.storage.keys():
530 if policy == DuplicatePolicy.FAIL:
531 raise DuplicateDocumentError(f"ID '{document.id}' already exists.")
532 if policy == DuplicatePolicy.SKIP:
533 logger.warning("ID '{document_id}' already exists", document_id=document.id)
534 written_documents -= 1
535 continue
537 # Since the statistics are updated in an incremental manner,
538 # we need to explicitly remove the existing document to revert
539 # the statistics before updating them with the new document.
540 if document.id in self.storage.keys():
541 self.delete_documents([document.id])
543 tokens = []
544 if document.content is not None:
545 tokens = self._tokenize_bm25(document.content)
547 self.storage[document.id] = document
549 self._bm25_attr[document.id] = BM25DocumentStats(Counter(tokens), len(tokens))
550 self._freq_vocab_for_idf.update(set(tokens))
551 # Update avg doc len based on the new document and the previous average
552 n_docs = len(self._bm25_attr)
553 self._avg_doc_len = (len(tokens) + self._avg_doc_len * (n_docs - 1)) / n_docs
554 return written_documents
556 def delete_documents(self, document_ids: list[str]) -> None:
557 """
558 Deletes all documents with matching document_ids from the DocumentStore.
560 :param document_ids: The document_ids to delete.
561 """
562 for doc_id in document_ids:
563 if doc_id not in self.storage.keys():
564 continue
565 del self.storage[doc_id]
567 # Update statistics accordingly
568 doc_stats = self._bm25_attr.pop(doc_id)
569 freq = doc_stats.freq_token
570 doc_len = doc_stats.doc_len
572 self._freq_vocab_for_idf.subtract(Counter(freq.keys()))
573 for token in freq:
574 if self._freq_vocab_for_idf[token] <= 0:
575 del self._freq_vocab_for_idf[token]
576 try:
577 self._avg_doc_len = (self._avg_doc_len * (len(self._bm25_attr) + 1) - doc_len) / len(self._bm25_attr)
578 except ZeroDivisionError:
579 self._avg_doc_len = 0
581 @staticmethod
582 def _validate_filters(filters: dict[str, Any] | None) -> None:
583 if filters and "operator" not in filters and "conditions" not in filters:
584 raise ValueError(
585 "Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details."
586 )
588 def delete_all_documents(self) -> None:
589 """
590 Deletes all documents in the document store.
591 """
592 if self._shared:
593 _STORAGES[self.index] = {}
594 _BM25_STATS_STORAGES[self.index] = {}
595 _AVERAGE_DOC_LEN_STORAGES[self.index] = 0.0
596 _FREQ_VOCAB_FOR_IDF_STORAGES[self.index] = Counter()
597 else:
598 self._local_storage = {}
599 self._local_bm25_attr = {}
600 self._local_avg_doc_len = 0.0
601 self._local_freq_vocab_for_idf = Counter()
603 def update_by_filter(self, filters: dict[str, Any], meta: dict[str, Any]) -> int:
604 """
605 Updates the metadata of all documents that match the provided filters.
607 :param filters: The filters to apply to select documents for updating.
608 For filter syntax, see filter_documents.
609 :param meta: The metadata fields to update. These will be merged with existing metadata.
610 :returns: The number of documents updated.
611 :raises ValueError: if filters have invalid syntax.
612 """
613 InMemoryDocumentStore._validate_filters(filters)
614 matching = [
615 doc
616 for doc in self.storage.values()
617 if document_matches_filter(
618 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison
619 )
620 ]
621 for doc in matching:
622 doc.meta.update(meta)
623 self.storage[doc.id] = doc
624 return len(matching)
626 def delete_by_filter(self, filters: dict[str, Any]) -> int:
627 """
628 Deletes all documents that match the provided filters.
630 :param filters: The filters to apply to select documents for deletion.
631 For filter syntax, see filter_documents.
632 :returns: The number of documents deleted.
633 :raises ValueError: if filters have invalid syntax.
634 """
635 InMemoryDocumentStore._validate_filters(filters)
636 matching = [
637 doc
638 for doc in self.storage.values()
639 if document_matches_filter(
640 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison
641 )
642 ]
643 doc_ids = [doc.id for doc in matching]
644 self.delete_documents(doc_ids)
645 return len(doc_ids)
647 def count_documents_by_filter(self, filters: dict[str, Any]) -> int:
648 """
649 Returns the number of documents that match the provided filters.
651 :param filters: The filters to apply.
652 For a detailed specification of the filters, refer to the
653 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering).
654 :returns: The number of documents that match the filters.
655 """
656 if filters:
657 InMemoryDocumentStore._validate_filters(filters)
658 return sum(
659 1
660 for doc in self.storage.values()
661 if document_matches_filter(
662 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison
663 )
664 )
665 return len(self.storage)
667 def count_unique_metadata_by_filter(self, filters: dict[str, Any], metadata_fields: list[str]) -> dict[str, int]:
668 """
669 Returns the number of unique values for each specified metadata field from documents matching the filters.
671 JSON-serializable metadata values, including nested lists and dictionaries, are supported.
673 :param filters: The filters to apply.
674 For a detailed specification of the filters, refer to the
675 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering).
676 :param metadata_fields: List of field names to count unique values for.
677 Field names can include or omit the "meta." prefix.
678 :returns: A dictionary mapping each metadata field name (without "meta." prefix)
679 to the count of its unique values among the filtered documents.
680 """
681 if filters:
682 InMemoryDocumentStore._validate_filters(filters)
683 docs = [
684 doc
685 for doc in self.storage.values()
686 if document_matches_filter(
687 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison
688 )
689 ]
690 else:
691 docs = list(self.storage.values())
693 result: dict[str, int] = {}
694 for field in metadata_fields:
695 key = field.removeprefix("meta.") if field.startswith("meta.") else field
696 values = {
697 _make_metadata_value_hashable(doc.meta.get(key))
698 for doc in docs
699 if key in doc.meta and doc.meta[key] is not None
700 }
701 result[key] = len(values)
702 return result
704 def get_metadata_fields_info(self) -> dict[str, dict[str, str]]:
705 """
706 Returns information about the metadata fields present in the stored documents.
708 Types are inferred from the stored values (keyword, int, float, boolean).
710 :returns: A dictionary mapping each metadata field name to a dict with a "type" key.
711 """
712 type_map: dict[str, str] = {}
713 for doc in self.storage.values():
714 for key, value in doc.meta.items():
715 if value is None:
716 continue
717 if isinstance(value, bool):
718 type_map[key] = "boolean"
719 elif isinstance(value, int):
720 type_map[key] = "int"
721 elif isinstance(value, float):
722 type_map[key] = "float"
723 else:
724 type_map[key] = "keyword"
725 return {k: {"type": v} for k, v in type_map.items()}
727 def get_metadata_field_min_max(self, metadata_field: str) -> dict[str, Any]:
728 """
729 Returns the minimum and maximum values for the given metadata field across all documents.
731 :param metadata_field: The metadata field name. Can include or omit the "meta." prefix.
732 :returns: A dictionary with "min" and "max" keys. Returns `{"min": None, "max": None}`
733 if the field is missing or has no values.
734 """
735 key = metadata_field.removeprefix("meta.") if metadata_field.startswith("meta.") else metadata_field
736 values = [
737 doc.meta[key]
738 for doc in self.storage.values()
739 if key in doc.meta and doc.meta[key] is not None and isinstance(doc.meta[key], (int, float, str))
740 ]
741 if not values:
742 return {"min": None, "max": None}
743 try:
744 return {"min": min(values), "max": max(values)}
745 except TypeError:
746 return {"min": None, "max": None}
748 def get_metadata_field_unique_values(
749 self,
750 metadata_field: str,
751 search_term: str | None = None,
752 from_: int = 0,
753 size: int = 10,
754 filters: dict[str, Any] | None = None,
755 ) -> tuple[list[Any], int]:
756 """
757 Returns unique values for a metadata field, optionally filtered by a search term, with pagination.
759 JSON-serializable metadata values, including nested lists and dictionaries, are supported.
761 :param metadata_field: The metadata field name. Can include or omit the "meta." prefix.
762 :param search_term: Optional search term to filter values, matched as a case-insensitive substring
763 against the metadata field's value.
764 :param from_: The offset to start returning values from (for pagination).
765 :param size: The maximum number of unique values to return.
766 :param filters: Optional filters to restrict the documents considered.
767 :returns: A tuple of (paginated list of unique values, total count of unique values).
768 """
769 key = metadata_field.removeprefix("meta.") if metadata_field.startswith("meta.") else metadata_field
770 unique_values: dict[tuple[str, Any], Any] = {}
771 for doc in self.filter_documents(filters=filters):
772 value = doc.meta.get(key)
773 if value is not None:
774 unique_values.setdefault((type(value).__name__, _make_metadata_value_hashable(value)), value)
776 if search_term:
777 search_term_lower = search_term.lower()
778 unique_values = {k: v for k, v in unique_values.items() if search_term_lower in str(v).lower()}
780 sorted_keys = sorted(unique_values, key=lambda k: (str(unique_values[k]), k[0]))
781 paginated_keys = sorted_keys[from_ : from_ + size]
782 return [unique_values[k] for k in paginated_keys], len(sorted_keys)
784 def bm25_retrieval(
785 self, query: str, filters: dict[str, Any] | None = None, top_k: int = 10, scale_score: bool = False
786 ) -> list[Document]:
787 """
788 Retrieves documents that are most relevant to the query using BM25 algorithm.
790 :param query: The query string.
791 :param filters: A dictionary with filters to narrow down the search space.
792 :param top_k: The number of top documents to retrieve. Default is 10.
793 :param scale_score: Whether to scale the scores of the retrieved documents. Default is False.
794 :returns: A list of the top_k documents most relevant to the query.
795 """
796 if not query:
797 raise ValueError("Query should be a non-empty string")
799 content_type_filter = {"field": "content", "operator": "!=", "value": None}
800 if filters:
801 if "operator" not in filters:
802 raise ValueError(
803 "Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details."
804 )
805 filters = {"operator": "AND", "conditions": [content_type_filter, filters]}
806 else:
807 filters = content_type_filter
809 all_documents = self.filter_documents(filters=filters)
810 if len(all_documents) == 0:
811 logger.info("No documents found for BM25 retrieval. Returning empty list.")
812 return []
814 # A tokenless corpus (every stored document has empty content) has no vocabulary and an
815 # average document length of zero, which would make all three BM25 algorithms divide by
816 # zero during scoring. Score every candidate as 0.0 instead; the non-positive-score
817 # handling below then keeps them for BM25Okapi and drops them for BM25L and BM25Plus.
818 if self._avg_doc_len == 0:
819 scored_documents = [(doc, 0.0) for doc in all_documents]
820 else:
821 scored_documents = self.bm25_algorithm_inst(query, all_documents)
823 results = sorted(scored_documents, key=lambda x: x[1], reverse=True)[:top_k]
825 # BM25Okapi can return meaningful negative values, so they should not be filtered out.
826 # It's the only algorithm supported by rank_bm25 at the time of writing (2024) that can return negative scores.
827 # see https://github.com/deepset-ai/haystack/pull/6889 for more context.
828 # BM25L and BM25Plus scores are 0 exactly when no query term occurs in the document. Filter on the raw score,
829 # because scaling maps 0 to 0.5.
830 drop_non_positive = self.bm25_algorithm != "BM25Okapi"
832 # Create documents with the BM25 score to return them
833 return_documents = []
834 for doc, score in results:
835 if drop_non_positive and score <= 0.0:
836 continue
838 if scale_score:
839 score = expit(score / BM25_SCALING_FACTOR)
841 doc_fields = doc.to_dict()
842 doc_fields["score"] = score
844 if not self.return_embedding and "embedding" in doc_fields:
845 doc_fields.pop("embedding")
847 return_document = Document.from_dict(doc_fields)
849 return_documents.append(return_document)
851 return return_documents
853 def embedding_retrieval(
854 self,
855 query_embedding: list[float],
856 filters: dict[str, Any] | None = None,
857 top_k: int = 10,
858 scale_score: bool = False,
859 return_embedding: bool | None = False,
860 ) -> list[Document]:
861 """
862 Retrieves documents that are most similar to the query embedding using a vector similarity metric.
864 :param query_embedding: Embedding of the query.
865 :param filters: A dictionary with filters to narrow down the search space.
866 :param top_k: The number of top documents to retrieve. Default is 10.
867 :param scale_score: Whether to scale the scores of the retrieved Documents. Default is False.
868 :param return_embedding: Whether to return the embedding of the retrieved Documents.
869 If not provided, the value of the `return_embedding` parameter set at component
870 initialization will be used. Default is False.
871 :returns: A list of the top_k documents most relevant to the query.
872 :raises ValueError: if filters have invalid syntax.
873 """
874 if len(query_embedding) == 0 or not isinstance(query_embedding[0], float):
875 raise ValueError("query_embedding should be a non-empty list of floats.")
877 if filters:
878 InMemoryDocumentStore._validate_filters(filters)
879 all_documents = [
880 doc
881 for doc in self.storage.values()
882 if document_matches_filter(
883 filters=filters, document=doc, strict_datetime_comparison=self.strict_datetime_comparison
884 )
885 ]
886 else:
887 all_documents = list(self.storage.values())
889 documents_with_embeddings = [doc for doc in all_documents if doc.embedding is not None]
890 if len(documents_with_embeddings) == 0:
891 logger.warning(
892 "No Documents found with embeddings. Returning empty list. "
893 "To generate embeddings, use a DocumentEmbedder."
894 )
895 return []
896 if len(documents_with_embeddings) < len(all_documents):
897 logger.info(
898 "Skipping some Documents that don't have an embedding. To generate embeddings, use a DocumentEmbedder."
899 )
901 scores = self._compute_query_embedding_similarity_scores(
902 embedding=query_embedding, documents=documents_with_embeddings, scale_score=scale_score
903 )
905 resolved_return_embedding = self.return_embedding if return_embedding is None else return_embedding
907 # create Documents with the similarity score for the top k results
908 top_documents = []
909 for doc, score in sorted(zip(documents_with_embeddings, scores, strict=True), key=lambda x: x[1], reverse=True)[
910 :top_k
911 ]:
912 doc_fields = doc.to_dict()
913 doc_fields["score"] = score
914 if resolved_return_embedding is False:
915 doc_fields["embedding"] = None
916 top_documents.append(Document.from_dict(doc_fields))
918 return top_documents
920 def _compute_query_embedding_similarity_scores(
921 self, embedding: list[float], documents: list[Document], scale_score: bool = False
922 ) -> list[float]:
923 """
924 Computes the similarity scores between the query embedding and the embeddings of the documents.
926 :param embedding: Embedding of the query.
927 :param documents: A list of Documents.
928 :param scale_score: Whether to scale the scores of the Documents. Default is False.
929 :returns: A list of scores.
930 """
932 query_embedding = np.array(embedding)
933 if query_embedding.ndim == 1:
934 query_embedding = np.expand_dims(a=query_embedding, axis=0)
936 try:
937 document_embeddings = np.array([doc.embedding for doc in documents])
938 except ValueError as e:
939 if "inhomogeneous shape" in str(e):
940 raise DocumentStoreError(
941 "The embedding size of all Documents should be the same. "
942 "Please make sure that the Documents have been embedded with the same model."
943 ) from e
944 raise e
945 if document_embeddings.ndim == 1:
946 document_embeddings = np.expand_dims(a=document_embeddings, axis=0)
948 if self.embedding_similarity_function == "cosine":
949 # cosine similarity is a normed dot product; guard against zero-norm vectors
950 # (e.g. a zero embedding) which would otherwise divide by zero and yield NaN scores.
951 query_norm = np.linalg.norm(x=query_embedding, axis=1, keepdims=True)
952 document_norms = np.linalg.norm(x=document_embeddings, axis=1, keepdims=True)
953 query_embedding /= np.where(query_norm == 0.0, 1.0, query_norm)
954 document_embeddings /= np.where(document_norms == 0.0, 1.0, document_norms)
956 try:
957 scores = np.dot(a=query_embedding, b=document_embeddings.T)[0].tolist()
958 except ValueError as e:
959 if "shapes" in str(e) and "not aligned" in str(e):
960 raise DocumentStoreError(
961 "The embedding size of the query should be the same as the embedding size of the Documents. "
962 "Please make sure that the query has been embedded with the same model as the Documents."
963 ) from e
964 raise e
966 if scale_score:
967 if self.embedding_similarity_function == "dot_product":
968 scores = [expit(float(score / DOT_PRODUCT_SCALING_FACTOR)) for score in scores]
969 elif self.embedding_similarity_function == "cosine":
970 scores = [(score + 1) / 2 for score in scores]
972 return scores
974 async def count_documents_async(self) -> int:
975 """
976 Returns the number of documents present in the DocumentStore.
977 """
978 return len(self.storage.keys())
980 async def filter_documents_async(self, filters: dict[str, Any] | None = None) -> list[Document]:
981 """
982 Returns the documents that match the filters provided.
984 :param filters: The filters to apply. For a detailed specification of the filters, refer to the
985 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering).
986 :returns: A list of Documents that match the given filters.
987 """
988 return await asyncio.get_running_loop().run_in_executor(
989 self.executor, lambda: self.filter_documents(filters=filters)
990 )
992 async def write_documents_async(
993 self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE
994 ) -> int:
995 """
996 Refer to the DocumentStore.write_documents() protocol documentation.
998 If `policy` is set to `DuplicatePolicy.NONE` defaults to `DuplicatePolicy.FAIL`.
999 """
1000 return await asyncio.get_running_loop().run_in_executor(
1001 self.executor, lambda: self.write_documents(documents=documents, policy=policy)
1002 )
1004 async def delete_documents_async(self, document_ids: list[str]) -> None:
1005 """
1006 Deletes all documents with matching document_ids from the DocumentStore.
1008 :param document_ids: The document_ids to delete.
1009 """
1010 await asyncio.get_running_loop().run_in_executor(
1011 self.executor, lambda: self.delete_documents(document_ids=document_ids)
1012 )
1014 async def update_by_filter_async(self, filters: dict[str, Any], meta: dict[str, Any]) -> int:
1015 """
1016 Updates the metadata of all documents that match the provided filters.
1018 :param filters: The filters to apply to select documents for updating.
1019 For filter syntax, see filter_documents.
1020 :param meta: The metadata fields to update. These will be merged with existing metadata.
1021 :returns: The number of documents updated.
1022 """
1023 return await asyncio.get_running_loop().run_in_executor(
1024 self.executor, lambda: self.update_by_filter(filters=filters, meta=meta)
1025 )
1027 async def count_documents_by_filter_async(self, filters: dict[str, Any]) -> int:
1028 """
1029 Returns the number of documents that match the provided filters.
1031 :param filters: The filters to apply.
1032 For a detailed specification of the filters, refer to the
1033 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering).
1034 :returns: The number of documents that match the filters.
1035 """
1036 return await asyncio.get_running_loop().run_in_executor(
1037 self.executor, lambda: self.count_documents_by_filter(filters=filters)
1038 )
1040 async def count_unique_metadata_by_filter_async(
1041 self, filters: dict[str, Any], metadata_fields: list[str]
1042 ) -> dict[str, int]:
1043 """
1044 Returns the number of unique values for each specified metadata field from documents matching the filters.
1046 JSON-serializable metadata values, including nested lists and dictionaries, are supported.
1048 :param filters: The filters to apply.
1049 For a detailed specification of the filters, refer to the
1050 [documentation](https://docs.haystack.deepset.ai/docs/metadata-filtering).
1051 :param metadata_fields: List of field names to count unique values for.
1052 Field names can include or omit the "meta." prefix.
1053 :returns: A dictionary mapping each metadata field name (without "meta." prefix)
1054 to the count of its unique values among the filtered documents.
1055 """
1056 return await asyncio.get_running_loop().run_in_executor(
1057 self.executor,
1058 lambda: self.count_unique_metadata_by_filter(filters=filters, metadata_fields=metadata_fields),
1059 )
1061 async def get_metadata_fields_info_async(self) -> dict[str, dict[str, str]]:
1062 """
1063 Returns information about the metadata fields present in the stored documents.
1065 Types are inferred from the stored values (keyword, int, float, boolean).
1067 :returns: A dictionary mapping each metadata field name to a dict with a "type" key.
1068 """
1069 return await asyncio.get_running_loop().run_in_executor(self.executor, self.get_metadata_fields_info)
1071 async def get_metadata_field_min_max_async(self, metadata_field: str) -> dict[str, Any]:
1072 """
1073 Returns the minimum and maximum values for the given metadata field across all documents.
1075 :param metadata_field: The metadata field name. Can include or omit the "meta." prefix.
1076 :returns: A dictionary with "min" and "max" keys. Returns `{"min": None, "max": None}`
1077 if the field is missing or has no values.
1078 """
1079 return await asyncio.get_running_loop().run_in_executor(
1080 self.executor, lambda: self.get_metadata_field_min_max(metadata_field=metadata_field)
1081 )
1083 async def get_metadata_field_unique_values_async(
1084 self,
1085 metadata_field: str,
1086 search_term: str | None = None,
1087 from_: int = 0,
1088 size: int = 10,
1089 filters: dict[str, Any] | None = None,
1090 ) -> tuple[list[Any], int]:
1091 """
1092 Returns unique values for a metadata field, optionally filtered by a search term, with pagination.
1094 JSON-serializable metadata values, including nested lists and dictionaries, are supported.
1096 :param metadata_field: The metadata field name. Can include or omit the "meta." prefix.
1097 :param search_term: Optional search term to filter values, matched as a case-insensitive substring
1098 against the metadata field's value.
1099 :param from_: The offset to start returning values from (for pagination).
1100 :param size: The maximum number of unique values to return.
1101 :param filters: Optional filters to restrict the documents considered.
1102 :returns: A tuple of (paginated list of unique values, total count of unique values).
1103 """
1104 return await asyncio.get_running_loop().run_in_executor(
1105 self.executor,
1106 lambda: self.get_metadata_field_unique_values(
1107 metadata_field=metadata_field, search_term=search_term, from_=from_, size=size, filters=filters
1108 ),
1109 )
1111 async def delete_all_documents_async(self) -> None:
1112 """
1113 Deletes all documents in the document store.
1114 """
1115 await asyncio.get_running_loop().run_in_executor(self.executor, self.delete_all_documents)
1117 async def bm25_retrieval_async(
1118 self, query: str, filters: dict[str, Any] | None = None, top_k: int = 10, scale_score: bool = False
1119 ) -> list[Document]:
1120 """
1121 Retrieves documents that are most relevant to the query using BM25 algorithm.
1123 :param query: The query string.
1124 :param filters: A dictionary with filters to narrow down the search space.
1125 :param top_k: The number of top documents to retrieve. Default is 10.
1126 :param scale_score: Whether to scale the scores of the retrieved documents. Default is False.
1127 :returns: A list of the top_k documents most relevant to the query.
1128 """
1129 return await asyncio.get_running_loop().run_in_executor(
1130 self.executor,
1131 lambda: self.bm25_retrieval(query=query, filters=filters, top_k=top_k, scale_score=scale_score),
1132 )
1134 async def embedding_retrieval_async(
1135 self,
1136 query_embedding: list[float],
1137 filters: dict[str, Any] | None = None,
1138 top_k: int = 10,
1139 scale_score: bool = False,
1140 return_embedding: bool = False,
1141 ) -> list[Document]:
1142 """
1143 Retrieves documents that are most similar to the query embedding using a vector similarity metric.
1145 :param query_embedding: Embedding of the query.
1146 :param filters: A dictionary with filters to narrow down the search space.
1147 :param top_k: The number of top documents to retrieve. Default is 10.
1148 :param scale_score: Whether to scale the scores of the retrieved Documents. Default is False.
1149 :param return_embedding: Whether to return the embedding of the retrieved Documents. Default is False.
1150 :returns: A list of the top_k documents most relevant to the query.
1151 """
1152 return await asyncio.get_running_loop().run_in_executor(
1153 self.executor,
1154 lambda: self.embedding_retrieval(
1155 query_embedding=query_embedding,
1156 filters=filters,
1157 top_k=top_k,
1158 scale_score=scale_score,
1159 return_embedding=return_embedding,
1160 ),
1161 )