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

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

from typing import Any

import pytest

from haystack import Document, Pipeline, component
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 TestMultiQueryTextRetrieverAsync:
    @pytest.fixture
    def sample_documents(self):
        return [
            Document(
                content="Solar energy is a type of green energy that is harnessed from the sun.",
                meta={"category": "solar"},
            ),
            Document(content="Solar panels convert sunlight directly into electricity.", meta={"category": "solar"}),
            Document(content="Photovoltaic cells are the building blocks of solar panels.", meta={"category": "solar"}),
            Document(
                content="Wind energy is another type of green energy that is generated by wind turbines.",
                meta={"category": "wind"},
            ),
            Document(
                content="Geothermal energy is heat that comes from the sub-surface of the earth.",
                meta={"category": "geo"},
            ),
            Document(
                content="Renewable energy is energy that is collected from renewable resources.",
                meta={"category": "renewable"},
            ),
            Document(
                content="Hydropower is a form of renewable energy using the flow of water.", meta={"category": "hydro"}
            ),
        ]

    @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

    @pytest.mark.asyncio
    async def test_run_async_with_empty_queries(self, document_store_with_docs):
        multi_retriever = MultiQueryTextRetriever(
            retriever=InMemoryBM25Retriever(document_store=document_store_with_docs)
        )
        result = await multi_retriever.run_async(queries=[])
        assert "documents" in result
        assert result["documents"] == []

    @pytest.mark.asyncio
    async def test_run_async_with_multiple_queries(self, document_store_with_docs):
        multi_retriever = MultiQueryTextRetriever(
            retriever=InMemoryBM25Retriever(document_store=document_store_with_docs)
        )
        result = await multi_retriever.run_async(queries=["renewable energy", "solar power"])

        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.asyncio
    async def test_run_async_deduplication(self):
        doc2 = Document(content="Wind energy is clean", id="doc2", score=0.8)
        # doc3 intentionally uses the duplicate id "doc1" to simulate deduplication across multiple queries
        doc3 = Document(content="Solar energy is renewable", id="doc1", score=0.7)

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

            @component.output_types(documents=list[Document])
            async def run_async(
                self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None, **kwargs: Any
            ) -> dict[str, list[Document]]:
                return {"documents": [doc3, doc2]}

        multi_retriever = MultiQueryTextRetriever(retriever=MockRetriever())
        result = await multi_retriever.run_async(queries=["query1", "query2"])

        assert "documents" in result
        assert len(result["documents"]) == 2
        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.asyncio
    async def test_run_async_falls_back_to_sync_when_no_run_async(self):
        @component
        class SyncOnlyRetriever:
            @component.output_types(documents=list[Document])
            def run(
                self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None
            ) -> dict[str, list[Document]]:
                return {"documents": [Document(content="Renewable energy", id="doc1", score=0.9)]}

        multi_retriever = MultiQueryTextRetriever(retriever=SyncOnlyRetriever())
        result = await multi_retriever.run_async(queries=["query"])
        assert "documents" in result
        assert len(result["documents"]) == 1
        assert result["documents"][0].content == "Renewable energy"

    @pytest.mark.asyncio
    @pytest.mark.integration
    async def test_run_async_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 = await multi_retriever.run_async(
            queries=["energy", "sunlight", "photovoltaic"], retriever_kwargs={"filters": filters}
        )
        assert "documents" in result
        assert len(result["documents"]) > 0
        assert all(doc.meta.get("category") == "solar" for doc in result["documents"])

    @pytest.mark.asyncio
    @pytest.mark.integration
    async def test_run_async_with_pipeline(self, document_store_with_docs):
        multi_retriever = MultiQueryTextRetriever(
            retriever=InMemoryBM25Retriever(document_store=document_store_with_docs)
        )
        pipeline = Pipeline()
        pipeline.add_component("retriever", multi_retriever)
        result = await pipeline.run_async(data={"retriever": {"queries": ["renewable energy", "solar power"]}})

        assert result
        assert "retriever" in result
        assert "documents" in result["retriever"]
        assert len(result["retriever"]["documents"]) > 0
        scores = [doc.score for doc in result["retriever"]["documents"] if doc.score is not None]
        assert scores == sorted(scores, reverse=True)