chore: import upstream snapshot with attribution

test/components/retrievers/test_multi_query_text_retriever_async.py

@@ -0,0 +1,145 @@
+# 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)