# 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.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)