# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 from typing import Any from haystack import Document, component, default_from_dict, default_to_dict from haystack.document_stores.in_memory import InMemoryDocumentStore from haystack.document_stores.types import FilterPolicy @component class InMemoryBM25Retriever: """ Retrieves documents that are most similar to the query using keyword-based algorithm. Use this retriever with the InMemoryDocumentStore. ### Usage example ```python from haystack import Document from haystack.components.retrievers.in_memory import InMemoryBM25Retriever from haystack.document_stores.in_memory import InMemoryDocumentStore docs = [ Document(content="Python is a popular programming language"), Document(content="python ist eine beliebte Programmiersprache"), ] doc_store = InMemoryDocumentStore() doc_store.write_documents(docs) retriever = InMemoryBM25Retriever(doc_store) result = retriever.run(query="Programmiersprache") print(result["documents"]) ``` """ def __init__( self, document_store: InMemoryDocumentStore, filters: dict[str, Any] | None = None, top_k: int = 10, scale_score: bool = False, filter_policy: FilterPolicy = FilterPolicy.REPLACE, ) -> None: """ Create the InMemoryBM25Retriever component. :param document_store: An instance of InMemoryDocumentStore where the retriever should search for relevant documents. :param filters: A dictionary with filters to narrow down the retriever's search space in the document store. :param top_k: The maximum number of documents to retrieve. :param scale_score: When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant. When `False`, uses raw similarity scores. :param filter_policy: The filter policy to apply during retrieval. Filter policy determines how filters are applied when retrieving documents. You can choose: - `REPLACE` (default): Overrides the initialization filters with the filters specified at runtime. Use this policy to dynamically change filtering for specific queries. - `MERGE`: Combines runtime filters with initialization filters to narrow down the search. :raises TypeError: If the document_store is not an instance of InMemoryDocumentStore. :raises ValueError: If the specified `top_k` is not > 0. """ if not isinstance(document_store, InMemoryDocumentStore): raise TypeError("document_store must be an instance of InMemoryDocumentStore") self.document_store = document_store if top_k <= 0: raise ValueError(f"top_k must be greater than 0. Currently, the top_k is {top_k}") self.filters = filters self.top_k = top_k self.scale_score = scale_score self.filter_policy = filter_policy def _get_telemetry_data(self) -> dict[str, Any]: """ Data that is sent to Posthog for usage analytics. """ return {"document_store": type(self.document_store).__name__} def to_dict(self) -> dict[str, Any]: """ Serializes the component to a dictionary. :returns: Dictionary with serialized data. """ return default_to_dict( self, document_store=self.document_store, filters=self.filters, top_k=self.top_k, scale_score=self.scale_score, filter_policy=self.filter_policy.value, ) @classmethod def from_dict(cls, data: dict[str, Any]) -> "InMemoryBM25Retriever": """ Deserializes the component from a dictionary. :param data: The dictionary to deserialize from. :returns: The deserialized component. """ init_params = data.get("init_parameters", {}) if "filter_policy" in init_params: init_params["filter_policy"] = FilterPolicy.from_str(init_params["filter_policy"]) return default_from_dict(cls, data) @component.output_types(documents=list[Document]) def run( self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None, scale_score: bool | None = None, ) -> dict[str, list[Document]]: """ Run the InMemoryBM25Retriever on the given input data. :param query: The query string for the Retriever. :param filters: A dictionary with filters to narrow down the search space when retrieving documents. :param top_k: The maximum number of documents to return. :param scale_score: When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant. When `False`, uses raw similarity scores. :returns: The retrieved documents. :raises ValueError: If the specified DocumentStore is not found or is not a InMemoryDocumentStore instance. """ if self.filter_policy == FilterPolicy.MERGE and filters: filters = {**(self.filters or {}), **filters} else: filters = filters or self.filters if top_k is None: top_k = self.top_k if scale_score is None: scale_score = self.scale_score docs = self.document_store.bm25_retrieval(query=query, filters=filters, top_k=top_k, scale_score=scale_score) return {"documents": docs} @component.output_types(documents=list[Document]) async def run_async( self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None, scale_score: bool | None = None, ) -> dict[str, list[Document]]: """ Run the InMemoryBM25Retriever on the given input data. :param query: The query string for the Retriever. :param filters: A dictionary with filters to narrow down the search space when retrieving documents. :param top_k: The maximum number of documents to return. :param scale_score: When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant. When `False`, uses raw similarity scores. :returns: The retrieved documents. :raises ValueError: If the specified DocumentStore is not found or is not a InMemoryDocumentStore instance. """ if self.filter_policy == FilterPolicy.MERGE and filters: filters = {**(self.filters or {}), **filters} else: filters = filters or self.filters if top_k is None: top_k = self.top_k if scale_score is None: scale_score = self.scale_score docs = await self.document_store.bm25_retrieval_async( query=query, filters=filters, top_k=top_k, scale_score=scale_score ) return {"documents": docs}