# SPDX-FileCopyrightText: 2022-present deepset GmbH # # 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()