# 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 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()