deepset-ai--haystack/haystack/components/retrievers/text_embedding_retriever.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

from typing import Any

from haystack import Document, component, default_from_dict, default_to_dict
from haystack.components.embedders.types.protocol import TextEmbedder
from haystack.components.retrievers.types import EmbeddingRetriever
from haystack.core.serialization import component_to_dict
from haystack.utils.async_utils import _execute_component_async


@component
class TextEmbeddingRetriever:
    """
    A component that retrieves documents using a query with an embedding-based retriever.

    This component takes a text query, converts it to an embedding using a text embedder, and then uses an
    embedding-based retriever to find relevant documents.
    The results are sorted by relevance score.

    ### Usage example

    ```python
    from haystack import Document
    from haystack.document_stores.in_memory import InMemoryDocumentStore
    from haystack.document_stores.types import DuplicatePolicy
    from haystack.components.embedders import OpenAITextEmbedder, OpenAIDocumentEmbedder
    from haystack.components.retrievers import InMemoryEmbeddingRetriever, TextEmbeddingRetriever
    from haystack.components.writers import DocumentWriter

    documents = [
        Document(content="Renewable energy is energy that is collected from renewable resources."),
        Document(content="Solar energy is a type of green energy that is harnessed from the sun."),
        Document(content="Wind energy is another type of green energy that is generated by wind turbines."),
        Document(content="Geothermal energy is heat that comes from the sub-surface of the earth."),
        Document(content="Biomass energy is produced from organic materials, such as plant and animal waste."),
        Document(content="Fossil fuels, such as coal, oil, and natural gas, are non-renewable energy sources."),
    ]

    # Populate the document store
    doc_store = InMemoryDocumentStore()
    doc_embedder = OpenAIDocumentEmbedder()
    doc_writer = DocumentWriter(document_store=doc_store, policy=DuplicatePolicy.SKIP)
    documents = doc_embedder.run(documents)["documents"]
    doc_writer.run(documents=documents)

    # Run the retriever
    in_memory_retriever = InMemoryEmbeddingRetriever(document_store=doc_store, top_k=1)
    text_embedder = OpenAITextEmbedder()
    retriever = TextEmbeddingRetriever(retriever=in_memory_retriever, text_embedder=text_embedder)
    result = retriever.run(query="Geothermal energy")

    for doc in result["documents"]:
        print(f"Content: {doc.content}, Score: {doc.score}")
    # >> Content: Geothermal energy is heat that comes from the sub-surface of the earth., Score: 0.8509603046266574
    ```
    """

    def __init__(self, *, retriever: EmbeddingRetriever, text_embedder: TextEmbedder) -> None:
        """
        Initialize TextEmbeddingRetriever.

        :param retriever: The embedding-based retriever to use for document retrieval.
        :param text_embedder: The text embedder to convert a text query to an embedding.
        """
        self.retriever = retriever
        self.text_embedder = text_embedder

    def warm_up(self) -> None:
        """
        Warm up the text embedder and the retriever.
        """
        for inner in (self.text_embedder, self.retriever):
            if hasattr(inner, "warm_up"):
                inner.warm_up()

    async def warm_up_async(self) -> None:
        """
        Warm up the text embedder and the retriever on the serving event loop.
        """
        for inner in (self.text_embedder, self.retriever):
            if hasattr(inner, "warm_up_async"):
                await inner.warm_up_async()
            elif hasattr(inner, "warm_up"):
                inner.warm_up()

    def close(self) -> None:
        """
        Release the text embedder's and the retriever's resources.
        """
        for inner in (self.text_embedder, self.retriever):
            if hasattr(inner, "close"):
                inner.close()

    async def close_async(self) -> None:
        """
        Release the text embedder's and the retriever's async resources.
        """
        for inner in (self.text_embedder, self.retriever):
            if hasattr(inner, "close_async"):
                await inner.close_async()
            elif hasattr(inner, "close"):
                inner.close()

    @component.output_types(documents=list[Document])
    def run(
        self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None
    ) -> dict[str, list[Document]]:
        """
        Retrieve documents using a single query.

        :param query: The query to retrieve documents for.
        :param filters: A dictionary of filters to apply when retrieving documents.
        :param top_k: The maximum number of documents to return.
        :returns:
            A dictionary containing:
                - `documents`: List of retrieved documents sorted by relevance score.
        """
        self.warm_up()

        embedding_result = self.text_embedder.run(text=query)
        result = self.retriever.run(query_embedding=embedding_result["embedding"], filters=filters, top_k=top_k)
        docs: list[Document] = result["documents"]

        # sort
        docs.sort(key=lambda x: x.score or 0.0, reverse=True)
        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
    ) -> dict[str, list[Document]]:
        """
        Retrieve documents using a single query asynchronously.

        Uses `run_async` on the text embedder and retriever if available, otherwise falls back to
        running `run` in a thread executor.

        :param query: The query to retrieve documents for.
        :param filters: A dictionary of filters to apply when retrieving documents.
        :param top_k: The maximum number of documents to return.
        :returns:
            A dictionary containing:
                - `documents`: List of retrieved documents sorted by relevance score.
        """
        await self.warm_up_async()

        embedding_result = await _execute_component_async(self.text_embedder, text=query)
        result = await _execute_component_async(
            self.retriever, query_embedding=embedding_result["embedding"], filters=filters, top_k=top_k
        )

        docs: list[Document] = result["documents"]
        docs.sort(key=lambda x: x.score or 0.0, reverse=True)
        return {"documents": docs}

    def to_dict(self) -> dict[str, Any]:
        """
        Serializes the component to a dictionary.

        :returns:
            A dictionary representing the serialized component.
        """
        return default_to_dict(
            self,
            retriever=component_to_dict(obj=self.retriever, name="retriever"),
            text_embedder=component_to_dict(obj=self.text_embedder, name="text_embedder"),
        )

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "TextEmbeddingRetriever":
        """
        Deserializes the component from a dictionary.

        :param data: The dictionary to deserialize from.
        :returns:
            The deserialized component.
        """
        return default_from_dict(cls, data)