deepset-ai--haystack/test/components/retrievers/test_text_embedding_retriever.py

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

import os
from typing import Any
from unittest.mock import ANY, AsyncMock, Mock

import numpy as np
import pytest

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


@component
class MockTextEmbedder:
    @component.output_types(embedding=list[float])
    def run(self, text: str) -> dict[str, list[float]]:
        return {"embedding": np.ones(384).tolist()}


class TestTextEmbeddingRetriever:
    @pytest.fixture
    def sample_documents(self):
        return [
            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="Fossil fuels, such as coal, oil, and natural gas, are non-renewable energy sources."),
        ]

    @pytest.fixture
    def document_store_with_embeddings(self, sample_documents):
        """Create a document store populated with embedded documents."""
        document_store = InMemoryDocumentStore()
        doc_embedder = OpenAIDocumentEmbedder()
        doc_writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP)

        embedded_docs = doc_embedder.run(sample_documents)["documents"]
        doc_writer.run(documents=embedded_docs)
        return document_store

    def test_init(self):
        embedding_retriever = InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore())
        text_embedder = MockTextEmbedder()
        retriever = TextEmbeddingRetriever(retriever=embedding_retriever, text_embedder=text_embedder)
        assert retriever.retriever == embedding_retriever
        assert retriever.text_embedder == text_embedder

    def test_run_with_empty_document_store(self):
        retriever = TextEmbeddingRetriever(
            retriever=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()),
            text_embedder=MockTextEmbedder(),
        )
        result = retriever.run(query="green energy")
        assert "documents" in result
        assert result["documents"] == []

    def test_run_returns_documents_sorted_by_score(self):
        doc_high = Document(content="Solar energy", id="doc1", score=0.9)
        doc_low = Document(content="Fossil fuels", id="doc2", score=0.3)
        doc_mid = Document(content="Wind energy", id="doc3", score=0.6)

        @component
        class MockRetriever:
            @component.output_types(documents=list[Document])
            def run(
                self, query_embedding: list[float], filters: dict[str, Any] | None = None, top_k: int | None = None
            ) -> dict[str, list[Document]]:
                return {"documents": [doc_low, doc_high, doc_mid]}

        retriever = TextEmbeddingRetriever(retriever=MockRetriever(), text_embedder=MockTextEmbedder())
        result = retriever.run(query="energy")

        scores = [doc.score for doc in result["documents"]]
        assert scores == sorted(scores, reverse=True)

    def test_to_dict(self):
        retriever = TextEmbeddingRetriever(
            retriever=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()),
            text_embedder=MockTextEmbedder(),
        )
        result = retriever.to_dict()
        assert result == {
            "type": "haystack.components.retrievers.text_embedding_retriever.TextEmbeddingRetriever",
            "init_parameters": {
                "retriever": {
                    "type": "haystack.components.retrievers.in_memory.embedding_retriever.InMemoryEmbeddingRetriever",
                    "init_parameters": {
                        "document_store": {
                            "type": "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore",
                            "init_parameters": {
                                "bm25_tokenization_regex": "(?u)\\b\\w+\\b",
                                "bm25_algorithm": "BM25L",
                                "bm25_parameters": {},
                                "embedding_similarity_function": "dot_product",
                                "index": ANY,
                                "shared": True,
                                "return_embedding": True,
                            },
                        },
                        "filters": None,
                        "top_k": 10,
                        "scale_score": False,
                        "return_embedding": False,
                        "filter_policy": "replace",
                    },
                },
                "text_embedder": {
                    "type": "retrievers.test_text_embedding_retriever.MockTextEmbedder",
                    "init_parameters": {},
                },
            },
        }

    def test_from_dict(self, monkeypatch):
        monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
        data = {
            "type": "haystack.components.retrievers.text_embedding_retriever.TextEmbeddingRetriever",
            "init_parameters": {
                "retriever": {
                    "type": "haystack.components.retrievers.in_memory.embedding_retriever.InMemoryEmbeddingRetriever",
                    "init_parameters": {
                        "document_store": {
                            "type": "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore",
                            "init_parameters": {
                                "bm25_tokenization_regex": "(?u)\\b\\w\\w+\\b",
                                "bm25_algorithm": "BM25L",
                                "bm25_parameters": {},
                                "embedding_similarity_function": "dot_product",
                                "index": "4bb5369d-779f-487b-9c16-3c40f503438b",
                                "shared": True,
                                "return_embedding": True,
                            },
                        },
                        "filters": None,
                        "top_k": 10,
                        "scale_score": False,
                        "return_embedding": False,
                        "filter_policy": "replace",
                    },
                },
                "text_embedder": {
                    "type": "haystack.components.embedders.openai_text_embedder.OpenAITextEmbedder",
                    "init_parameters": {
                        "api_key": {"env_vars": ["OPENAI_API_KEY"], "strict": True, "type": "env_var"},
                        "api_base_url": None,
                        "dimensions": None,
                        "model": "text-embedding-ada-002",
                        "organization": None,
                        "http_client_kwargs": None,
                        "prefix": "",
                        "suffix": "",
                        "timeout": None,
                        "max_retries": None,
                    },
                },
            },
        }
        result = TextEmbeddingRetriever.from_dict(data)
        assert isinstance(result, TextEmbeddingRetriever)
        assert isinstance(result.retriever, InMemoryEmbeddingRetriever)
        assert isinstance(result.text_embedder, OpenAITextEmbedder)

    @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
    @pytest.mark.integration
    def test_run_with_filters(self, document_store_with_embeddings):
        retriever = TextEmbeddingRetriever(
            retriever=InMemoryEmbeddingRetriever(document_store=document_store_with_embeddings),
            text_embedder=OpenAITextEmbedder(),
        )
        result = retriever.run(query="energy", filters={"field": "meta.category", "operator": "==", "value": "solar"})
        assert "documents" in result
        assert all(doc.meta.get("category") == "solar" for doc in result["documents"])

    @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
    @pytest.mark.integration
    def test_run_with_top_k(self, document_store_with_embeddings):
        retriever = TextEmbeddingRetriever(
            retriever=InMemoryEmbeddingRetriever(document_store=document_store_with_embeddings),
            text_embedder=OpenAITextEmbedder(),
        )
        result = retriever.run(query="energy", top_k=2)
        assert "documents" in result
        assert len(result["documents"]) <= 2


