deepset-ai--haystack/test/components/retrievers/test_multi_query_embedding_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 numpy as np
import pytest

from haystack import Document, Pipeline, component
from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.query import QueryExpander
from haystack.components.retrievers import InMemoryEmbeddingRetriever, MultiQueryEmbeddingRetriever
from haystack.components.writers import DocumentWriter
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.document_stores.types import DuplicatePolicy


@component
class MockQueryEmbedder:
    @component.output_types(embedding=list[float])
    def run(self, text: str) -> dict[str, list[float]]:
        return {"embedding": np.ones(384).tolist()}


class TestMultiQueryEmbeddingRetriever:
    @pytest.fixture
    def sample_documents(self):
        return [
            Document(
                content="Renewable energy is energy that is collected from renewable resources.",
                meta={"category": None},
            ),
            Document(
                content="Solar energy is a type of green energy that is harnessed from the sun.",
                meta={"category": "solar"},
            ),
            Document(
                content="Wind energy is another type of green energy that is generated by wind turbines",
                meta={"category": "wind"},
            ),
            Document(
                content="Hydropower is a form of renewable energy using the flow of water to generate electricity.",
                meta={"category": "hydro"},
            ),
            Document(
                content="Geothermal energy is heat that comes from the sub-surface of the earth.",
                meta={"category": "geo"},
            ),
            Document(
                content="Fossil fuels, such as coal, oil, and natural gas, are non-renewable energy sources.",
                meta={"category": "fossil"},
            ),
            Document(
                content="Nuclear energy is produced through nuclear reactions, typically using uranium or "
                "plutonium as fuel.",
                meta={"category": "nuclear"},
            ),
        ]

    @pytest.fixture
    def document_store_with_embeddings(self, 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

    def test_init_with_default_parameters(self, in_memory_doc_store):
        embedding_retriever = InMemoryEmbeddingRetriever(document_store=in_memory_doc_store)
        query_embedder = MockQueryEmbedder()
        retriever = MultiQueryEmbeddingRetriever(retriever=embedding_retriever, query_embedder=query_embedder)
        assert retriever.retriever == embedding_retriever
        assert retriever.query_embedder == query_embedder
        assert retriever.max_workers == 3

    def test_init_with_custom_parameters(self, in_memory_doc_store):
        embedding_retriever = InMemoryEmbeddingRetriever(document_store=in_memory_doc_store)
        query_embedder = MockQueryEmbedder()
        retriever = MultiQueryEmbeddingRetriever(
            retriever=embedding_retriever, query_embedder=query_embedder, max_workers=2
        )
        assert retriever.retriever == embedding_retriever
        assert retriever.query_embedder == query_embedder
        assert retriever.max_workers == 2

    def test_run_with_empty_queries(self, in_memory_doc_store):
        multi_retriever = MultiQueryEmbeddingRetriever(
            retriever=InMemoryEmbeddingRetriever(document_store=in_memory_doc_store), query_embedder=MockQueryEmbedder()
        )
        result = multi_retriever.run(queries=[])
        assert "documents" in result
        assert result["documents"] == []

    def test_run_with_empty_results(self, in_memory_doc_store):
        multi_retriever = MultiQueryEmbeddingRetriever(
            retriever=InMemoryEmbeddingRetriever(document_store=in_memory_doc_store), query_embedder=MockQueryEmbedder()
        )
        result = multi_retriever.run(queries=["query"])
        assert "documents" in result
        assert result["documents"] == []

    def test_to_dict(self, in_memory_doc_store):
        multi_retriever = MultiQueryEmbeddingRetriever(
            retriever=InMemoryEmbeddingRetriever(document_store=in_memory_doc_store),
            query_embedder=MockQueryEmbedder(),
            max_workers=2,
        )
        result = multi_retriever.to_dict()
        assert result == {
            "type": "haystack.components.retrievers.multi_query_embedding_retriever.MultiQueryEmbeddingRetriever",
            "init_parameters": {
                "retriever": {
                    "type": "haystack.components.retrievers.in_memory.embedding_retriever.InMemoryEmbeddingRetriever",
                    "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,
                        "return_embedding": False,
                        "filter_policy": "replace",
                    },
                },
                "query_embedder": {
                    "type": "retrievers.test_multi_query_embedding_retriever.MockQueryEmbedder",
                    "init_parameters": {},
                },
                "max_workers": 2,
            },
        }

    def test_from_dict(self, monkeypatch):
        monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
        data = {
            "type": "haystack.components.retrievers.multi_query_embedding_retriever.MultiQueryEmbeddingRetriever",  # noqa E501
            "init_parameters": {
                "retriever": {
                    "type": "haystack.components.retrievers.in_memory.embedding_retriever.InMemoryEmbeddingRetriever",
                    "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,
                        "return_embedding": False,
                        "filter_policy": "replace",
                    },
                },
                "query_embedder": {
                    "type": "haystack.components.embedders.openai_text_embedder.OpenAITextEmbedder",
                    "init_parameters": {
                        "api_key": {"env_vars": ["OPENAI_API_KEY"], "strict": True, "type": "env_var"},
                        "api_base_url": None,
                        "dimensions": None,
                        "model": "text-embedding-ada-002",
                        "organization": None,
                        "http_client_kwargs": None,
                        "prefix": "",
                        "suffix": "",
                        "timeout": None,
                        "max_retries": None,
                    },
                },
                "max_workers": 2,
            },
        }
        result = MultiQueryEmbeddingRetriever.from_dict(data)
        assert isinstance(result, MultiQueryEmbeddingRetriever)
        assert result.max_workers == 2

    def test_deduplication_with_overlapping_results(self):
        doc1 = Document(content="Solar energy is renewable", id="doc1", score=0.9)
        doc2 = Document(content="Wind energy is clean", id="doc2", score=0.8)
        # same as doc1 w/ different score
        doc3 = Document(content="Solar energy is renewable", id="doc1", score=0.7)

        call_count = 0

        @component
        class MockRetriever:
            @component.output_types(documents=list[Document])
            def run(
                self,
                query_embedding: list[float],
                filters: dict[str, Any] | None = None,
                top_k: int | None = None,
                **kwargs: Any,
            ) -> dict[str, list[Document]]:
                nonlocal call_count
                if call_count == 1:
                    return {"documents": [doc1, doc2]}
                return {"documents": [doc3, doc2]}

        multi_retriever = MultiQueryEmbeddingRetriever(
            retriever=MockRetriever(), query_embedder=MockQueryEmbedder(), max_workers=1
        )
        result = multi_retriever.run(queries=["query1", "query2"])

        assert "documents" in result
        assert len(result["documents"]) == 2  # Only 2 unique documents (doc1/doc3 and doc2)

        contents = [doc.content for doc in result["documents"]]
        assert contents.count("Solar energy is renewable") == 1
        assert contents.count("Wind energy is clean") == 1

    @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, document_store_with_embeddings):
        in_memory_retriever = InMemoryEmbeddingRetriever(document_store=document_store_with_embeddings)
        query_embedder = OpenAITextEmbedder()
        multi_retriever = MultiQueryEmbeddingRetriever(retriever=in_memory_retriever, query_embedder=query_embedder)
        result = multi_retriever.run(["energy"], {"filters": {"field": "category", "operator": "==", "value": "solar"}})
        assert "documents" in result
        assert all(doc.meta.get("category") == "solar" for doc in result["documents"])

    @pytest.mark.skipif(
        not os.environ.get("OPENAI_API_KEY", None),
        reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
    )
    @pytest.mark.integration
    def test_pipeline_integration(self, document_store_with_embeddings):
        expander = QueryExpander(
            chat_generator=OpenAIChatGenerator(model="gpt-4.1-nano"), n_expansions=3, include_original_query=True
        )
        in_memory_retriever = InMemoryEmbeddingRetriever(document_store=document_store_with_embeddings)
        query_embedder = OpenAITextEmbedder()
        multiquery_retriever = MultiQueryEmbeddingRetriever(
            retriever=in_memory_retriever, query_embedder=query_embedder, max_workers=3
        )

        pipeline = Pipeline()
        pipeline.add_component("query_expander", expander)
        pipeline.add_component("multiquery_retriever", multiquery_retriever)
        pipeline.connect("query_expander.queries", "multiquery_retriever.queries")

        data = {
            "query_expander": {"query": "green energy sources"},
            "multiquery_retriever": {"retriever_kwargs": {"top_k": 3}},
        }
        results = pipeline.run(data=data, include_outputs_from={"query_expander", "multiquery_retriever"})

        assert "multiquery_retriever" in results
        assert "documents" in results["multiquery_retriever"]
        assert len(results["multiquery_retriever"]["documents"]) > 0
        assert "query_expander" in results
        assert "queries" in results["query_expander"]
        assert len(results["query_expander"]["queries"]) == 4

        # assert that documents are sorted by score (highest first)
        scores = [doc.score for doc in results["multiquery_retriever"]["documents"] if doc.score is not None]
        assert scores == sorted(scores, reverse=True)

        # assert there are not duplicates
        ids = [doc.id for doc in results["multiquery_retriever"]["documents"]]
        assert len(ids) == len(set(ids))


