deepset-ai--haystack/test/components/retrievers/test_auto_merging_retriever_async.py

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

import pytest

from haystack import Document
from haystack.components.preprocessors import HierarchicalDocumentSplitter
from haystack.components.retrievers.auto_merging_retriever import AutoMergingRetriever


class TestAutoMergingRetrieverAsync:
    @pytest.mark.asyncio
    async def test_run_missing_parent_id(self, in_memory_doc_store):
        docs = [Document(content="test", meta={"__level": 1, "__block_size": 10})]
        retriever = AutoMergingRetriever(in_memory_doc_store)
        with pytest.raises(
            ValueError, match="The matched leaf documents do not have the required meta field '__parent_id'"
        ):
            await retriever.run_async(documents=docs)

    @pytest.mark.asyncio
    async def test_run_missing_level(self, in_memory_doc_store):
        docs = [Document(content="test", meta={"__parent_id": "parent1", "__block_size": 10})]

        retriever = AutoMergingRetriever(in_memory_doc_store)
        with pytest.raises(
            ValueError, match="The matched leaf documents do not have the required meta field '__level'"
        ):
            await retriever.run_async(documents=docs)

    @pytest.mark.asyncio
    async def test_run_missing_block_size(self, in_memory_doc_store):
        docs = [Document(content="test", meta={"__parent_id": "parent1", "__level": 1})]

        retriever = AutoMergingRetriever(in_memory_doc_store)
        with pytest.raises(
            ValueError, match="The matched leaf documents do not have the required meta field '__block_size'"
        ):
            await retriever.run_async(documents=docs)

    @pytest.mark.asyncio
    async def test_run_mixed_valid_and_invalid_documents(self, in_memory_doc_store):
        docs = [
            Document(content="valid", meta={"__parent_id": "parent1", "__level": 1, "__block_size": 10}),
            Document(content="invalid", meta={"__level": 1, "__block_size": 10}),
        ]
        retriever = AutoMergingRetriever(in_memory_doc_store)
        with pytest.raises(
            ValueError, match="The matched leaf documents do not have the required meta field '__parent_id'"
        ):
            await retriever.run_async(documents=docs)

    @pytest.mark.asyncio
    async def test_run_parent_not_found(self, in_memory_doc_store):
        retriever = AutoMergingRetriever(in_memory_doc_store, threshold=0.5)

        # a leaf document with a non-existent parent_id
        leaf_doc = Document(
            content="test", meta={"__parent_id": "non_existent_parent", "__level": 1, "__block_size": 10}
        )

        with pytest.raises(ValueError, match="Expected 1 parent document with id non_existent_parent, found 0"):
            await retriever.run_async([leaf_doc])

    @pytest.mark.asyncio
    async def test_run_parent_without_children_metadata(self, in_memory_doc_store):
        """Test case where a parent document exists but doesn't have the __children_ids metadata field"""
        # Create and store a parent document without __children_ids metadata
        parent_doc = Document(
            content="parent content",
            id="parent1",
            meta={
                "__level": 1,  # Add other required metadata
                "__block_size": 10,
            },
        )
        in_memory_doc_store.write_documents([parent_doc])

        retriever = AutoMergingRetriever(in_memory_doc_store, threshold=0.5)

        # Create a leaf document that points to this parent
        leaf_doc = Document(content="leaf content", meta={"__parent_id": "parent1", "__level": 2, "__block_size": 5})

        with pytest.raises(ValueError, match="Parent document with id parent1 does not have any children"):
            await retriever.run_async([leaf_doc])

    @pytest.mark.asyncio
    async def test_run_empty_documents(self, in_memory_doc_store):
        retriever = AutoMergingRetriever(in_memory_doc_store)
        assert await retriever.run_async([]) == {"documents": []}

    @pytest.mark.asyncio
    async def test_run_return_parent_document(self, in_memory_doc_store):
        text = "The sun rose early in the morning. It cast a warm glow over the trees. Birds began to sing."

        docs = [Document(content=text)]
        builder = HierarchicalDocumentSplitter(block_sizes={10, 3}, split_overlap=0, split_by="word")
        docs = builder.run(docs)

        # store all non-leaf documents
        for doc in docs["documents"]:
            if doc.meta["__children_ids"]:
                in_memory_doc_store.write_documents([doc])
        retriever = AutoMergingRetriever(in_memory_doc_store, threshold=0.5)

        # assume we retrieved 2 leaf docs from the same parent, the parent document should be returned,
        # since it has 3 children and the threshold=0.5, and we retrieved 2 children (2/3 > 0.66(6))
        leaf_docs = [doc for doc in docs["documents"] if not doc.meta["__children_ids"]]
        docs = await retriever.run_async(leaf_docs[4:6])
        assert len(docs["documents"]) == 1
        assert docs["documents"][0].content == "warm glow over the trees. Birds began to sing."
        assert len(docs["documents"][0].meta["__children_ids"]) == 3

    @pytest.mark.asyncio
    async def test_run_return_leafs_document(self, in_memory_doc_store):
        docs = [Document(content="The monarch of the wild blue yonder rises from the eastern side of the horizon.")]
        builder = HierarchicalDocumentSplitter(block_sizes={10, 3}, split_overlap=0, split_by="word")
        docs = builder.run(docs)

        for doc in docs["documents"]:
            if doc.meta["__level"] == 1:
                in_memory_doc_store.write_documents([doc])

        leaf_docs = [doc for doc in docs["documents"] if not doc.meta["__children_ids"]]
        retriever = AutoMergingRetriever(in_memory_doc_store, threshold=0.6)
        result = await retriever.run_async([leaf_docs[4]])

        assert len(result["documents"]) == 1
        assert result["documents"][0].content == "eastern side of "
        assert result["documents"][0].meta["__parent_id"] == docs["documents"][2].id

    @pytest.mark.asyncio
    async def test_run_return_leafs_document_different_parents(self, in_memory_doc_store):
        docs = [Document(content="The monarch of the wild blue yonder rises from the eastern side of the horizon.")]
        builder = HierarchicalDocumentSplitter(block_sizes={10, 3}, split_overlap=0, split_by="word")
        docs = builder.run(docs)

        for doc in docs["documents"]:
            if doc.meta["__level"] == 1:
                in_memory_doc_store.write_documents([doc])

        leaf_docs = [doc for doc in docs["documents"] if not doc.meta["__children_ids"]]
        retriever = AutoMergingRetriever(in_memory_doc_store, threshold=0.6)
        result = await retriever.run_async([leaf_docs[4], leaf_docs[3]])

        assert len(result["documents"]) == 2
        assert result["documents"][0].meta["__parent_id"] != result["documents"][1].meta["__parent_id"]

    @pytest.mark.asyncio
    async def test_run_go_up_hierarchy_multiple_levels(self, in_memory_doc_store):
        """
        Test if the retriever can go up the hierarchy multiple levels to find the parent document.

        Simulate a scenario where we have 4 leaf-documents that matched some initial query. The leaf-documents
        are continuously merged up the hierarchy until the threshold is no longer met.
        In this case it goes from the 4th level in the hierarchy up the 1st level.
        """
        text = "The sun rose early in the morning. It cast a warm glow over the trees. Birds began to sing."

        docs = [Document(content=text)]
        builder = HierarchicalDocumentSplitter(block_sizes={6, 4, 2, 1}, split_overlap=0, split_by="word")
        docs = builder.run(docs)

        # store all non-leaf documents
        for doc in docs["documents"]:
            if doc.meta["__children_ids"]:
                in_memory_doc_store.write_documents([doc])
        retriever = AutoMergingRetriever(in_memory_doc_store, threshold=0.4)

        # simulate a scenario where we have 4 leaf-documents that matched some initial query
        retrieved_leaf_docs = [d for d in docs["documents"] if d.content in {"The ", "sun ", "rose ", "early "}]

        result = await retriever.run_async(retrieved_leaf_docs)

        assert len(result["documents"]) == 1
        assert result["documents"][0].content == "The sun rose early in the "

    @pytest.mark.asyncio
    async def test_run_go_up_hierarchy_multiple_levels_hit_root_document(self, in_memory_doc_store):
        """
        Test case where we go up hierarchy until the root document, so the root document is returned.

        It's the only document in the hierarchy which has no parent.
        """
        text = "The sun rose early in the morning. It cast a warm glow over the trees. Birds began to sing."

        docs = [Document(content=text)]
        builder = HierarchicalDocumentSplitter(block_sizes={6, 4}, split_overlap=0, split_by="word")
        docs = builder.run(docs)

        # store all non-leaf documents
        for doc in docs["documents"]:
            if doc.meta["__children_ids"]:
                in_memory_doc_store.write_documents([doc])
        retriever = AutoMergingRetriever(in_memory_doc_store, threshold=0.1)  # set a low threshold to hit root document

        # simulate a scenario where we have 4 leaf-documents that matched some initial query
        retrieved_leaf_docs = [
            d
            for d in docs["documents"]
            if d.content in {"The sun rose early ", "in the ", "morning. It cast a ", "over the trees. Birds "}
        ]

        result = await retriever.run_async(retrieved_leaf_docs)

        assert len(result["documents"]) == 1
        assert result["documents"][0].meta["__level"] == 0  # hit root document