feat(astra): add run_async to AstraEmbeddingRetriever

integrations/astra/pyproject.toml

@@ -159,4 +159,5 @@ exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
 [tool.pytest.ini_options]
 minversion = "6.0"
+asyncio_mode = "auto"
 markers = ["integration: integration tests"]

integrations/astra/src/haystack_integrations/components/retrievers/astra/retriever.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import asyncio
 from typing import Any
 
 from haystack import Document, component, default_from_dict, default_to_dict
@@ -82,6 +83,32 @@ def run(
 
         return {"documents": self.document_store.search(query_embedding, top_k, filters=filters)}
 
+    @component.output_types(documents=list[Document])
+    async def run_async(
+        self,
+        query_embedding: list[float],
+        filters: dict[str, Any] | None = None,
+        top_k: int | None = None,
+    ) -> dict[str, list[Document]]:
+        """
+        Retrieve documents from the AstraDocumentStore asynchronously.
+
+        Runs the sync search in a thread pool to avoid blocking the event loop.
+
+        :param query_embedding: floats representing the query embedding
+        :param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on
+                        the `filter_policy` chosen at retriever initialization. See init method docstring for more
+                        details.
+        :param top_k: the maximum number of documents to retrieve.
+        :returns: a dictionary with the following keys:
+            - `documents`: A list of documents retrieved from the AstraDocumentStore.
+        """
+        filters = apply_filter_policy(self.filter_policy, self.filters, filters)
+        top_k = top_k or self.top_k
+
+        documents = await asyncio.to_thread(self.document_store.search, query_embedding, top_k, filters=filters)
+        return {"documents": documents}
+
     def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.

integrations/astra/tests/test_retriever.py

@@ -4,6 +4,7 @@
 from unittest.mock import patch
 
 import pytest
+from haystack import Document
 from haystack.document_stores.types import FilterPolicy
 
 from haystack_integrations.components.retrievers.astra import AstraEmbeddingRetriever
@@ -122,3 +123,58 @@ def test_retriever_from_json_no_filter_policy(*_):
     assert retriever.top_k == 42
     assert retriever.filters == {"bar": "baz"}
     assert retriever.filter_policy == FilterPolicy.REPLACE  # defaults to REPLACE
+
+
+@patch.dict(
+    "os.environ",
+    {"ASTRA_DB_APPLICATION_TOKEN": "fake-token", "ASTRA_DB_API_ENDPOINT": "http://fake-url.apps.astra.datastax.com"},
+)
+@patch("haystack_integrations.document_stores.astra.document_store.AstraClient")
+@pytest.mark.asyncio
+async def test_run_async(*_):
+    ds = AstraDocumentStore()
+    mock_doc = Document(content="test", id="1")
+    with patch.object(ds, "search", return_value=[mock_doc]):
+        retriever = AstraEmbeddingRetriever(ds, top_k=5)
+        result = await retriever.run_async(query_embedding=[0.1] * 768)
+        assert result["documents"] == [mock_doc]
+        ds.search.assert_called_once()
+        call_args = ds.search.call_args[0]
+        assert call_args[0] == [0.1] * 768
+        assert call_args[1] == 5
+        assert ds.search.call_args.kwargs["filters"] == {}
+
+
+@patch.dict(
+    "os.environ",
+    {"ASTRA_DB_APPLICATION_TOKEN": "fake-token", "ASTRA_DB_API_ENDPOINT": "http://fake-url.apps.astra.datastax.com"},
+)
+@patch("haystack_integrations.document_stores.astra.document_store.AstraClient")
+@pytest.mark.asyncio
+async def test_run_async_filters_replace(*_):
+    ds = AstraDocumentStore()
+    mock_doc = Document(content="test", id="1")
+    with patch.object(ds, "search", return_value=[mock_doc]):
+        retriever = AstraEmbeddingRetriever(ds, top_k=5, filters={"lang": "en"}, filter_policy=FilterPolicy.REPLACE)
+        await retriever.run_async(query_embedding=[0.1] * 768, filters={"year": 2024})
+        assert ds.search.call_args.kwargs["filters"] == {"year": 2024}
+
+
+@patch.dict(
+    "os.environ",
+    {"ASTRA_DB_APPLICATION_TOKEN": "fake-token", "ASTRA_DB_API_ENDPOINT": "http://fake-url.apps.astra.datastax.com"},
+)
+@patch("haystack_integrations.document_stores.astra.document_store.AstraClient")
+@pytest.mark.asyncio
+async def test_run_async_filters_merge(*_):
+    ds = AstraDocumentStore()
+    mock_doc = Document(content="test", id="1")
+    init_filters = {"field": "lang", "operator": "==", "value": "en"}
+    runtime_filters = {"field": "year", "operator": "==", "value": 2024}
+    with patch.object(ds, "search", return_value=[mock_doc]):
+        retriever = AstraEmbeddingRetriever(ds, top_k=5, filters=init_filters, filter_policy=FilterPolicy.MERGE)
+        await retriever.run_async(query_embedding=[0.1] * 768, filters=runtime_filters)
+        merged = ds.search.call_args.kwargs["filters"]
+        assert merged["operator"] == "AND"
+        assert init_filters in merged["conditions"]
+        assert runtime_filters in merged["conditions"]