# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 import os from typing import Any from unittest.mock import ANY, AsyncMock, Mock import pytest from haystack import Document, component from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder from haystack.components.retrievers import ( InMemoryBM25Retriever, InMemoryEmbeddingRetriever, MultiRetriever, TextEmbeddingRetriever, ) from haystack.components.writers import DocumentWriter from haystack.document_stores.in_memory import InMemoryDocumentStore from haystack.document_stores.types import DuplicatePolicy from haystack.utils.experimental import ExperimentalWarning pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") @component class MockRetriever: def __init__(self, documents: list[Document] | None = None): self.documents = documents or [] @component.output_types(documents=list[Document]) def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None): return {"documents": self.documents} @component class FailingRetriever: @component.output_types(documents=list[Document]) def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None): raise RuntimeError("connection error") @pytest.fixture def sample_documents(): return [ Document( content="Renewable energy is energy that is collected from renewable resources.", meta={"category": "renewable"}, id="doc1", ), Document( content="Solar energy is a type of green energy that is harnessed from the sun.", meta={"category": "solar"}, id="doc2", ), Document( content="Wind energy is another type of green energy that is generated by wind turbines.", meta={"category": "wind"}, id="doc3", ), Document( content="Geothermal energy is heat that comes from the sub-surface of the earth.", meta={"category": "geothermal"}, id="doc4", ), Document( content="Fossil fuels, such as coal, oil, and natural gas, are non-renewable energy sources.", meta={"category": "fossil"}, id="doc5", ), ] @pytest.fixture def document_store_with_embeddings(sample_documents): """Create a document store populated with embedded documents.""" document_store = InMemoryDocumentStore() doc_embedder = OpenAIDocumentEmbedder() doc_writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP) embedded_docs = doc_embedder.run(sample_documents)["documents"] doc_writer.run(documents=embedded_docs) return document_store @pytest.fixture def bm25_retriever(document_store_with_embeddings): return InMemoryBM25Retriever(document_store=document_store_with_embeddings) @pytest.fixture def embedding_retriever(document_store_with_embeddings): return TextEmbeddingRetriever( retriever=InMemoryEmbeddingRetriever(document_store=document_store_with_embeddings), text_embedder=OpenAITextEmbedder(), ) class TestMultiRetriever: def test_init_default_parameters(self): retrievers = {"mock": MockRetriever()} retriever = MultiRetriever(retrievers=retrievers) assert retriever.retrievers == retrievers assert retriever.filters is None assert retriever.top_k_per_retriever is None assert retriever.top_k is None assert retriever.max_workers == 4 assert retriever.join_mode == "reciprocal_rank_fusion" def test_init_custom_parameters(self): retrievers = {"mock": MockRetriever()} retriever = MultiRetriever( retrievers=retrievers, filters={"field": "meta.category"}, top_k=5, max_workers=2, join_mode="concatenate" ) assert retriever.retrievers == retrievers assert retriever.filters == {"field": "meta.category"} assert retriever.top_k == 5 assert retriever.max_workers == 2 assert retriever.join_mode == "concatenate" def test_run_rrf_assigns_scores_and_sorts(self, sample_documents): docs_a = [sample_documents[0], sample_documents[1], sample_documents[2]] docs_b = [sample_documents[2], sample_documents[0], sample_documents[3]] retriever = MultiRetriever( retrievers={"a": MockRetriever(docs_a), "b": MockRetriever(docs_b)}, join_mode="reciprocal_rank_fusion" ) result = retriever.run(query="energy") assert all(doc.score is not None for doc in result["documents"]) scores = [doc.score for doc in result["documents"]] assert scores == sorted(scores, reverse=True) # doc1 ranked 1st in a and 2nd in b, doc3 ranked 3rd in a and 1st in b — doc1 should beat doc3 ids = [doc.id for doc in result["documents"]] assert ids.index("doc1") < ids.index("doc3") def test_run_with_empty_document_store(self): retriever = MultiRetriever(retrievers={"mock": MockRetriever()}) result = retriever.run(query="green energy") assert "documents" in result assert result["documents"] == [] def test_run_combines_results_from_multiple_retrievers(self, sample_documents): retriever = MultiRetriever( retrievers={ "a": MockRetriever(documents=[sample_documents[0]]), "b": MockRetriever(documents=[sample_documents[1]]), }, max_workers=2, ) result = retriever.run(query="energy") assert len(result["documents"]) == 2 assert {doc.id for doc in result["documents"]} == {"doc1", "doc2"} def test_run_deduplicates_results(self, sample_documents): retriever = MultiRetriever( retrievers={ "c": MockRetriever(documents=[sample_documents[0], sample_documents[1]]), "d": MockRetriever(documents=[sample_documents[0]]), }, max_workers=2, ) result = retriever.run(query="energy") assert len(result["documents"]) == 2 ids = [doc.id for doc in result["documents"]] assert ids.count("doc1") == 1 def test_run_resolves_filters_and_top_k_per_retriever(self): received: dict = {} @component class CapturingRetriever: @component.output_types(documents=list[Document]) def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None): received["filters"] = filters received["top_k"] = top_k return {"documents": []} retriever = MultiRetriever( retrievers={"capturing": CapturingRetriever()}, filters={"field": "meta.category"}, top_k_per_retriever=5 ) # Should use init-time values when not overridden (top_k_per_retriever is forwarded as the retriever's top_k) retriever.run(query="energy") assert received["filters"] == {"field": "meta.category"} assert received["top_k"] == 5 # Should prefer run-time values when provided retriever.run(query="energy", filters={"field": "meta.other"}, top_k_per_retriever=2) assert received["filters"] == {"field": "meta.other"} assert received["top_k"] == 2 def test_run_forwards_top_k_per_retriever_not_overall_top_k(self): received: dict = {} @component class CapturingRetriever: @component.output_types(documents=list[Document]) def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None): received["top_k"] = top_k return {"documents": []} retriever = MultiRetriever(retrievers={"capturing": CapturingRetriever()}) # top_k_per_retriever is forwarded to each retriever as its top_k retriever.run(query="energy", top_k_per_retriever=3) assert received["top_k"] == 3 # the overall top_k is applied at merge-time only, not forwarded to retrievers received.clear() retriever.run(query="energy", top_k=5) assert received.get("top_k") is None def test_run_top_k_truncates_merged_results(self, sample_documents): retriever = MultiRetriever( retrievers={ "a": MockRetriever(documents=sample_documents[:3]), "b": MockRetriever(documents=sample_documents[2:5]), }, max_workers=2, ) result = retriever.run(query="energy", top_k=2) assert len(result["documents"]) == 2 scores = [doc.score for doc in result["documents"]] assert all(score is not None for score in scores) assert scores == sorted(scores, reverse=True) def test_run_top_k_forces_rrf_in_concatenate_mode(self, sample_documents): # In concatenate mode there is no global ranking, so setting top_k falls back to RRF to truncate consistently retriever = MultiRetriever( retrievers={ "a": MockRetriever(documents=sample_documents[:3]), "b": MockRetriever(documents=sample_documents[1:4]), }, join_mode="concatenate", max_workers=2, ) result = retriever.run(query="energy", top_k=2) assert len(result["documents"]) == 2 assert all(doc.score is not None for doc in result["documents"]) def test_run_with_active_retrievers(self, sample_documents): retriever = MultiRetriever( retrievers={"a": MockRetriever([sample_documents[0]]), "b": MockRetriever([sample_documents[1]])} ) # Only run retriever "a" result = retriever.run(query="energy", active_retrievers=["a"]) assert len(result["documents"]) == 1 assert result["documents"][0].id == "doc1" def test_run_with_unknown_active_retriever_raises(self): retriever = MultiRetriever(retrievers={"mock": MockRetriever()}) with pytest.raises(ValueError, match="Unknown retriever name"): retriever.run(query="energy", active_retrievers=["nonexistent"]) def test_run_retriever_failure_raises_with_name(self): retriever = MultiRetriever(retrievers={"failing": FailingRetriever()}) with pytest.raises(RuntimeError, match="Retriever 'failing' failed"): retriever.run(query="energy") def test_to_dict(self): retriever = MultiRetriever( retrievers={"bm25": InMemoryBM25Retriever(document_store=InMemoryDocumentStore())}, filters=None, top_k_per_retriever=3, top_k=5, max_workers=2, ) result = retriever.to_dict() assert result == { "type": "haystack.components.retrievers.multi_retriever.MultiRetriever", "init_parameters": { "retrievers": { "bm25": { "type": "haystack.components.retrievers.in_memory.bm25_retriever.InMemoryBM25Retriever", "init_parameters": { "document_store": { "type": "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore", "init_parameters": { "bm25_tokenization_regex": "(?u)\\b\\w+\\b", "bm25_algorithm": "BM25L", "bm25_parameters": {}, "embedding_similarity_function": "dot_product", "index": ANY, "shared": True, "return_embedding": True, }, }, "filters": None, "top_k": 10, "scale_score": False, "filter_policy": "replace", }, } }, "filters": None, "top_k_per_retriever": 3, "top_k": 5, "max_workers": 2, "join_mode": "reciprocal_rank_fusion", }, } def test_from_dict(self): data = { "type": "haystack.components.retrievers.multi_retriever.MultiRetriever", "init_parameters": { "retrievers": { "bm25": { "type": "haystack.components.retrievers.in_memory.bm25_retriever.InMemoryBM25Retriever", "init_parameters": { "document_store": { "type": "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore", "init_parameters": { "bm25_tokenization_regex": "(?u)\\b\\w\\w+\\b", "bm25_algorithm": "BM25L", "bm25_parameters": {}, "embedding_similarity_function": "dot_product", "index": "4bb5369d-779f-487b-9c16-3c40f503438b", "shared": True, "return_embedding": True, }, }, "filters": None, "top_k": 10, "scale_score": False, "filter_policy": "replace", }, } }, "filters": None, "top_k_per_retriever": 3, "top_k": 5, "max_workers": 2, "join_mode": "concatenate", }, } result = MultiRetriever.from_dict(data) assert isinstance(result, MultiRetriever) assert len(result.retrievers) == 1 assert "bm25" in result.retrievers assert isinstance(result.retrievers["bm25"], InMemoryBM25Retriever) assert result.top_k_per_retriever == 3 assert result.top_k == 5 assert result.max_workers == 2 assert result.join_mode == "concatenate" def test_from_dict_with_no_retrievers(self): data = { "type": "haystack.components.retrievers.multi_retriever.MultiRetriever", "init_parameters": {"retrievers": {}, "filters": None, "top_k": 10, "max_workers": 4}, } result = MultiRetriever.from_dict(data) assert isinstance(result, MultiRetriever) assert result.retrievers == {} def test_from_dict_with_unknown_retriever_type_raises(self): data = { "type": "haystack.components.retrievers.multi_retriever.MultiRetriever", "init_parameters": { "retrievers": { "bad": {"type": "haystack.components.retrievers.NonExistentRetriever", "init_parameters": {}} }, "filters": None, "top_k": 10, "max_workers": 4, }, } with pytest.raises(ImportError, match="Could not import class"): MultiRetriever.from_dict(data) @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set") @pytest.mark.integration def test_run_with_filters(self, bm25_retriever, embedding_retriever): retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever}) result = retriever.run(query="energy", filters={"field": "meta.category", "operator": "==", "value": "solar"}) assert len(result["documents"]) == 1 assert result["documents"][0].meta["category"] == "solar" @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set") @pytest.mark.integration def test_run_with_top_k(self, bm25_retriever, embedding_retriever): retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever}) result = retriever.run(query="energy", top_k=2) assert len(result["documents"]) == 2 @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set") @pytest.mark.integration def test_run_with_active_retrievers_integration(self, bm25_retriever, embedding_retriever): retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever}) result_bm25_active = retriever.run(query="energy", active_retrievers=["bm25"]) result_bm25 = bm25_retriever.run(query="energy") # Scores differ because MultiRetriever applies join_mode processing (e.g. RRF) even for a single retriever. assert [doc.id for doc in result_bm25_active["documents"]] == [doc.id for doc in result_bm25["documents"]] class TestMultiRetrieverAsync: @pytest.mark.asyncio async def test_run_async_with_empty_results(self): retriever = MultiRetriever(retrievers={"mock": MockRetriever()}) result = await retriever.run_async(query="green energy") assert "documents" in result assert result["documents"] == [] @pytest.mark.asyncio async def test_run_async_combines_results_from_multiple_retrievers(self, sample_documents): retriever = MultiRetriever( retrievers={ "a": MockRetriever(documents=[sample_documents[0]]), "b": MockRetriever(documents=[sample_documents[1]]), } ) result = await retriever.run_async(query="energy") assert len(result["documents"]) == 2 assert {doc.id for doc in result["documents"]} == {"doc1", "doc2"} @pytest.mark.asyncio async def test_run_async_deduplicates_results(self, sample_documents): retriever = MultiRetriever( retrievers={ "c": MockRetriever(documents=[sample_documents[0], sample_documents[1]]), "d": MockRetriever(documents=[sample_documents[0]]), } ) result = await retriever.run_async(query="energy") assert len(result["documents"]) == 2 assert [doc.id for doc in result["documents"]].count("doc1") == 1 @pytest.mark.asyncio async def test_run_async_rrf_assigns_scores_and_sorts(self, sample_documents): docs_a = [sample_documents[0], sample_documents[1], sample_documents[2]] docs_b = [sample_documents[2], sample_documents[0], sample_documents[3]] retriever = MultiRetriever( retrievers={"a": MockRetriever(docs_a), "b": MockRetriever(docs_b)}, join_mode="reciprocal_rank_fusion" ) result = await retriever.run_async(query="energy") assert all(doc.score is not None for doc in result["documents"]) scores = [doc.score for doc in result["documents"]] assert scores == sorted(scores, reverse=True) ids = [doc.id for doc in result["documents"]] assert ids.index("doc1") < ids.index("doc3") @pytest.mark.asyncio async def test_run_async_resolves_filters_and_top_k_per_retriever(self): received: dict = {} @component class CapturingRetriever: @component.output_types(documents=list[Document]) def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None): received["filters"] = filters received["top_k"] = top_k return {"documents": []} retriever = MultiRetriever( retrievers={"capturing": CapturingRetriever()}, filters={"field": "meta.category"}, top_k_per_retriever=5 ) # top_k_per_retriever is forwarded as the retriever's top_k await retriever.run_async(query="energy") assert received["filters"] == {"field": "meta.category"} assert received["top_k"] == 5 await retriever.run_async(query="energy", filters={"field": "meta.other"}, top_k_per_retriever=2) assert received["filters"] == {"field": "meta.other"} assert received["top_k"] == 2 @pytest.mark.asyncio async def test_run_async_forwards_top_k_per_retriever_not_overall_top_k(self): received: dict = {} @component class CapturingRetriever: @component.output_types(documents=list[Document]) def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None): received["top_k"] = top_k return {"documents": []} retriever = MultiRetriever(retrievers={"capturing": CapturingRetriever()}) # top_k_per_retriever is forwarded to each retriever as its top_k await retriever.run_async(query="energy", top_k_per_retriever=3) assert received["top_k"] == 3 # the overall top_k is applied at merge-time only, not forwarded to retrievers received.clear() await retriever.run_async(query="energy", top_k=5) assert received.get("top_k") is None @pytest.mark.asyncio async def test_run_async_top_k_truncates_merged_results(self, sample_documents): retriever = MultiRetriever( retrievers={ "a": MockRetriever(documents=sample_documents[:3]), "b": MockRetriever(documents=sample_documents[2:5]), } ) result = await retriever.run_async(query="energy", top_k=2) assert len(result["documents"]) == 2 scores = [doc.score for doc in result["documents"]] assert all(score is not None for score in scores) assert scores == sorted(scores, reverse=True) @pytest.mark.asyncio async def test_run_async_top_k_forces_rrf_in_concatenate_mode(self, sample_documents): # In concatenate mode there is no global ranking, so setting top_k falls back to RRF to truncate consistently retriever = MultiRetriever( retrievers={ "a": MockRetriever(documents=sample_documents[:3]), "b": MockRetriever(documents=sample_documents[1:4]), }, join_mode="concatenate", ) result = await retriever.run_async(query="energy", top_k=2) assert len(result["documents"]) == 2 assert all(doc.score is not None for doc in result["documents"]) @pytest.mark.asyncio async def test_run_async_with_active_retrievers(self, sample_documents): retriever = MultiRetriever( retrievers={"a": MockRetriever([sample_documents[0]]), "b": MockRetriever([sample_documents[1]])} ) result = await retriever.run_async(query="energy", active_retrievers=["a"]) assert len(result["documents"]) == 1 assert result["documents"][0].id == "doc1" @pytest.mark.asyncio async def test_run_async_with_unknown_active_retriever_raises(self): retriever = MultiRetriever(retrievers={"mock": MockRetriever()}) with pytest.raises(ValueError, match="Unknown retriever name"): await retriever.run_async(query="energy", active_retrievers=["nonexistent"]) @pytest.mark.asyncio async def test_run_async_retriever_failure_raises_with_name(self): retriever = MultiRetriever(retrievers={"failing": FailingRetriever()}) with pytest.raises(RuntimeError, match="Retriever 'failing' failed"): await retriever.run_async(query="energy") @pytest.mark.asyncio async def test_run_async_uses_run_async_on_retriever_if_available(self): @component class AsyncCapableRetriever: def __init__(self): self.used_async = False @component.output_types(documents=list[Document]) def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None): return {"documents": []} @component.output_types(documents=list[Document]) async def run_async(self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None): self.used_async = True return {"documents": [Document(content="async result", id="async1")]} inner = AsyncCapableRetriever() retriever = MultiRetriever(retrievers={"async_capable": inner}) result = await retriever.run_async(query="energy") assert inner.used_async is True assert len(result["documents"]) == 1 assert result["documents"][0].id == "async1" @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set") @pytest.mark.integration @pytest.mark.asyncio async def test_run_async_with_filters(self, bm25_retriever, embedding_retriever): retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever}) result = await retriever.run_async( query="energy", filters={"field": "meta.category", "operator": "==", "value": "solar"} ) assert len(result["documents"]) == 1 assert result["documents"][0].meta["category"] == "solar" @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set") @pytest.mark.integration @pytest.mark.asyncio async def test_run_async_with_top_k(self, bm25_retriever, embedding_retriever): retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever}) result = await retriever.run_async(query="energy", top_k=2) assert len(result["documents"]) == 2 @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set") @pytest.mark.integration @pytest.mark.asyncio async def test_run_async_with_active_retrievers_integration(self, bm25_retriever, embedding_retriever): retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever}) result_bm25_active = await retriever.run_async(query="energy", active_retrievers=["bm25"]) result_bm25 = await bm25_retriever.run_async(query="energy") # Scores differ because MultiRetriever applies join_mode processing (e.g. RRF) even for a single retriever. assert [doc.id for doc in result_bm25_active["documents"]] == [doc.id for doc in result_bm25["documents"]] class TestMultiRetrieverExperimental: @pytest.mark.filterwarnings("always::haystack.utils.experimental.ExperimentalWarning") def test_emits_experimental_warning_on_init(self): with pytest.warns(ExperimentalWarning, match="MultiRetriever.*experimental"): MultiRetriever(retrievers={"mock": MockRetriever()}) @pytest.mark.filterwarnings("always::haystack.utils.experimental.ExperimentalWarning") def test_experimental_attribute_is_set(self): assert getattr(MultiRetriever, "__experimental__", False) is True class TestComponentLifecycle: def test_warm_up_delegates_to_all_retrievers(self): a = Mock(spec=["run", "warm_up"]) b = Mock(spec=["run", "warm_up"]) retriever = MultiRetriever(retrievers={"a": a, "b": b}) retriever.warm_up() a.warm_up.assert_called_once() b.warm_up.assert_called_once() async def test_warm_up_async_delegates_to_all_retrievers(self): a = Mock(spec=["run", "warm_up_async"]) a.warm_up_async = AsyncMock() b = Mock(spec=["run", "warm_up_async"]) b.warm_up_async = AsyncMock() retriever = MultiRetriever(retrievers={"a": a, "b": b}) await retriever.warm_up_async() a.warm_up_async.assert_awaited_once() b.warm_up_async.assert_awaited_once() async def test_warm_up_async_falls_back_to_sync_warm_up(self): a = Mock(spec=["run", "warm_up"]) b = Mock(spec=["run", "warm_up"]) retriever = MultiRetriever(retrievers={"a": a, "b": b}) await retriever.warm_up_async() a.warm_up.assert_called_once() b.warm_up.assert_called_once() def test_close_delegates_to_all_retrievers(self): a = Mock(spec=["run", "close"]) b = Mock(spec=["run", "close"]) retriever = MultiRetriever(retrievers={"a": a, "b": b}) retriever.close() a.close.assert_called_once() b.close.assert_called_once() async def test_close_async_delegates_to_all_retrievers(self): a = Mock(spec=["run", "close_async"]) a.close_async = AsyncMock() b = Mock(spec=["run", "close_async"]) b.close_async = AsyncMock() retriever = MultiRetriever(retrievers={"a": a, "b": b}) await retriever.close_async() a.close_async.assert_awaited_once() b.close_async.assert_awaited_once() async def test_close_async_falls_back_to_sync_close(self): a = Mock(spec=["run", "close"]) b = Mock(spec=["run", "close"]) retriever = MultiRetriever(retrievers={"a": a, "b": b}) await retriever.close_async() a.close.assert_called_once() b.close.assert_called_once() async def test_lifecycle_is_safe_when_retrievers_lack_methods(self): a = Mock(spec=["run"]) b = Mock(spec=["run"]) retriever = MultiRetriever(retrievers={"a": a, "b": b}) retriever.warm_up() await retriever.warm_up_async() retriever.close() await retriever.close_async()