Skip to content

Commit

Permalink
Feat: Add filters to run function in retrievers of elasticsearch (#440)
Browse files Browse the repository at this point in the history
* feat: add filters to run function of bm_25 retriever in elastic search

* feat: add filters to run function of embedding retriever in elastic search

* docs: add docstring for filters in run function of retrievers in elasticsearch
  • Loading branch information
sebastian-weisshaar authored Feb 19, 2024
1 parent df0090f commit 9f4e1ec
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -89,17 +89,18 @@ 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, top_k: Optional[int] = None):
def run(self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None):
"""
Retrieve documents using the BM25 keyword-based algorithm.
:param query: String to search in Documents' text.
:param filters: Filters applied to the retrieved Documents.
: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,
filters=filters or self._filters,
fuzziness=self._fuzziness,
top_k=top_k or self._top_k,
scale_score=self._scale_score,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,17 +63,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], top_k: Optional[int] = None):
def run(self, query_embedding: List[float], filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None):
"""
Retrieve documents using a vector similarity metric.
:param query_embedding: Embedding of the query.
:param filters: Filters applied to the retrieved Documents.
: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,
filters=filters or self._filters,
top_k=top_k or self._top_k,
num_candidates=self._num_candidates,
)
Expand Down

0 comments on commit 9f4e1ec

Please sign in to comment.