class TestComponentLifecycle:
    def test_warm_up_delegates_to_inner_components(self):
        text_embedder = Mock(spec=["run", "warm_up"])
        retriever = Mock(spec=["run", "warm_up"])
        component = TextEmbeddingRetriever(retriever=retriever, text_embedder=text_embedder)
        component.warm_up()
        text_embedder.warm_up.assert_called_once()
        retriever.warm_up.assert_called_once()

    async def test_warm_up_async_delegates_to_inner_components(self):
        text_embedder = Mock(spec=["run", "warm_up_async"])
        text_embedder.warm_up_async = AsyncMock()
        retriever = Mock(spec=["run", "warm_up_async"])
        retriever.warm_up_async = AsyncMock()
        component = TextEmbeddingRetriever(retriever=retriever, text_embedder=text_embedder)
        await component.warm_up_async()
        text_embedder.warm_up_async.assert_awaited_once()
        retriever.warm_up_async.assert_awaited_once()

    async def test_warm_up_async_falls_back_to_sync_warm_up(self):
        text_embedder = Mock(spec=["run", "warm_up"])
        retriever = Mock(spec=["run", "warm_up"])
        component = TextEmbeddingRetriever(retriever=retriever, text_embedder=text_embedder)
        await component.warm_up_async()
        text_embedder.warm_up.assert_called_once()
        retriever.warm_up.assert_called_once()

    def test_close_delegates_to_inner_components(self):
        text_embedder = Mock(spec=["run", "close"])
        retriever = Mock(spec=["run", "close"])
        component = TextEmbeddingRetriever(retriever=retriever, text_embedder=text_embedder)
        component.close()
        text_embedder.close.assert_called_once()
        retriever.close.assert_called_once()

    async def test_close_async_delegates_to_inner_components(self):
        text_embedder = Mock(spec=["run", "close_async"])
        text_embedder.close_async = AsyncMock()
        retriever = Mock(spec=["run", "close_async"])
        retriever.close_async = AsyncMock()
        component = TextEmbeddingRetriever(retriever=retriever, text_embedder=text_embedder)
        await component.close_async()
        text_embedder.close_async.assert_awaited_once()
        retriever.close_async.assert_awaited_once()

    async def test_close_async_falls_back_to_sync_close(self):
        text_embedder = Mock(spec=["run", "close"])
        retriever = Mock(spec=["run", "close"])
        component = TextEmbeddingRetriever(retriever=retriever, text_embedder=text_embedder)
        await component.close_async()
        text_embedder.close.assert_called_once()
        retriever.close.assert_called_once()

    async def test_lifecycle_is_safe_when_inner_components_lack_methods(self):
        text_embedder = Mock(spec=["run"])
        retriever = Mock(spec=["run"])
        component = TextEmbeddingRetriever(retriever=retriever, text_embedder=text_embedder)
        component.warm_up()
        await component.warm_up_async()
        component.close()
        await component.close_async()