feat: support additional search_kwargs with OpenSearchEmbeddingRetriever (#2825)

* feat: support additional  with

* fix lint

* apply feedback

* add support in OpenSearchHybridRetriever as well

* fix tests

* fix tests

* fix tests

* fix tests

* fix tests

* apply feedback

Author: tstadel

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

@@ -34,6 +34,7 @@ def __init__(
         custom_query: dict[str, Any] | None = None,
         raise_on_failure: bool = True,
         efficient_filtering: bool = False,
+        search_kwargs: dict[str, Any] | None = None,
     ):
         """
         Create the OpenSearchEmbeddingRetriever component.
@@ -90,6 +91,18 @@ def __init__(
             If `False`, logs a warning and returns an empty list.
         :param efficient_filtering: If `True`, the filter will be applied during the approximate kNN search.
             This is only supported for knn engines "faiss" and "lucene" and does not work with the default "nmslib".
+        :param search_kwargs: Additional keyword arguments for finetuning the embedding search.
+            E.g., to specify `k` and `ef_search`
+            ```python
+            {
+                "k": 20, # See https://docs.opensearch.org/latest/vector-search/vector-search-techniques/approximate-knn/#the-number-of-returned-results
+                "method_parameters": {
+                    "ef_search": 512, # See https://docs.opensearch.org/latest/query-dsl/specialized/k-nn/index/#ef_search
+                }
+            }
+            ```
+            For a full list of available parameters, see the OpenSearch documentation:
+            https://docs.opensearch.org/latest/query-dsl/specialized/k-nn/index/#request-body-fields
 
         :raises ValueError: If `document_store` is not an instance of OpenSearchDocumentStore.
         """
@@ -106,6 +119,7 @@ def __init__(
         self._custom_query = custom_query
         self._raise_on_failure = raise_on_failure
         self._efficient_filtering = efficient_filtering
+        self._search_kwargs = search_kwargs
 
     def to_dict(self) -> dict[str, Any]:
         """
@@ -123,6 +137,7 @@ def to_dict(self) -> dict[str, Any]:
             custom_query=self._custom_query,
             raise_on_failure=self._raise_on_failure,
             efficient_filtering=self._efficient_filtering,
+            search_kwargs=self._search_kwargs,
         )
 
     @classmethod
@@ -155,6 +170,7 @@ def run(
         custom_query: dict[str, Any] | None = None,
         efficient_filtering: bool | None = None,
         document_store: OpenSearchDocumentStore | None = None,
+        search_kwargs: dict[str, Any] | None = None,
     ) -> dict[str, list[Document]]:
         """
         Retrieve documents using a vector similarity metric.
@@ -208,6 +224,19 @@ def run(
         :param efficient_filtering: If `True`, the filter will be applied during the approximate kNN search.
             This is only supported for knn engines "faiss" and "lucene" and does not work with the default "nmslib".
         :param document_store: Optional instance of OpenSearchDocumentStore to use with the Retriever.
+        :param search_kwargs: Additional keyword arguments for finetuning the embedding search. If not provided,
+            defaults to the parameter set at initialization (if any).
+            E.g., to specify `k` and `ef_search`
+            ```python
+            {
+                "k": 20, # See https://docs.opensearch.org/latest/vector-search/vector-search-techniques/approximate-knn/#the-number-of-returned-results
+                "method_parameters": {
+                    "ef_search": 512, # See https://docs.opensearch.org/latest/query-dsl/specialized/k-nn/index/#ef_search
+                }
+            }
+            ```
+            For a full list of available parameters, see the OpenSearch documentation:
+            https://docs.opensearch.org/latest/query-dsl/specialized/k-nn/index/#request-body-fields
 
         :returns:
             Dictionary with key "documents" containing the retrieved Documents.
@@ -223,6 +252,8 @@ def run(
             custom_query = self._custom_query
         if efficient_filtering is None:
             efficient_filtering = self._efficient_filtering
+        if search_kwargs is None:
+            search_kwargs = self._search_kwargs
 
         docs: list[Document] = []
 
@@ -241,6 +272,7 @@ def run(
                 top_k=top_k,
                 custom_query=custom_query,
                 efficient_filtering=efficient_filtering,
+                search_kwargs=search_kwargs,
             )
         except Exception as e:
             if self._raise_on_failure:
@@ -264,6 +296,7 @@ async def run_async(
         custom_query: dict[str, Any] | None = None,
         efficient_filtering: bool | None = None,
         document_store: OpenSearchDocumentStore | None = None,
+        search_kwargs: dict[str, Any] | None = None,
     ) -> dict[str, list[Document]]:
         """
         Asynchronously retrieve documents using a vector similarity metric.
@@ -317,6 +350,19 @@ async def run_async(
         :param efficient_filtering: If `True`, the filter will be applied during the approximate kNN search.
             This is only supported for knn engines "faiss" and "lucene" and does not work with the default "nmslib".
         :param document_store: Optional instance of OpenSearchDocumentStore to use with the Retriever.
+        :param search_kwargs: Additional keyword arguments for finetuning the embedding search. If not provided,
+            defaults to the parameter set at initialization (if any).
+            E.g., to specify `k` and `ef_search`
+            ```python
+            {
+                "k": 20, # See https://docs.opensearch.org/latest/vector-search/vector-search-techniques/approximate-knn/#the-number-of-returned-results
+                "method_parameters": {
+                    "ef_search": 512, # See https://docs.opensearch.org/latest/query-dsl/specialized/k-nn/index/#ef_search
+                }
+            }
+            ```
+            For a full list of available parameters, see the OpenSearch documentation:
+            https://docs.opensearch.org/latest/query-dsl/specialized/k-nn/index/#request-body-fields
 
         :returns:
             Dictionary with key "documents" containing the retrieved Documents.
@@ -332,6 +378,8 @@ async def run_async(
             custom_query = self._custom_query
         if efficient_filtering is None:
             efficient_filtering = self._efficient_filtering
+        if search_kwargs is None:
+            search_kwargs = self._search_kwargs
 
         docs: list[Document] = []
 
@@ -350,6 +398,7 @@ async def run_async(
                 top_k=top_k,
                 custom_query=custom_query,
                 efficient_filtering=efficient_filtering,
+                search_kwargs=search_kwargs,
             )
         except Exception as e:
             if self._raise_on_failure:

integrations/opensearch/src/haystack_integrations/components/retrievers/opensearch/open_search_hybrid_retriever.py

@@ -104,6 +104,7 @@ def __init__(
         top_k_embedding: int = 10,
         filter_policy_embedding: str | FilterPolicy = FilterPolicy.REPLACE,
         custom_query_embedding: dict[str, Any] | None = None,
+        search_kwargs_embedding: dict[str, Any] | None = None,
         # DocumentJoiner
         join_mode: str | JoinMode = JoinMode.RECIPROCAL_RANK_FUSION,
         weights: list[float] | None = None,
@@ -153,6 +154,8 @@ def __init__(
             The filter policy for the embedding retriever.
         :param custom_query_embedding:
             A custom query for the embedding retriever.
+        :param search_kwargs_embedding:
+            Additional search kwargs for the embedding retriever.
         :param join_mode:
             The mode to use for joining the results from the BM25 and embedding retrievers.
         :param weights:
@@ -185,6 +188,7 @@ def __init__(
         self.top_k_embedding = top_k_embedding
         self.filter_policy_embedding = filter_policy_embedding
         self.custom_query_embedding = custom_query_embedding
+        self.search_kwargs_embedding = search_kwargs_embedding
 
         # DocumentJoiner
         self.join_mode = join_mode
@@ -209,6 +213,7 @@ def __init__(
                 "top_k": self.top_k_embedding,
                 "filter_policy": self.filter_policy_embedding,
                 "custom_query": self.custom_query_embedding,
+                "search_kwargs": self.search_kwargs_embedding,
             },
             "document_joiner": {
                 "join_mode": self.join_mode,
@@ -311,6 +316,7 @@ def to_dict(self):
                 else self.filter_policy_embedding
             ),
             custom_query_embedding=self.custom_query_embedding,
+            search_kwargs_embedding=self.search_kwargs_embedding,
             # DocumentJoiner
             join_mode=(self.join_mode.value if isinstance(self.join_mode, JoinMode) else self.join_mode),
             weights=self.weights,

integrations/opensearch/src/haystack_integrations/document_stores/opensearch/document_store.py

@@ -1409,6 +1409,7 @@ def _prepare_embedding_search_request(
         top_k: int,
         custom_query: dict[str, Any] | None,
         efficient_filtering: bool = False,
+        search_kwargs: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         if not query_embedding:
             msg = "query_embedding must be a non-empty list of floats"
@@ -1434,6 +1435,7 @@ def _prepare_embedding_search_request(
                                     "embedding": {
                                         "vector": query_embedding,
                                         "k": top_k,
+                                        **(search_kwargs or {}),
                                     }
                                 }
                             }
@@ -1465,6 +1467,7 @@ def _embedding_retrieval(
         top_k: int = 10,
         custom_query: dict[str, Any] | None = None,
         efficient_filtering: bool = False,
+        search_kwargs: dict[str, Any] | None = None,
     ) -> list[Document]:
         """
         Retrieves documents that are most similar to the query embedding using a vector similarity metric.
@@ -1484,6 +1487,7 @@ def _embedding_retrieval(
             top_k=top_k,
             custom_query=custom_query,
             efficient_filtering=efficient_filtering,
+            search_kwargs=search_kwargs,
         )
         return self._search_documents(search_params)
 
@@ -1495,6 +1499,7 @@ async def _embedding_retrieval_async(
         top_k: int = 10,
         custom_query: dict[str, Any] | None = None,
         efficient_filtering: bool = False,
+        search_kwargs: dict[str, Any] | None = None,
     ) -> list[Document]:
         """
         Asynchronously retrieves documents that are most similar to the query embedding using a vector similarity
@@ -1515,6 +1520,7 @@ async def _embedding_retrieval_async(
             top_k=top_k,
             custom_query=custom_query,
             efficient_filtering=efficient_filtering,
+            search_kwargs=search_kwargs,
         )
         return await self._search_documents_async(search_params)
 

integrations/opensearch/tests/test_embedding_retriever.py

@@ -85,6 +85,7 @@ def test_to_dict(_mock_opensearch_client):
             "custom_query": {"some": "custom query"},
             "raise_on_failure": True,
             "efficient_filtering": False,
+            "search_kwargs": None,
         },
     }
 
@@ -145,6 +146,7 @@ def test_run():
         top_k=10,
         custom_query=None,
         efficient_filtering=False,
+        search_kwargs=None,
     )
     assert len(res) == 1
     assert len(res["documents"]) == 1
@@ -164,6 +166,7 @@ async def test_run_async():
         top_k=10,
         custom_query=None,
         efficient_filtering=False,
+        search_kwargs=None,
     )
     assert len(res) == 1
     assert len(res["documents"]) == 1
@@ -180,6 +183,7 @@ def test_run_init_params():
         top_k=11,
         custom_query="custom_query",
         efficient_filtering=True,
+        search_kwargs={"k": 10},
     )
     res = retriever.run(query_embedding=[0.5, 0.7])
     mock_store._embedding_retrieval.assert_called_once_with(
@@ -188,6 +192,7 @@ def test_run_init_params():
         top_k=11,
         custom_query="custom_query",
         efficient_filtering=True,
+        search_kwargs={"k": 10},
     )
     assert len(res) == 1
     assert len(res["documents"]) == 1
@@ -204,6 +209,7 @@ async def test_run_async_init_params():
         filters={"from": "init"},
         top_k=11,
         custom_query="custom_query",
+        search_kwargs={"k": 10},
     )
     res = await retriever.run_async(query_embedding=[0.5, 0.7])
     mock_store._embedding_retrieval_async.assert_called_once_with(
@@ -212,6 +218,7 @@ async def test_run_async_init_params():
         top_k=11,
         custom_query="custom_query",
         efficient_filtering=False,
+        search_kwargs={"k": 10},
     )
     assert len(res) == 1
     assert len(res["documents"]) == 1
@@ -222,14 +229,19 @@ async def test_run_async_init_params():
 def test_run_time_params():
     mock_store = Mock(spec=OpenSearchDocumentStore)
     mock_store._embedding_retrieval.return_value = [Document(content="Test doc", embedding=[0.1, 0.2])]
-    retriever = OpenSearchEmbeddingRetriever(document_store=mock_store, filters={"from": "init"}, top_k=11)
-    res = retriever.run(query_embedding=[0.5, 0.7], filters={"from": "run"}, top_k=9, efficient_filtering=True)
+    retriever = OpenSearchEmbeddingRetriever(
+        document_store=mock_store, filters={"from": "init"}, top_k=11, search_kwargs={"k": 10}
+    )
+    res = retriever.run(
+        query_embedding=[0.5, 0.7], filters={"from": "run"}, top_k=9, efficient_filtering=True, search_kwargs={"k": 9}
+    )
     mock_store._embedding_retrieval.assert_called_once_with(
         query_embedding=[0.5, 0.7],
         filters={"from": "run"},
         top_k=9,
         custom_query=None,
         efficient_filtering=True,
+        search_kwargs={"k": 9},
     )
     assert len(res) == 1
     assert len(res["documents"]) == 1
@@ -249,6 +261,7 @@ async def test_run_async_time_params():
         top_k=9,
         custom_query=None,
         efficient_filtering=False,
+        search_kwargs=None,
     )
     assert len(res) == 1
     assert len(res["documents"]) == 1
@@ -288,6 +301,7 @@ def test_run_with_runtime_document_store():
         top_k=10,
         custom_query=None,
         efficient_filtering=False,
+        search_kwargs=None,
     )
     initial_store._embedding_retrieval.assert_not_called()
 
@@ -325,6 +339,7 @@ async def test_run_async_with_runtime_document_store():
         top_k=10,
         custom_query=None,
         efficient_filtering=False,
+        search_kwargs=None,
     )
     initial_store._embedding_retrieval_async.assert_not_called()
 

integrations/opensearch/tests/test_open_search_hybrid_retriever.py

@@ -88,6 +88,7 @@ class TestOpenSearchHybridRetriever:
             "weights": None,
             "top_k": None,
             "sort_by_score": True,
+            "search_kwargs_embedding": None,
         },
     }
 
@@ -224,6 +225,7 @@ def test_run_with_extra_runtime_params(self, mock_embedder):
             top_k=1,
             custom_query=None,
             efficient_filtering=False,
+            search_kwargs=None,
         )
 
     def test_run_in_pipeline(self, mock_embedder):
@@ -256,4 +258,5 @@ def test_run_in_pipeline(self, mock_embedder):
             top_k=10,
             custom_query=None,
             efficient_filtering=False,
+            search_kwargs=None,
         )