class TestComponentLifecycle:
    def test_warm_up_delegates_to_inner_components(self):
        query_embedder = Mock(spec=["run", "warm_up"])
        retriever = Mock(spec=["run", "warm_up"])
        component = MultiQueryEmbeddingRetriever(retriever=retriever, query_embedder=query_embedder)
        component.warm_up()
        query_embedder.warm_up.assert_called_once()
        retriever.warm_up.assert_called_once()

    async def test_warm_up_async_delegates_to_inner_components(self):
        query_embedder = Mock(spec=["run", "warm_up_async"])
        query_embedder.warm_up_async = AsyncMock()
        retriever = Mock(spec=["run", "warm_up_async"])
        retriever.warm_up_async = AsyncMock()
        component = MultiQueryEmbeddingRetriever(retriever=retriever, query_embedder=query_embedder)
        await component.warm_up_async()
        query_embedder.warm_up_async.assert_awaited_once()
        retriever.warm_up_async.assert_awaited_once()

    async def test_warm_up_async_falls_back_to_sync_warm_up(self):
        query_embedder = Mock(spec=["run", "warm_up"])
        retriever = Mock(spec=["run", "warm_up"])
        component = MultiQueryEmbeddingRetriever(retriever=retriever, query_embedder=query_embedder)
        await component.warm_up_async()
        query_embedder.warm_up.assert_called_once()
        retriever.warm_up.assert_called_once()

    def test_close_delegates_to_inner_components(self):
        query_embedder = Mock(spec=["run", "close"])
        retriever = Mock(spec=["run", "close"])
        component = MultiQueryEmbeddingRetriever(retriever=retriever, query_embedder=query_embedder)
        component.close()
        query_embedder.close.assert_called_once()
        retriever.close.assert_called_once()

    async def test_close_async_delegates_to_inner_components(self):
        query_embedder = Mock(spec=["run", "close_async"])
        query_embedder.close_async = AsyncMock()
        retriever = Mock(spec=["run", "close_async"])
        retriever.close_async = AsyncMock()
        component = MultiQueryEmbeddingRetriever(retriever=retriever, query_embedder=query_embedder)
        await component.close_async()
        query_embedder.close_async.assert_awaited_once()
        retriever.close_async.assert_awaited_once()

    async def test_close_async_falls_back_to_sync_close(self):
        query_embedder = Mock(spec=["run", "close"])
        retriever = Mock(spec=["run", "close"])
        component = MultiQueryEmbeddingRetriever(retriever=retriever, query_embedder=query_embedder)
        await component.close_async()
        query_embedder.close.assert_called_once()
        retriever.close.assert_called_once()

    async def test_lifecycle_is_safe_when_inner_components_lack_methods(self):
        query_embedder = Mock(spec=["run"])
        retriever = Mock(spec=["run"])
        component = MultiQueryEmbeddingRetriever(retriever=retriever, query_embedder=query_embedder)
        component.warm_up()
        await component.warm_up_async()
        component.close()
        await component.close_async()