feat: adding async version of InMemoryDocumentStore and associated retrievers (#8963)

* adding classes from experimental

* adding release notes

* adding tests

* merging all into a single class

* adding async retriever methods

* Update haystack/document_stores/in_memory/document_store.py

Co-authored-by: Stefano Fiorucci <[email protected]>

* adding missed tests

---------

Co-authored-by: Stefano Fiorucci <[email protected]>

Author: davidsbatista

haystack/components/retrievers/in_memory/bm25_retriever.py

@@ -161,3 +161,43 @@ def run(
 
         docs = self.document_store.bm25_retrieval(query=query, filters=filters, top_k=top_k, scale_score=scale_score)
         return {"documents": docs}
+
+    @component.output_types(documents=List[Document])
+    async def run_async(
+        self,
+        query: str,
+        filters: Optional[Dict[str, Any]] = None,
+        top_k: Optional[int] = None,
+        scale_score: Optional[bool] = None,
+    ):
+        """
+        Run the InMemoryBM25Retriever on the given input data.
+
+        :param query:
+            The query string for the Retriever.
+        :param filters:
+            A dictionary with filters to narrow down the search space when retrieving documents.
+        :param top_k:
+            The maximum number of documents to return.
+        :param scale_score:
+            When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant.
+            When `False`, uses raw similarity scores.
+        :returns:
+            The retrieved documents.
+
+        :raises ValueError:
+            If the specified DocumentStore is not found or is not a InMemoryDocumentStore instance.
+        """
+        if self.filter_policy == FilterPolicy.MERGE and filters:
+            filters = {**(self.filters or {}), **filters}
+        else:
+            filters = filters or self.filters
+        if top_k is None:
+            top_k = self.top_k
+        if scale_score is None:
+            scale_score = self.scale_score
+
+        docs = await self.document_store.bm25_retrieval_async(
+            query=query, filters=filters, top_k=top_k, scale_score=scale_score
+        )
+        return {"documents": docs}

haystack/components/retrievers/in_memory/embedding_retriever.py

@@ -192,3 +192,54 @@ def run(  # pylint: disable=too-many-positional-arguments
         )
 
         return {"documents": docs}
+
+    @component.output_types(documents=List[Document])
+    async def run_async(  # pylint: disable=too-many-positional-arguments
+        self,
+        query_embedding: List[float],
+        filters: Optional[Dict[str, Any]] = None,
+        top_k: Optional[int] = None,
+        scale_score: Optional[bool] = None,
+        return_embedding: Optional[bool] = None,
+    ):
+        """
+        Run the InMemoryEmbeddingRetriever on the given input data.
+
+        :param query_embedding:
+            Embedding of the query.
+        :param filters:
+            A dictionary with filters to narrow down the search space when retrieving documents.
+        :param top_k:
+            The maximum number of documents to return.
+        :param scale_score:
+            When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant.
+            When `False`, uses raw similarity scores.
+        :param return_embedding:
+            When `True`, returns the embedding of the retrieved documents.
+            When `False`, returns just the documents, without their embeddings.
+        :returns:
+            The retrieved documents.
+
+        :raises ValueError:
+            If the specified DocumentStore is not found or is not an InMemoryDocumentStore instance.
+        """
+        if self.filter_policy == FilterPolicy.MERGE and filters:
+            filters = {**(self.filters or {}), **filters}
+        else:
+            filters = filters or self.filters
+        if top_k is None:
+            top_k = self.top_k
+        if scale_score is None:
+            scale_score = self.scale_score
+        if return_embedding is None:
+            return_embedding = self.return_embedding
+
+        docs = await self.document_store.embedding_retrieval_async(
+            query_embedding=query_embedding,
+            filters=filters,
+            top_k=top_k,
+            scale_score=scale_score,
+            return_embedding=return_embedding,
+        )
+
+        return {"documents": docs}

haystack/document_stores/in_memory/document_store.py

@@ -2,11 +2,13 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import asyncio
 import json
 import math
 import re
 import uuid
 from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple
@@ -65,6 +67,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         bm25_parameters: Optional[Dict] = None,
         embedding_similarity_function: Literal["dot_product", "cosine"] = "dot_product",
         index: Optional[str] = None,
+        async_executor: Optional[ThreadPoolExecutor] = None,
     ):
         """
         Initializes the DocumentStore.
@@ -79,6 +82,9 @@ def __init__(  # pylint: disable=too-many-positional-arguments
             about your embedding model.
         :param index: A specific index to store the documents. If not specified, a random UUID is used.
             Using the same index allows you to store documents across multiple InMemoryDocumentStore instances.
+        :param async_executor:
+            Optional ThreadPoolExecutor to use for async calls. If not provided, a single-threaded
+            executor will be initialized and used.
         """
         self.bm25_tokenization_regex = bm25_tokenization_regex
         self.tokenizer = re.compile(bm25_tokenization_regex).findall
@@ -105,6 +111,12 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         if self.index not in _FREQ_VOCAB_FOR_IDF_STORAGES:
             _FREQ_VOCAB_FOR_IDF_STORAGES[self.index] = Counter()
 
+        self.executor = (
+            ThreadPoolExecutor(thread_name_prefix=f"async-inmemory-docstore-executor-{id(self)}", max_workers=1)
+            if async_executor is None
+            else async_executor
+        )
+
     @property
     def storage(self) -> Dict[str, Document]:
         """
@@ -620,3 +632,91 @@ def _compute_query_embedding_similarity_scores(
                 scores = [(score + 1) / 2 for score in scores]
 
         return scores
+
+    async def count_documents_async(self) -> int:
+        """
+        Returns the number of how many documents are present in the DocumentStore.
+        """
+        return len(self.storage.keys())
+
+    async def filter_documents_async(self, filters: Optional[Dict[str, Any]] = None) -> List[Document]:
+        """
+        Returns the documents that match the filters provided.
+
+        For a detailed specification of the filters, refer to the DocumentStore.filter_documents() protocol
+        documentation.
+
+        :param filters: The filters to apply to the document list.
+        :returns: A list of Documents that match the given filters.
+        """
+        return await asyncio.get_event_loop().run_in_executor(
+            self.executor, lambda: self.filter_documents(filters=filters)
+        )
+
+    async def write_documents_async(
+        self, documents: List[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE
+    ) -> int:
+        """
+        Refer to the DocumentStore.write_documents() protocol documentation.
+
+        If `policy` is set to `DuplicatePolicy.NONE` defaults to `DuplicatePolicy.FAIL`.
+        """
+        return await asyncio.get_event_loop().run_in_executor(
+            self.executor, lambda: self.write_documents(documents=documents, policy=policy)
+        )
+
+    async def delete_documents_async(self, document_ids: List[str]) -> None:
+        """
+        Deletes all documents with matching document_ids from the DocumentStore.
+
+        :param document_ids: The object_ids to delete.
+        """
+        await asyncio.get_event_loop().run_in_executor(
+            self.executor, lambda: self.delete_documents(document_ids=document_ids)
+        )
+
+    async def bm25_retrieval_async(
+        self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: int = 10, scale_score: bool = False
+    ) -> List[Document]:
+        """
+        Retrieves documents that are most relevant to the query using BM25 algorithm.
+
+        :param query: The query string.
+        :param filters: A dictionary with filters to narrow down the search space.
+        :param top_k: The number of top documents to retrieve. Default is 10.
+        :param scale_score: Whether to scale the scores of the retrieved documents. Default is False.
+        :returns: A list of the top_k documents most relevant to the query.
+        """
+        return await asyncio.get_event_loop().run_in_executor(
+            self.executor,
+            lambda: self.bm25_retrieval(query=query, filters=filters, top_k=top_k, scale_score=scale_score),
+        )
+
+    async def embedding_retrieval_async(  # pylint: disable=too-many-positional-arguments
+        self,
+        query_embedding: List[float],
+        filters: Optional[Dict[str, Any]] = None,
+        top_k: int = 10,
+        scale_score: bool = False,
+        return_embedding: bool = False,
+    ) -> List[Document]:
+        """
+        Retrieves documents that are most similar to the query embedding using a vector similarity metric.
+
+        :param query_embedding: Embedding of the query.
+        :param filters: A dictionary with filters to narrow down the search space.
+        :param top_k: The number of top documents to retrieve. Default is 10.
+        :param scale_score: Whether to scale the scores of the retrieved Documents. Default is False.
+        :param return_embedding: Whether to return the embedding of the retrieved Documents. Default is False.
+        :returns: A list of the top_k documents most relevant to the query.
+        """
+        return await asyncio.get_event_loop().run_in_executor(
+            self.executor,
+            lambda: self.embedding_retrieval(
+                query_embedding=query_embedding,
+                filters=filters,
+                top_k=top_k,
+                scale_score=scale_score,
+                return_embedding=return_embedding,
+            ),
+        )

releasenotes/notes/adding-async-inmemory-doc-store-3e5d108f921f2b70.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    The `InMemoryDocumentStore` and the associated `InMemoryBM25Retriever` and `InMemoryEmbeddingRetriever` retrievers now support async mode.