# SPDX-FileCopyrightText: 2022-present deepset GmbH # # 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