deepset-ai--haystack/test/document_stores/test_in_memory.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import asyncio
import gc
import logging
import math
import tempfile
from typing import Literal, cast
from unittest.mock import patch

import pytest

from haystack import Document
from haystack.dataclasses import ByteStream, SparseEmbedding
from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.document_stores.in_memory import document_store as in_memory_module
from haystack.testing.document_store import (
    CountDocumentsByFilterTest,
    CountUniqueMetadataByFilterTest,
    DocumentStoreBaseExtendedTests,
    DocumentStoreBaseTests,
    FilterableDocsFixtureMixin,
    GetMetadataFieldMinMaxTest,
    GetMetadataFieldsInfoTest,
    GetMetadataFieldUniqueValuesTest,
)
from haystack.testing.document_store_async import (
    CountDocumentsAsyncTest,
    CountDocumentsByFilterAsyncTest,
    CountUniqueMetadataByFilterAsyncTest,
    DeleteAllAsyncTest,
    DeleteDocumentsAsyncTest,
    FilterDocumentsAsyncTest,
    GetMetadataFieldMinMaxAsyncTest,
    GetMetadataFieldsInfoAsyncTest,
    GetMetadataFieldUniqueValuesAsyncTest,
    UpdateByFilterAsyncTest,
    WriteDocumentsAsyncTest,
)


