# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 import os from unittest.mock import ANY, AsyncMock, Mock import pytest from haystack import Document, Pipeline from haystack.components.generators.chat import OpenAIChatGenerator from haystack.components.query import QueryExpander from haystack.components.retrievers import InMemoryBM25Retriever, MultiQueryTextRetriever from haystack.components.writers import DocumentWriter from haystack.document_stores.in_memory import InMemoryDocumentStore from haystack.document_stores.types import DuplicatePolicy class TestMultiQueryTextRetriever: @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_docs(self, sample_documents): document_store = InMemoryDocumentStore() doc_writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP) doc_writer.run(documents=sample_documents) return document_store def test_init_with_default_parameters(self, in_memory_doc_store): in_memory_retriever = InMemoryBM25Retriever(document_store=in_memory_doc_store) retriever = MultiQueryTextRetriever(retriever=in_memory_retriever) assert retriever.retriever == in_memory_retriever assert retriever.max_workers == 3 def test_init_with_custom_parameters(self, in_memory_doc_store): in_memory_retriever = InMemoryBM25Retriever(document_store=in_memory_doc_store) retriever = MultiQueryTextRetriever(retriever=in_memory_retriever, max_workers=2) assert retriever.retriever == in_memory_retriever assert retriever.max_workers == 2 def test_to_dict(self, in_memory_doc_store): in_memory_retriever = InMemoryBM25Retriever(document_store=in_memory_doc_store) multi_retriever = MultiQueryTextRetriever(retriever=in_memory_retriever, max_workers=2) result = multi_retriever.to_dict() assert result == { "type": "haystack.components.retrievers.multi_query_text_retriever.MultiQueryTextRetriever", "init_parameters": { "retriever": { "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", }, }, "max_workers": 2, }, } def test_from_dict(self): data = { "type": "haystack.components.retrievers.multi_query_text_retriever.MultiQueryTextRetriever", "init_parameters": { "retriever": { "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": "88144fa9-6e45-4e5d-8647-4c4002d8b6db", "shared": True, "return_embedding": True, }, }, "filters": None, "top_k": 10, "scale_score": False, "filter_policy": "replace", }, }, "max_workers": 3, }, } result = MultiQueryTextRetriever.from_dict(data) assert isinstance(result, MultiQueryTextRetriever) assert result.retriever.__class__.__name__ == "InMemoryBM25Retriever" assert result.max_workers == 3 def test_run_with_multiple_queries(self, document_store_with_docs): in_memory_retriever = InMemoryBM25Retriever(document_store=document_store_with_docs) multi_retriever = MultiQueryTextRetriever(retriever=in_memory_retriever) queries = ["renewable energy", "solar power", "wind turbines"] result = multi_retriever.run(queries=queries) assert "documents" in result assert len(result["documents"]) > 0 assert all(isinstance(doc, Document) for doc in result["documents"]) scores = [doc.score for doc in result["documents"] if doc.score is not None] assert scores == sorted(scores, reverse=True) @pytest.mark.integration def test_run_with_filters(self, document_store_with_docs): in_memory_retriever = InMemoryBM25Retriever(document_store=document_store_with_docs) filters = {"field": "category", "operator": "==", "value": "solar"} multi_retriever = MultiQueryTextRetriever(retriever=in_memory_retriever) result = multi_retriever.run(queries=["energy"], retriever_kwargs={"filters": filters}) 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_docs): pipeline = Pipeline() expander = QueryExpander( chat_generator=OpenAIChatGenerator(model="gpt-4.1-nano"), n_expansions=3, include_original_query=True ) in_memory_retriever = InMemoryBM25Retriever(document_store=document_store_with_docs) multiquery_retriever = MultiQueryTextRetriever(retriever=in_memory_retriever, max_workers=3) 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 contents = [doc.content for doc in results["multiquery_retriever"]["documents"]] assert len(contents) == len(set(contents)) class TestComponentLifecycle: def test_warm_up_delegates_to_retriever(self): retriever = Mock(spec=["run", "warm_up"]) component = MultiQueryTextRetriever(retriever=retriever) component.warm_up() retriever.warm_up.assert_called_once() async def test_warm_up_async_delegates_to_retriever(self): retriever = Mock(spec=["run", "warm_up_async"]) retriever.warm_up_async = AsyncMock() component = MultiQueryTextRetriever(retriever=retriever) await component.warm_up_async() retriever.warm_up_async.assert_awaited_once() async def test_warm_up_async_falls_back_to_sync_warm_up(self): retriever = Mock(spec=["run", "warm_up"]) component = MultiQueryTextRetriever(retriever=retriever) await component.warm_up_async() retriever.warm_up.assert_called_once() def test_close_delegates_to_retriever(self): retriever = Mock(spec=["run", "close"]) component = MultiQueryTextRetriever(retriever=retriever) component.close() retriever.close.assert_called_once() async def test_close_async_delegates_to_retriever(self): retriever = Mock(spec=["run", "close_async"]) retriever.close_async = AsyncMock() component = MultiQueryTextRetriever(retriever=retriever) await component.close_async() retriever.close_async.assert_awaited_once() async def test_close_async_falls_back_to_sync_close(self): retriever = Mock(spec=["run", "close"]) component = MultiQueryTextRetriever(retriever=retriever) await component.close_async() retriever.close.assert_called_once() async def test_lifecycle_is_safe_when_retriever_lacks_methods(self): retriever = Mock(spec=["run"]) component = MultiQueryTextRetriever(retriever=retriever) component.warm_up() await component.warm_up_async() component.close() await component.close_async()