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

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# 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()