class TestMemoryDocumentStore(
    DocumentStoreBaseExtendedTests,
    UpdateByFilterAsyncTest,
    CountDocumentsByFilterAsyncTest,
    CountDocumentsAsyncTest,
    WriteDocumentsAsyncTest,
    DeleteAllAsyncTest,
    DeleteDocumentsAsyncTest,
    CountDocumentsByFilterTest,
    CountUniqueMetadataByFilterAsyncTest,
    CountUniqueMetadataByFilterTest,
    FilterDocumentsAsyncTest,
    FilterableDocsFixtureMixin,
    GetMetadataFieldMinMaxTest,
    GetMetadataFieldUniqueValuesTest,
    GetMetadataFieldsInfoTest,
    GetMetadataFieldsInfoAsyncTest,
    GetMetadataFieldMinMaxAsyncTest,
    GetMetadataFieldUniqueValuesAsyncTest,
):
    """
    Test InMemoryDocumentStore's specific features
    """

    @pytest.fixture
    def tmp_dir(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            yield tmp_dir

    @pytest.fixture
    def document_store(self):
        store = InMemoryDocumentStore(bm25_algorithm="BM25L")
        yield store
        store.shutdown()

    @pytest.fixture
    def cosine_document_store(self):
        store = InMemoryDocumentStore(embedding_similarity_function="cosine")
        yield store
        store.shutdown()

    def test_to_dict(self, in_memory_doc_store):
        data = in_memory_doc_store.to_dict()
        assert data == {
            "type": "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore",
            "init_parameters": {
                "bm25_tokenization_regex": r"(?u)\b\w+\b",
                "bm25_algorithm": "BM25L",
                "bm25_parameters": {},
                "embedding_similarity_function": "dot_product",
                "index": in_memory_doc_store.index,
                "shared": True,
                "return_embedding": True,
            },
        }

    def test_to_dict_with_custom_init_parameters(self):
        store = InMemoryDocumentStore(
            bm25_tokenization_regex="custom_regex",
            bm25_algorithm="BM25Plus",
            bm25_parameters={"key": "value"},
            embedding_similarity_function="cosine",
            index="my_cool_index",
            return_embedding=True,
        )
        data = store.to_dict()
        assert data == {
            "type": "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore",
            "init_parameters": {
                "bm25_tokenization_regex": "custom_regex",
                "bm25_algorithm": "BM25Plus",
                "bm25_parameters": {"key": "value"},
                "embedding_similarity_function": "cosine",
                "index": "my_cool_index",
                "shared": True,
                "return_embedding": True,
            },
        }

    @patch("haystack.document_stores.in_memory.document_store.re")
    def test_from_dict(self, mock_regex):
        data = {
            "type": "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore",
            "init_parameters": {
                "bm25_tokenization_regex": "custom_regex",
                "bm25_algorithm": "BM25Plus",
                "bm25_parameters": {"key": "value"},
                "index": "my_cool_index",
            },
        }
        store = InMemoryDocumentStore.from_dict(data)
        mock_regex.compile.assert_called_with("custom_regex")
        assert store.tokenizer is not None
        assert store.bm25_algorithm == "BM25Plus"
        assert store.bm25_parameters == {"key": "value"}
        assert store.index == "my_cool_index"

    def test_save_to_disk_and_load_from_disk(self, in_memory_doc_store: InMemoryDocumentStore, tmp_dir: str) -> None:
        docs = [Document(content="Hello world"), Document(content="Haystack supports multiple languages")]
        in_memory_doc_store.write_documents(docs)
        tmp_dir = tmp_dir + "/in_memory_doc_store.json"
        in_memory_doc_store.save_to_disk(tmp_dir)
        document_store_loaded = InMemoryDocumentStore.load_from_disk(tmp_dir)

        assert document_store_loaded.count_documents() == 2
        assert list(document_store_loaded.storage.values()) == docs
        assert document_store_loaded.to_dict() == in_memory_doc_store.to_dict()

    def test_save_to_disk_and_load_from_disk_with_blob_and_sparse_embedding(
        self, in_memory_doc_store: InMemoryDocumentStore, tmp_dir: str
    ) -> None:
        doc = Document(
            content="document with binary data",
            blob=ByteStream(data=b"binary data", mime_type="image/png"),
            sparse_embedding=SparseEmbedding(indices=[0, 5], values=[0.1, 0.9]),
        )
        in_memory_doc_store.write_documents([doc])
        save_path = tmp_dir + "/in_memory_doc_store.json"
        in_memory_doc_store.save_to_disk(save_path)
        document_store_loaded = InMemoryDocumentStore.load_from_disk(save_path)

        loaded_doc = document_store_loaded.filter_documents()[0]
        assert isinstance(loaded_doc.blob, ByteStream)
        assert isinstance(loaded_doc.sparse_embedding, SparseEmbedding)
        assert loaded_doc == doc
        # The loaded store must be savable again
        document_store_loaded.save_to_disk(save_path)

    def test_invalid_bm25_algorithm(self):
        with pytest.raises(ValueError, match="BM25 algorithm 'invalid' is not supported"):
            InMemoryDocumentStore(bm25_algorithm="invalid")  # type: ignore[arg-type]

    def test_write_documents(self, document_store):
        docs = [Document(id="1")]
        assert document_store.write_documents(docs) == 1
        with pytest.raises(DuplicateDocumentError):
            document_store.write_documents(docs)

    def test_bm25_retrieval(self, document_store: InMemoryDocumentStore) -> None:
        # Tests if the bm25_retrieval method returns the correct document based on the input query.
        docs = [Document(content="Hello world"), Document(content="Haystack supports multiple languages")]
        document_store.write_documents(docs)
        results = document_store.bm25_retrieval(query="What languages?", top_k=1)
        assert len(results) == 1
        assert results[0].content == "Haystack supports multiple languages"

    def test_bm25_retrieval_with_empty_document_store(
        self, document_store: InMemoryDocumentStore, caplog: pytest.LogCaptureFixture
    ) -> None:
        caplog.set_level(logging.INFO)
        # Tests if the bm25_retrieval method correctly returns an empty list when there are no documents in the
        # DocumentStore.
        results = document_store.bm25_retrieval(query="How to test this?", top_k=2)
        assert len(results) == 0
        assert "No documents found for BM25 retrieval. Returning empty list." in caplog.text

    @pytest.mark.parametrize("bm25_algorithm", ["BM25Okapi", "BM25L", "BM25Plus"])
    def test_bm25_retrieval_with_tokenless_corpus(
        self, bm25_algorithm: Literal["BM25Okapi", "BM25L", "BM25Plus"]
    ) -> None:
        # Regression test for #11598: a corpus where every document has empty (but not None)
        # content must not raise ZeroDivisionError at query time.
        store = InMemoryDocumentStore(bm25_algorithm=bm25_algorithm)
        store.write_documents([Document(content="", meta={"i": 1}), Document(content="", meta={"i": 2})])
        results = store.bm25_retrieval(query="anything")
        if bm25_algorithm == "BM25Okapi":
            # Unscaled BM25Okapi keeps non-positive scores, so documents are returned with score 0.0.
            assert len(results) == 2
            assert all(doc.score == 0.0 for doc in results)
        else:
            # BM25L / BM25Plus filter out non-positive scores.
            assert results == []

    def test_bm25_retrieval_empty_query(self, document_store: InMemoryDocumentStore) -> None:
        # Tests if the bm25_retrieval method returns a document when the query is an empty string.
        docs = [Document(content="Hello world"), Document(content="Haystack supports multiple languages")]
        document_store.write_documents(docs)
        with pytest.raises(ValueError, match="Query should be a non-empty string"):
            document_store.bm25_retrieval(query="", top_k=1)

    def test_bm25_retrieval_with_different_top_k(self, document_store: InMemoryDocumentStore) -> None:
        # Tests if the bm25_retrieval method correctly changes the number of returned documents
        # based on the top_k parameter.
        docs = [
            Document(content="Hello world"),
            Document(content="Haystack supports multiple languages"),
            Document(content="Python is a popular programming language"),
        ]
        document_store.write_documents(docs)

        # top_k = 2
        results = document_store.bm25_retrieval(query="language", top_k=2)
        assert len(results) == 2

        # top_k = 3
        results = document_store.bm25_retrieval(query="languages", top_k=3)
        assert len(results) == 3

    def test_bm25_plus_retrieval(self):
        doc_store = InMemoryDocumentStore(bm25_algorithm="BM25Plus")
        docs = [
            Document(content="Hello world"),
            Document(content="Haystack supports multiple languages"),
            Document(content="Python is a popular programming language"),
        ]
        doc_store.write_documents(docs)

        results = doc_store.bm25_retrieval(query="language", top_k=1)
        assert len(results) == 1
        assert results[0].content == "Python is a popular programming language"

    def test_bm25_retrieval_with_two_queries(self, document_store: InMemoryDocumentStore) -> None:
        # Tests if the bm25_retrieval method returns different documents for different queries.
        docs = [
            Document(content="Javascript is a popular programming language"),
            Document(content="Java is a popular programming language"),
            Document(content="Python is a popular programming language"),
            Document(content="Ruby is a popular programming language"),
            Document(content="PHP is a popular programming language"),
        ]
        document_store.write_documents(docs)

        results = document_store.bm25_retrieval(query="Java", top_k=1)
        assert results[0].content == "Java is a popular programming language"

        results = document_store.bm25_retrieval(query="Python", top_k=1)
        assert results[0].content == "Python is a popular programming language"

    # Test a query, add a new document and make sure results are appropriately updated

    def test_bm25_retrieval_with_updated_docs(self, document_store: InMemoryDocumentStore) -> None:
        # Tests if the bm25_retrieval method correctly updates the retrieved documents when new
        # documents are added to the DocumentStore.
        docs = [Document(content="Hello world")]
        document_store.write_documents(docs)

        results = document_store.bm25_retrieval(query="Python", top_k=1)
        assert len(results) == 0

        document_store.write_documents([Document(content="Python is a popular programming language")])
        results = document_store.bm25_retrieval(query="Python", top_k=1)
        assert len(results) == 1
        assert results[0].content == "Python is a popular programming language"

        document_store.write_documents([Document(content="Java is a popular programming language")])
        results = document_store.bm25_retrieval(query="Python", top_k=1)
        assert len(results) == 1
        assert results[0].content == "Python is a popular programming language"

    def test_bm25_retrieval_with_scale_score(self, document_store: InMemoryDocumentStore) -> None:
        docs = [Document(content="Python programming"), Document(content="Java programming")]
        document_store.write_documents(docs)

        results1 = document_store.bm25_retrieval(query="Python", top_k=1, scale_score=True)
        # Confirm that score is scaled between 0 and 1
        assert results1[0].score is not None
        assert 0.0 <= results1[0].score <= 1.0

        # Same query, different scale, scores differ when not scaled
        results = document_store.bm25_retrieval(query="Python", top_k=1, scale_score=False)
        assert results[0].score != results1[0].score

    def test_bm25_retrieval_with_non_scaled_BM25Okapi(self):
        # Highly repetitive documents make BM25Okapi return negative scores, which should not be filtered if the
        # scores are not scaled
        docs = [
            Document(
                content="""Use pip to install a basic version of Haystack's latest release: pip install
                farm-haystack. All the core Haystack components live in the haystack repo. But there's also the
                haystack-extras repo which contains components that are not as widely used, and you need to
                install them separately."""
            ),
            Document(
                content="""Use pip to install a basic version of Haystack's latest release: pip install
                farm-haystack[inference]. All the core Haystack components live in the haystack repo. But there's
                also the haystack-extras repo which contains components that are not as widely used, and you need
                to install them separately."""
            ),
            Document(
                content="""Use pip to install only the Haystack 2.0 code: pip install haystack-ai. The haystack-ai
                package is built on the main branch which is an unstable beta version, but it's useful if you want
                to try the new features as soon as they are merged."""
            ),
        ]
        document_store = InMemoryDocumentStore(bm25_algorithm="BM25Okapi")
        document_store.write_documents(docs)

        results1 = document_store.bm25_retrieval(query="Haystack installation", top_k=10, scale_score=False)
        assert len(results1) == 3
        assert all(res.score < 0.0 for res in results1 if res.score)

        results2 = document_store.bm25_retrieval(query="Haystack installation", top_k=10, scale_score=True)
        assert len(results2) == 3
        assert all(0.0 <= res.score <= 1.0 for res in results2 if res.score)

    def test_bm25_retrieval_default_filter(self, document_store: InMemoryDocumentStore) -> None:
        docs = [Document(), Document(content="Gardening"), Document(content="Bird watching")]
        document_store.write_documents(docs)
        results = document_store.bm25_retrieval(query="doesn't matter, top_k is 10", top_k=10)
        assert len(results) == 0

    def test_embedding_retrieval_return_embedding_false_on_store(self):
        # Initialize InMemoryDocumentStore with return_embedding=False
        docstore = InMemoryDocumentStore(embedding_similarity_function="cosine", return_embedding=False)
        docs = [
            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
            Document(content="Haystack supports multiple languages", embedding=[1.0, 1.0, 1.0, 1.0]),
        ]
        docstore.write_documents(docs)

        # embedding_retrieval should not return embeddings in the documents
        results = docstore.embedding_retrieval(query_embedding=[0.1, 0.1, 0.1, 0.1], top_k=2)
        assert all(doc.embedding is None for doc in results)

        # bm25_retrieval should also not return embeddings
        bm25_results = docstore.bm25_retrieval(query="languages", top_k=2)
        assert all(doc.embedding is None for doc in bm25_results)

        # filter_documents should not return embeddings
        filtered_docs = docstore.filter_documents()
        assert all(doc.embedding is None for doc in filtered_docs)

    def test_embedding_retrieval_override_return_embedding(self):
        docstore = InMemoryDocumentStore(embedding_similarity_function="cosine", return_embedding=False)
        docs = [
            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
            Document(content="Haystack supports multiple languages", embedding=[1.0, 1.0, 1.0, 1.0]),
        ]
        docstore.write_documents(docs)

        # Overriding return_embedding to True should return embeddings
        # Query for the embedding that matches both documents by cosine similarity
        results_with_embedding = docstore.embedding_retrieval(
            query_embedding=[0.1, 0.2, 0.3, 0.4], top_k=2, return_embedding=True
        )

        # Assert that the retrieved documents have the expected embeddings
        assert len(results_with_embedding) == 2

        assert results_with_embedding[0].embedding in ([1.0, 1.0, 1.0, 1.0], [0.1, 0.2, 0.3, 0.4])
        assert results_with_embedding[1].embedding in ([1.0, 1.0, 1.0, 1.0], [0.1, 0.2, 0.3, 0.4])

    def test_embedding_retrieval(self):
        docstore = InMemoryDocumentStore(embedding_similarity_function="cosine")
        # Tests if the embedding retrieval method returns the correct document based on the input query embedding.
        docs = [
            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
            Document(content="Haystack supports multiple languages", embedding=[1.0, 1.0, 1.0, 1.0]),
        ]
        docstore.write_documents(docs)
        results = docstore.embedding_retrieval(
            query_embedding=[0.1, 0.1, 0.1, 0.1], top_k=1, filters={}, scale_score=False
        )
        assert len(results) == 1
        assert results[0].content == "Haystack supports multiple languages"

    def test_embedding_retrieval_with_zero_vector_does_not_produce_nan(self):
        # A zero embedding has no direction; normalizing it must not divide by zero and
        # produce NaN cosine scores, which would silently corrupt ranking.
        docstore = InMemoryDocumentStore(embedding_similarity_function="cosine")
        docstore.write_documents(
            [Document(content="zero", embedding=[0.0, 0.0, 0.0]), Document(content="normal", embedding=[1.0, 0.0, 0.0])]
        )
        results = docstore.embedding_retrieval(query_embedding=[1.0, 0.0, 0.0], scale_score=False)
        scores = {doc.content: doc.score for doc in results}
        assert all(score is not None and not math.isnan(score) for score in scores.values())
        assert scores["zero"] == 0.0

    def test_embedding_retrieval_invalid_query(self, in_memory_doc_store):
        with pytest.raises(ValueError, match="query_embedding should be a non-empty list of floats"):
            in_memory_doc_store.embedding_retrieval(query_embedding=[])
        with pytest.raises(ValueError, match="query_embedding should be a non-empty list of floats"):
            in_memory_doc_store.embedding_retrieval(query_embedding=["invalid", "list", "of", "strings"])

    def test_embedding_retrieval_no_embeddings(self, in_memory_doc_store, caplog):
        caplog.set_level(logging.WARNING)
        docs = [Document(content="Hello world"), Document(content="Haystack supports multiple languages")]
        in_memory_doc_store.write_documents(docs)
        results = in_memory_doc_store.embedding_retrieval(query_embedding=[0.1, 0.1, 0.1, 0.1])
        assert len(results) == 0
        assert "No Documents found with embeddings. Returning empty list." in caplog.text

    def test_embedding_retrieval_some_documents_wo_embeddings(self, in_memory_doc_store, caplog):
        caplog.set_level(logging.INFO)
        docs = [
            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
            Document(content="Haystack supports multiple languages"),
        ]
        in_memory_doc_store.write_documents(docs)
        in_memory_doc_store.embedding_retrieval(query_embedding=[0.1, 0.1, 0.1, 0.1])
        assert "Skipping some Documents that don't have an embedding." in caplog.text

    def test_embedding_retrieval_documents_different_embedding_sizes(self, in_memory_doc_store):
        docs = [
            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
            Document(content="Haystack supports multiple languages", embedding=[1.0, 1.0]),
        ]
        in_memory_doc_store.write_documents(docs)

        with pytest.raises(DocumentStoreError, match="The embedding size of all Documents should be the same."):
            in_memory_doc_store.embedding_retrieval(query_embedding=[0.1, 0.1, 0.1, 0.1])

    def test_embedding_retrieval_query_documents_different_embedding_sizes(self, in_memory_doc_store):
        docs = [Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4])]
        in_memory_doc_store.write_documents(docs)

        with pytest.raises(
            DocumentStoreError,
            match="The embedding size of the query should be the same as the embedding size of the Documents.",
        ):
            in_memory_doc_store.embedding_retrieval(query_embedding=[0.1, 0.1])

    def test_embedding_retrieval_with_different_top_k(self, in_memory_doc_store):
        docs = [
            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
            Document(content="Haystack supports multiple languages", embedding=[1.0, 1.0, 1.0, 1.0]),
            Document(content="Python is a popular programming language", embedding=[0.5, 0.5, 0.5, 0.5]),
        ]
        in_memory_doc_store.write_documents(docs)

        results = in_memory_doc_store.embedding_retrieval(query_embedding=[0.1, 0.1, 0.1, 0.1], top_k=2)
        assert len(results) == 2

        results = in_memory_doc_store.embedding_retrieval(query_embedding=[0.1, 0.1, 0.1, 0.1], top_k=3)
        assert len(results) == 3

    def test_embedding_retrieval_with_scale_score(self, in_memory_doc_store):
        docs = [
            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
            Document(content="Haystack supports multiple languages", embedding=[1.0, 1.0, 1.0, 1.0]),
            Document(content="Python is a popular programming language", embedding=[0.5, 0.5, 0.5, 0.5]),
        ]
        in_memory_doc_store.write_documents(docs)

        results1 = in_memory_doc_store.embedding_retrieval(
            query_embedding=[0.1, 0.1, 0.1, 0.1], top_k=1, scale_score=True
        )
        # Confirm that score is scaled between 0 and 1
        assert results1[0].score is not None
        assert 0.0 <= results1[0].score <= 1.0

        # Same query, different scale, scores differ when not scaled
        results = in_memory_doc_store.embedding_retrieval(
            query_embedding=[0.1, 0.1, 0.1, 0.1], top_k=1, scale_score=False
        )
        assert results[0].score != results1[0].score

    def test_embedding_retrieval_return_embedding(self):
        docstore = InMemoryDocumentStore(embedding_similarity_function="cosine")
        docs = [
            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
            Document(content="Haystack supports multiple languages", embedding=[1.0, 1.0, 1.0, 1.0]),
        ]
        docstore.write_documents(docs)

        results = docstore.embedding_retrieval(query_embedding=[0.1, 0.1, 0.1, 0.1], top_k=1, return_embedding=False)
        assert results[0].embedding is None

        results = docstore.embedding_retrieval(query_embedding=[0.1, 0.1, 0.1, 0.1], top_k=1, return_embedding=True)
        assert results[0].embedding == [1.0, 1.0, 1.0, 1.0]

    def test_compute_cosine_similarity_scores(self):
        docstore = InMemoryDocumentStore(embedding_similarity_function="cosine")
        docs = [
            Document(content="Document 1", embedding=[1.0, 0.0, 0.0, 0.0]),
            Document(content="Document 2", embedding=[1.0, 1.0, 1.0, 1.0]),
        ]

        scores = docstore._compute_query_embedding_similarity_scores(
            embedding=[0.1, 0.1, 0.1, 0.1], documents=docs, scale_score=False
        )
        assert scores == [0.5, 1.0]

    def test_compute_dot_product_similarity_scores(self):
        docstore = InMemoryDocumentStore(embedding_similarity_function="dot_product")
        docs = [
            Document(content="Document 1", embedding=[1.0, 0.0, 0.0, 0.0]),
            Document(content="Document 2", embedding=[1.0, 1.0, 1.0, 1.0]),
        ]

        scores = docstore._compute_query_embedding_similarity_scores(
            embedding=[0.1, 0.1, 0.1, 0.1], documents=docs, scale_score=False
        )
        assert scores == [0.1, 0.4]

    def test_multiple_document_stores_using_same_index(self):
        index = "test_multiple_document_stores_using_same_index"
        document_store_1 = InMemoryDocumentStore(index=index)
        document_store_2 = InMemoryDocumentStore(index=index)

        assert document_store_1.count_documents() == document_store_2.count_documents() == 0

        doc_1 = Document(content="Hello world")
        document_store_1.write_documents([doc_1])
        assert document_store_1.count_documents() == document_store_2.count_documents() == 1

        assert document_store_1.filter_documents() == document_store_2.filter_documents() == [doc_1]

        doc_2 = Document(content="Hello another world")
        document_store_2.write_documents([doc_2])
        assert document_store_1.count_documents() == document_store_2.count_documents() == 2

        assert document_store_1.filter_documents() == document_store_2.filter_documents() == [doc_1, doc_2]

        document_store_1.delete_documents([doc_2.id])
        assert document_store_1.count_documents() == document_store_2.count_documents() == 1

        document_store_2.delete_documents([doc_1.id])
        assert document_store_1.count_documents() == document_store_2.count_documents() == 0

    # Test async/await methods and concurrency

    @pytest.mark.asyncio
    async def test_write_documents_async(self, document_store: InMemoryDocumentStore) -> None:
        docs = [Document(id="1")]
        assert await document_store.write_documents_async(docs) == 1
        with pytest.raises(DuplicateDocumentError):
            await document_store.write_documents_async(docs)

    @pytest.mark.asyncio
    async def test_filter_documents(self, document_store: InMemoryDocumentStore) -> None:
        filterable_docs = [Document(content="1", meta={"number": -10}), Document(content="2", meta={"number": 100})]
        await document_store.write_documents_async(filterable_docs)
        result = await document_store.filter_documents_async(
            filters={"field": "meta.number", "operator": "==", "value": 100}
        )
        DocumentStoreBaseTests().assert_documents_are_equal(
            result, [d for d in filterable_docs if d.meta.get("number") == 100]
        )

    @pytest.mark.asyncio
    async def test_bm25_retrieval_async(self, document_store: InMemoryDocumentStore) -> None:
        # Tests if the bm25_retrieval method returns the correct document based on the input query.
        docs = [Document(content="Hello world"), Document(content="Haystack supports multiple languages")]
        await document_store.write_documents_async(docs)
        results = await document_store.bm25_retrieval_async(query="What languages?", top_k=1)
        assert len(results) == 1
        assert results[0].content == "Haystack supports multiple languages"

    @pytest.mark.asyncio
    async def test_embedding_retrieval_async(self, cosine_document_store):
        # Tests if the embedding retrieval method returns the correct document based on the input query embedding.
        docs = [
            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
            Document(content="Haystack supports multiple languages", embedding=[1.0, 1.0, 1.0, 1.0]),
        ]
        await cosine_document_store.write_documents_async(docs)
        results = await cosine_document_store.embedding_retrieval_async(
            query_embedding=[0.1, 0.1, 0.1, 0.1], top_k=1, filters={}, scale_score=False
        )
        assert len(results) == 1
        assert results[0].content == "Haystack supports multiple languages"

    @pytest.mark.asyncio
    async def test_concurrent_bm25_retrievals(self, document_store: InMemoryDocumentStore) -> None:
        # Test multiple concurrent BM25 retrievals
        docs = [
            Document(content="Python is a popular programming language"),
            Document(content="Java is a popular programming language"),
            Document(content="JavaScript is a popular programming language"),
            Document(content="Ruby is a popular programming language"),
        ]
        await document_store.write_documents_async(docs)

        # Create multiple concurrent retrievals
        queries = ["Python", "Java", "JavaScript", "Ruby"]
        tasks = [document_store.bm25_retrieval_async(query=query, top_k=1) for query in queries]
        results = await asyncio.gather(*tasks)

        # Verify each result matches the expected content
        for query, result in zip(queries, results, strict=True):
            assert len(result) == 1
            assert result[0].content == f"{query} is a popular programming language"

    @pytest.mark.asyncio
    async def test_concurrent_embedding_retrievals(self, cosine_document_store):
        # Test multiple concurrent embedding retrievals
        docs = [
            Document(content="Python programming", embedding=[1.0, 0.0, 0.0, 0.0]),
            Document(content="Java programming", embedding=[0.0, 1.0, 0.0, 0.0]),
            Document(content="JavaScript programming", embedding=[0.0, 0.0, 1.0, 0.0]),
            Document(content="Ruby programming", embedding=[0.0, 0.0, 0.0, 1.0]),
        ]
        await cosine_document_store.write_documents_async(docs)

        # Create multiple concurrent retrievals with different query embeddings
        query_embeddings = [
            [1.0, 0.0, 0.0, 0.0],  # Should match Python
            [0.0, 1.0, 0.0, 0.0],  # Should match Java
            [0.0, 0.0, 1.0, 0.0],  # Should match JavaScript
            [0.0, 0.0, 0.0, 1.0],  # Should match Ruby
        ]
        tasks = [
            cosine_document_store.embedding_retrieval_async(query_embedding=emb, top_k=1) for emb in query_embeddings
        ]
        results = await asyncio.gather(*tasks)

        # Verify each result matches the expected content
        expected_contents = ["Python programming", "Java programming", "JavaScript programming", "Ruby programming"]
        for result, expected in zip(results, expected_contents, strict=True):
            assert len(result) == 1
            assert result[0].content == expected

    @pytest.mark.asyncio
    async def test_mixed_concurrent_operations(self, document_store: InMemoryDocumentStore) -> None:
        # Test a mix of concurrent operations including writes and retrievals
        docs = [
            Document(content="First document"),
            Document(content="Second document"),
            Document(content="Third document"),
        ]
        await document_store.write_documents_async(docs)

        # Create a mix of concurrent operations
        tasks = [
            document_store.bm25_retrieval_async(query="First", top_k=1),
            document_store.write_documents_async([Document(content="Fourth document")]),
            document_store.bm25_retrieval_async(query="Fourth", top_k=1),
            document_store.filter_documents_async(),
        ]
        results = cast(tuple[list[Document], int, list[Document], list[Document]], await asyncio.gather(*tasks))

        # Verify results
        assert len(results[0]) == 1  # First retrieval
        assert results[1] == 1  # Write operation
        assert len(results[2]) == 1  # Fourth retrieval
        assert len(results[3]) == 4  # Filter operation

    @pytest.mark.asyncio
    async def test_concurrent_operations_with_errors(self, document_store: InMemoryDocumentStore) -> None:
        # Test concurrent operations where some might fail
        docs = [Document(content="Test document")]
        await document_store.write_documents_async(docs)

        # Create tasks including some that should fail
        tasks = [
            document_store.bm25_retrieval_async(query="Test", top_k=1),  # Should succeed
            document_store.bm25_retrieval_async(query="", top_k=1),  # Should fail
            document_store.embedding_retrieval_async(query_embedding=[], top_k=1),  # Should fail
        ]

        # Gather results and expect some to raise exceptions
        with pytest.raises(ValueError):
            await asyncio.gather(*tasks)

    @pytest.mark.asyncio
    async def test_concurrent_operations_with_large_dataset(self, document_store: InMemoryDocumentStore) -> None:
        # Test concurrent operations with a larger dataset
        # Create 100 documents with different content
        docs = [Document(content=f"Document {i} content") for i in range(100)]
        await document_store.write_documents_async(docs)

        # Create multiple concurrent retrievals
        queries = [f"Document {i}" for i in range(0, 100, 10)]  # Query every 10th document
        tasks = [document_store.bm25_retrieval_async(query=query, top_k=1) for query in queries]
        results = await asyncio.gather(*tasks)

        # Verify results
        for i, result in enumerate(results):
            assert len(result) == 1
            assert result[0].content == f"Document {i * 10} content"

    def test_executor_shutdown(self):
        doc_store = InMemoryDocumentStore()
        executor = doc_store.executor
        with patch.object(executor, "shutdown", wraps=executor.shutdown) as mock_shutdown:
            del doc_store
            gc.collect()
            mock_shutdown.assert_called_once_with(wait=True)

    def test_bm25_tokenization_includes_single_char_tokens(self, in_memory_doc_store):
        tokens = in_memory_doc_store._tokenize_bm25("Luna is a dog")
        assert tokens == ["luna", "is", "a", "dog"]

    def test_bm25_retrieval_with_single_char_query(self, in_memory_doc_store):
        docs = [
            Document(content="C programming language"),
            Document(content="Java programming language"),
            Document(content="Python programming language"),
        ]
        in_memory_doc_store.write_documents(docs)

        results = in_memory_doc_store.bm25_retrieval(query="C", top_k=1)
        assert len(results) == 1
        assert results[0].content == "C programming language"

    def test_bm25_retrieval_single_char_content_token(self, in_memory_doc_store):
        docs = [Document(content="I like R"), Document(content="I like Python")]
        in_memory_doc_store.write_documents(docs)

        results = in_memory_doc_store.bm25_retrieval(query="R programming", top_k=1)
        assert len(results) == 1
        assert results[0].content == "I like R"

    def test_bm25_avg_doc_len_correctness(self, in_memory_doc_store):
        """Average document length should be computed correctly after writes."""
        # Write documents with known token counts.
        # "hello world" -> 2 tokens, "foo bar baz" -> 3 tokens, "go" -> 1 token
        in_memory_doc_store.write_documents(
            [
                Document(content="hello world", id="d1"),
                Document(content="foo bar baz", id="d2"),
                Document(content="go", id="d3"),
            ]
        )
        # Average should be (2 + 3 + 1) / 3 = 2.0
        assert in_memory_doc_store._avg_doc_len == pytest.approx(2.0)

    def test_bm25_avg_doc_len_after_delete(self, in_memory_doc_store):
        """Average document length should remain correct after deletion."""
        in_memory_doc_store.write_documents(
            [
                Document(content="hello world", id="d1"),  # 2 tokens
                Document(content="foo bar baz", id="d2"),  # 3 tokens
            ]
        )
        assert in_memory_doc_store._avg_doc_len == pytest.approx(2.5)
        in_memory_doc_store.delete_documents(["d1"])
        # After removing "hello world" (2 tokens), only "foo bar baz" (3 tokens) remains
        assert in_memory_doc_store._avg_doc_len == pytest.approx(3.0)


