From 82d2af3b02b295aa203435310391c3da1672925e Mon Sep 17 00:00:00 2001
From: sahusiddharth <siddharth.sahu@plaksha.edu.in>
Date: Wed, 20 Dec 2023 15:43:03 +0530
Subject: [PATCH] added top_k arguement in the run function of
 ElasticSearcBM25Retriever

---
 .../src/elasticsearch_haystack/bm25_retriever.py             | 4 ++--
 .../src/elasticsearch_haystack/embedding_retriever.py        | 5 +++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/integrations/elasticsearch/src/elasticsearch_haystack/bm25_retriever.py b/integrations/elasticsearch/src/elasticsearch_haystack/bm25_retriever.py
index df85cf1a1..1965f281d 100644
--- a/integrations/elasticsearch/src/elasticsearch_haystack/bm25_retriever.py
+++ b/integrations/elasticsearch/src/elasticsearch_haystack/bm25_retriever.py
@@ -48,12 +48,12 @@ def from_dict(cls, data: Dict[str, Any]) -> "ElasticsearchBM25Retriever":
         return default_from_dict(cls, data)
 
     @component.output_types(documents=List[Document])
-    def run(self, query: str):
+    def run(self, query: str, top_k: int=None):
         docs = self._document_store._bm25_retrieval(
             query=query,
             filters=self._filters,
             fuzziness=self._fuzziness,
-            top_k=self._top_k,
+            top_k=self._top_k if top_k == None else top_k,
             scale_score=self._scale_score,
         )
         return {"documents": docs}
diff --git a/integrations/elasticsearch/src/elasticsearch_haystack/embedding_retriever.py b/integrations/elasticsearch/src/elasticsearch_haystack/embedding_retriever.py
index 2aaba382d..50faa07cf 100644
--- a/integrations/elasticsearch/src/elasticsearch_haystack/embedding_retriever.py
+++ b/integrations/elasticsearch/src/elasticsearch_haystack/embedding_retriever.py
@@ -64,17 +64,18 @@ def from_dict(cls, data: Dict[str, Any]) -> "ElasticsearchEmbeddingRetriever":
         return default_from_dict(cls, data)
 
     @component.output_types(documents=List[Document])
-    def run(self, query_embedding: List[float]):
+    def run(self, query_embedding: List[float], top_k:int = None):
         """
         Retrieve documents using a vector similarity metric.
 
         :param query_embedding: Embedding of the query.
+        :param top_k: Maximum number of Documents to return
         :return: List of Document similar to `query_embedding`.
         """
         docs = self._document_store._embedding_retrieval(
             query_embedding=query_embedding,
             filters=self._filters,
-            top_k=self._top_k,
+            top_k=self._top_k if top_k == None else top_k,
             num_candidates=self._num_candidates,
         )
         return {"documents": docs}