feat: add run_async to all Weaviate retrievers (#2607)

* add run_async to all retrievers

* add :returns: docstrings

* embed user values in error

* add check for alpha range

* fix tests

* add tests for alpha checks

* Apply suggestions from code review

Co-authored-by: Stefano Fiorucci <stefanofiorucci@gmail.com>

* fix double backticks and wrong URL syntax

---------

Co-authored-by: David S. Batista <dsbatista@gmail.com>
Co-authored-by: Stefano Fiorucci <stefanofiorucci@gmail.com>

Author: 3 people

integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/bm25_retriever.py

@@ -102,9 +102,34 @@ def run(
                         details.
         :param top_k:
             The maximum number of documents to return.
+        :returns: A dictionary with the following keys:
+            - `documents`: List of documents returned by the search engine.
         """
         filters = apply_filter_policy(self._filter_policy, self._filters, filters)
 
         top_k = top_k or self._top_k
         documents = self._document_store._bm25_retrieval(query=query, filters=filters, top_k=top_k)
         return {"documents": documents}
+
+    @component.output_types(documents=list[Document])
+    async def run_async(
+        self, query: str, filters: Optional[dict[str, Any]] = None, top_k: Optional[int] = None
+    ) -> dict[str, list[Document]]:
+        """
+        Asynchronously retrieves documents from Weaviate using the BM25 algorithm.
+
+        :param query:
+            The query text.
+        :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 return.
+        :returns: A dictionary with the following keys:
+            - `documents`: List of documents returned by the search engine.
+        """
+        filters = apply_filter_policy(self._filter_policy, self._filters, filters)
+
+        top_k = top_k or self._top_k
+        documents = await self._document_store._bm25_retrieval_async(query=query, filters=filters, top_k=top_k)
+        return {"documents": documents}

integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/embedding_retriever.py

@@ -48,7 +48,7 @@ def __init__(
             `distance` and `certainty` parameters.
         """
         if distance is not None and certainty is not None:
-            msg = "Can't use 'distance' and 'certainty' parameters together"
+            msg = f"Can't use 'distance' ({distance}) and 'certainty' ({certainty}) parameters together"
             raise ValueError(msg)
 
         self._document_store = document_store
@@ -121,6 +121,8 @@ def run(
             The maximum allowed distance between Documents' embeddings.
         :param certainty:
             Normalized distance between the result item and the search vector.
+        :returns: A dictionary with the following keys:
+            - `documents`: List of documents returned by the search engine.
         :raises ValueError:
             If both `distance` and `certainty` are provided.
             See https://weaviate.io/developers/weaviate/api/graphql/search-operators#variables to learn more about
@@ -132,7 +134,7 @@ def run(
         distance = distance or self._distance
         certainty = certainty or self._certainty
         if distance is not None and certainty is not None:
-            msg = "Can't use 'distance' and 'certainty' parameters together"
+            msg = f"Can't use 'distance' ({distance}) and 'certainty' ({certainty}) parameters together"
             raise ValueError(msg)
 
         documents = self._document_store._embedding_retrieval(
@@ -143,3 +145,51 @@ def run(
             certainty=certainty,
         )
         return {"documents": documents}
+
+    @component.output_types(documents=list[Document])
+    async def run_async(
+        self,
+        query_embedding: list[float],
+        filters: Optional[dict[str, Any]] = None,
+        top_k: Optional[int] = None,
+        distance: Optional[float] = None,
+        certainty: Optional[float] = None,
+    ) -> dict[str, list[Document]]:
+        """
+        Asynchronously retrieves documents from Weaviate using the vector search.
+
+        :param query_embedding:
+            Embedding of the query.
+        :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 return.
+        :param distance:
+            The maximum allowed distance between Documents' embeddings.
+        :param certainty:
+            Normalized distance between the result item and the search vector.
+        :returns: A dictionary with the following keys:
+            - `documents`: List of documents returned by the search engine.
+        :raises ValueError:
+            If both `distance` and `certainty` are provided.
+            See https://weaviate.io/developers/weaviate/api/graphql/search-operators#variables to learn more about
+            `distance` and `certainty` parameters.
+        """
+        filters = apply_filter_policy(self._filter_policy, self._filters, filters)
+        top_k = top_k or self._top_k
+
+        distance = distance or self._distance
+        certainty = certainty or self._certainty
+        if distance is not None and certainty is not None:
+            msg = f"Can't use 'distance' ({distance}) and 'certainty' ({certainty}) parameters together"
+            raise ValueError(msg)
+
+        documents = await self._document_store._embedding_retrieval_async(
+            query_embedding=query_embedding,
+            filters=filters,
+            top_k=top_k,
+            distance=distance,
+            certainty=certainty,
+        )
+        return {"documents": documents}

integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/hybrid_retriever.py

@@ -37,35 +37,39 @@ def __init__(
         :param top_k:
             Maximum number of documents to return.
         :param alpha:
-            Blending factor for hybrid retrieval in Weaviate. Must be in the range ``[0.0, 1.0]``.
+            Blending factor for hybrid retrieval in Weaviate. Must be in the range `[0.0, 1.0]`.
 
-            Weaviate hybrid search combines keyword (BM25) and vector scores into a single ranking. ``alpha`` controls
+            Weaviate hybrid search combines keyword (BM25) and vector scores into a single ranking. `alpha` controls
             how much each part contributes to the final score:
 
-            - ``alpha = 0.0``: only keyword (BM25) scoring is used.
-            - ``alpha = 1.0``: only vector similarity scoring is used.
+            - `alpha = 0.0`: only keyword (BM25) scoring is used.
+            - `alpha = 1.0`: only vector similarity scoring is used.
             - Values in between blend the two; higher values favor the vector score, lower values favor BM25.
 
-            If ``None``, the Weaviate server default is used.
+            If `None`, the Weaviate server default is used.
 
             See the official Weaviate docs on Hybrid Search parameters for more details:
-            `Hybrid search parameters <https://weaviate.io/developers/weaviate/search/hybrid#parameters>`_
-            `Hybrid Search <https://docs.weaviate.io/weaviate/concepts/search/hybrid-search>`_
+            - [Hybrid search parameters](https://weaviate.io/developers/weaviate/search/hybrid#parameters)
+            - [Hybrid Search](https://docs.weaviate.io/weaviate/concepts/search/hybrid-search)
         :param max_vector_distance:
             Optional threshold that restricts the vector part of the hybrid search to candidates within a maximum
             vector distance. Candidates with a distance larger than this threshold are excluded from the vector portion
             before blending.
 
-            Use this to prune low-quality vector matches while still benefitting from keyword recall. Leave ``None`` to
+            Use this to prune low-quality vector matches while still benefitting from keyword recall. Leave `None` to
             use Weaviate's default behavior without an explicit cutoff.
 
             See the official Weaviate docs on Hybrid Search parameters for more details:
-            - `Hybrid search parameters <https://weaviate.io/developers/weaviate/search/hybrid#parameters>`_
-            - `Hybrid Search <https://docs.weaviate.io/weaviate/concepts/search/hybrid-search>`_
+            - [Hybrid search parameters](https://weaviate.io/developers/weaviate/search/hybrid#parameters)
+            - [Hybrid Search](https://docs.weaviate.io/weaviate/concepts/search/hybrid-search)
         :param filter_policy:
             Policy to determine how filters are applied.
         """
 
+        if alpha is not None and not 0.0 <= alpha <= 1.0:
+            msg = f"alpha ({alpha}) must be in the range [0.0, 1.0]"
+            raise ValueError(msg)
+
         self._document_store = document_store
         self._filters = filters or {}
         self._top_k = top_k
@@ -134,37 +138,43 @@ def run(
         :param top_k:
             The maximum number of documents to return.
         :param alpha:
-            Blending factor for hybrid retrieval in Weaviate. Must be in the range ``[0.0, 1.0]``.
+            Blending factor for hybrid retrieval in Weaviate. Must be in the range `[0.0, 1.0]`.
 
-            Weaviate hybrid search combines keyword (BM25) and vector scores into a single ranking. ``alpha`` controls
+            Weaviate hybrid search combines keyword (BM25) and vector scores into a single ranking. `alpha` controls
             how much each part contributes to the final score:
 
-            - ``alpha = 0.0``: only keyword (BM25) scoring is used.
-            - ``alpha = 1.0``: only vector similarity scoring is used.
+            - `alpha = 0.0`: only keyword (BM25) scoring is used.
+            - `alpha = 1.0`: only vector similarity scoring is used.
             - Values in between blend the two; higher values favor the vector score, lower values favor BM25.
 
-            If ``None``, the Weaviate server default is used.
+            If `None`, the Weaviate server default is used.
 
             See the official Weaviate docs on Hybrid Search parameters for more details:
-            `Hybrid search parameters <https://weaviate.io/developers/weaviate/search/hybrid#parameters>`_
-            `Hybrid Search <https://docs.weaviate.io/weaviate/concepts/search/hybrid-search>`_
+            - [Hybrid search parameters](https://weaviate.io/developers/weaviate/search/hybrid#parameters)
+            - [Hybrid Search](https://docs.weaviate.io/weaviate/concepts/search/hybrid-search)
         :param max_vector_distance:
             Optional threshold that restricts the vector part of the hybrid search to candidates within a maximum
             vector distance. Candidates with a distance larger than this threshold are excluded from the vector portion
             before blending.
 
-            Use this to prune low-quality vector matches while still benefitting from keyword recall. Leave ``None`` to
+            Use this to prune low-quality vector matches while still benefitting from keyword recall. Leave `None` to
             use Weaviate's default behavior without an explicit cutoff.
 
             See the official Weaviate docs on Hybrid Search parameters for more details:
-            - `Hybrid search parameters <https://weaviate.io/developers/weaviate/search/hybrid#parameters>`_
-            - `Hybrid Search <https://docs.weaviate.io/weaviate/concepts/search/hybrid-search>`_
+            - [Hybrid search parameters](https://weaviate.io/developers/weaviate/search/hybrid#parameters)
+            - [Hybrid Search](https://docs.weaviate.io/weaviate/concepts/search/hybrid-search)
+        :returns: A dictionary with the following keys:
+            - `documents`: List of documents returned by the search engine.
         """
         filters = apply_filter_policy(self._filter_policy, self._filters, filters)
         top_k = self._top_k if top_k is None else top_k
         alpha = self._alpha if alpha is None else alpha
         max_vector_distance = self._max_vector_distance if max_vector_distance is None else max_vector_distance
 
+        if alpha is not None and not 0.0 <= alpha <= 1.0:
+            msg = f"alpha ({alpha}) must be in the range [0.0, 1.0]"
+            raise ValueError(msg)
+
         documents = self._document_store._hybrid_retrieval(
             query=query,
             query_embedding=query_embedding,
@@ -174,3 +184,73 @@ def run(
             max_vector_distance=max_vector_distance,
         )
         return {"documents": documents}
+
+    @component.output_types(documents=list[Document])
+    async def run_async(
+        self,
+        query: str,
+        query_embedding: list[float],
+        filters: Optional[dict[str, Any]] = None,
+        top_k: Optional[int] = None,
+        alpha: Optional[float] = None,
+        max_vector_distance: Optional[float] = None,
+    ) -> dict[str, list[Document]]:
+        """
+        Asynchronously retrieves documents from Weaviate using hybrid search.
+
+        :param query:
+            The query text.
+        :param query_embedding:
+            Embedding of the query.
+        :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 return.
+        :param alpha:
+            Blending factor for hybrid retrieval in Weaviate. Must be in the range `[0.0, 1.0]`.
+
+            Weaviate hybrid search combines keyword (BM25) and vector scores into a single ranking. `alpha` controls
+            how much each part contributes to the final score:
+
+            - `alpha = 0.0`: only keyword (BM25) scoring is used.
+            - `alpha = 1.0`: only vector similarity scoring is used.
+            - Values in between blend the two; higher values favor the vector score, lower values favor BM25.
+
+            If `None`, the Weaviate server default is used.
+
+            See the official Weaviate docs on Hybrid Search parameters for more details:
+            - [Hybrid search parameters](https://weaviate.io/developers/weaviate/search/hybrid#parameters)
+            - [Hybrid Search](https://docs.weaviate.io/weaviate/concepts/search/hybrid-search)
+        :param max_vector_distance:
+            Optional threshold that restricts the vector part of the hybrid search to candidates within a maximum
+            vector distance. Candidates with a distance larger than this threshold are excluded from the vector portion
+            before blending.
+
+            Use this to prune low-quality vector matches while still benefitting from keyword recall. Leave `None` to
+            use Weaviate's default behavior without an explicit cutoff.
+
+            See the official Weaviate docs on Hybrid Search parameters for more details:
+            - [Hybrid search parameters](https://weaviate.io/developers/weaviate/search/hybrid#parameters)
+            - [Hybrid Search](https://docs.weaviate.io/weaviate/concepts/search/hybrid-search)
+        :returns: A dictionary with the following keys:
+            - `documents`: List of documents returned by the search engine.
+        """
+        filters = apply_filter_policy(self._filter_policy, self._filters, filters)
+        top_k = self._top_k if top_k is None else top_k
+        alpha = self._alpha if alpha is None else alpha
+        max_vector_distance = self._max_vector_distance if max_vector_distance is None else max_vector_distance
+
+        if alpha is not None and not 0.0 <= alpha <= 1.0:
+            msg = f"alpha ({alpha}) must be in the range [0.0, 1.0]"
+            raise ValueError(msg)
+
+        documents = await self._document_store._hybrid_retrieval_async(
+            query=query,
+            query_embedding=query_embedding,
+            filters=filters,
+            top_k=top_k,
+            alpha=alpha,
+            max_vector_distance=max_vector_distance,
+        )
+        return {"documents": documents}