class TestMemoryDocumentStoreNotShared(TestMemoryDocumentStore):
    """
    Runs the full DocumentStore conformance suite against a non-shared (instance-local) store.

    A store created with shared=False keeps its data on the instance instead of in the process-global storage,
    so this re-runs every protocol test to confirm all operations behave identically through the instance-local
    code path. It also holds the tests specific to the shared/non-shared storage behavior.
    """

    @pytest.fixture
    def document_store(self):
        store = InMemoryDocumentStore(bm25_algorithm="BM25L", shared=False)
        yield store
        store.shutdown()

    @pytest.fixture
    def cosine_document_store(self):
        store = InMemoryDocumentStore(embedding_similarity_function="cosine", shared=False)
        yield store
        store.shutdown()

    def test_default_store_is_shared_and_registers_global_storage(self):
        index = "test_default_store_is_shared_and_registers_global_storage"
        store = InMemoryDocumentStore(index=index)
        try:
            assert store._shared is True
            assert index in in_memory_module._STORAGES
        finally:
            store.shutdown()
            for storage in (
                in_memory_module._STORAGES,
                in_memory_module._BM25_STATS_STORAGES,
                in_memory_module._AVERAGE_DOC_LEN_STORAGES,
                in_memory_module._FREQ_VOCAB_FOR_IDF_STORAGES,
            ):
                storage.pop(index, None)

    def test_shared_false_keeps_storage_instance_local(self):
        index = "test_shared_false_keeps_storage_instance_local"
        store = InMemoryDocumentStore(index=index, shared=False)
        assert store._shared is False

        store.write_documents([Document(content="Hello world")])
        assert store.count_documents() == 1
        # Nothing is registered in the process-global storage.
        assert index not in in_memory_module._STORAGES

        # A second store with the same index does not see the first one's documents (no sharing).
        other = InMemoryDocumentStore(index=index, shared=False)
        assert other.count_documents() == 0