diff --git a/integrations/elasticsearch/src/elasticsearch_haystack/bm25_retriever.py b/integrations/elasticsearch/src/elasticsearch_haystack/bm25_retriever.py index df85cf1a1..41be4f351 100644 --- a/integrations/elasticsearch/src/elasticsearch_haystack/bm25_retriever.py +++ b/integrations/elasticsearch/src/elasticsearch_haystack/bm25_retriever.py @@ -48,12 +48,19 @@ 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: Optional[int] = None): + """ + Retrieve documents using the BM25 keyword-based algorithm. + + :param query: String to search in Documents' text. + :param top_k: Maximum number of Documents to return. + :return: List of Documents that match the query. + """ docs = self._document_store._bm25_retrieval( query=query, filters=self._filters, fuzziness=self._fuzziness, - top_k=self._top_k, + top_k=top_k or self._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..40c455a4f 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: Optional[int] = None): """ Retrieve documents using a vector similarity metric. :param query_embedding: Embedding of the query. - :return: List of Document similar to `query_embedding`. + :param top_k: Maximum number of Documents to return. + :return: List of Documents similar to `query_embedding`. """ docs = self._document_store._embedding_retrieval( query_embedding=query_embedding, filters=self._filters, - top_k=self._top_k, + top_k=top_k or self._top_k, num_candidates=self._num_candidates, ) return {"documents": docs}