From 9a5fb71b3266d29850d9c3feb38aaadeec3d8703 Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Fri, 24 Nov 2023 11:52:55 +0100 Subject: [PATCH 01/17] Remove preview folder in test/ --------- Co-authored-by: Silvano Cerza --- test/components/embedders/__init__.py | 0 .../test_openai_document_embedder.py | 288 ++++++++++++++++++ .../embedders/test_openai_text_embedder.py | 118 +++++++ ...sentence_transformers_document_embedder.py | 210 +++++++++++++ ...sentence_transformers_embedding_backend.py | 42 +++ ...est_sentence_transformers_text_embedder.py | 151 +++++++++ 6 files changed, 809 insertions(+) create mode 100644 test/components/embedders/__init__.py create mode 100644 test/components/embedders/test_openai_document_embedder.py create mode 100644 test/components/embedders/test_openai_text_embedder.py create mode 100644 test/components/embedders/test_sentence_transformers_document_embedder.py create mode 100644 test/components/embedders/test_sentence_transformers_embedding_backend.py create mode 100644 test/components/embedders/test_sentence_transformers_text_embedder.py diff --git a/test/components/embedders/__init__.py b/test/components/embedders/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/test/components/embedders/test_openai_document_embedder.py b/test/components/embedders/test_openai_document_embedder.py new file mode 100644 index 0000000000..954846c480 --- /dev/null +++ b/test/components/embedders/test_openai_document_embedder.py @@ -0,0 +1,288 @@ +from unittest.mock import patch +from typing import List, cast + +import pytest +import numpy as np +import openai +from openai.util import convert_to_openai_object +from openai.openai_object import OpenAIObject + +from haystack.preview import Document +from haystack.preview.components.embedders.openai_document_embedder import OpenAIDocumentEmbedder + + +def mock_openai_response(input: List[str], model: str = "text-embedding-ada-002", **kwargs) -> OpenAIObject: + dict_response = { + "object": "list", + "data": [ + {"object": "embedding", "index": i, "embedding": np.random.rand(1536).tolist()} for i in range(len(input)) + ], + "model": model, + "usage": {"prompt_tokens": 4, "total_tokens": 4}, + } + + return cast(OpenAIObject, convert_to_openai_object(dict_response)) + + +class TestOpenAIDocumentEmbedder: + @pytest.mark.unit + def test_init_default(self, monkeypatch): + openai.api_key = None + monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") + embedder = OpenAIDocumentEmbedder() + + assert openai.api_key == "fake-api-key" + + assert embedder.model_name == "text-embedding-ada-002" + assert embedder.organization is None + assert embedder.prefix == "" + assert embedder.suffix == "" + assert embedder.batch_size == 32 + assert embedder.progress_bar is True + assert embedder.metadata_fields_to_embed == [] + assert embedder.embedding_separator == "\n" + + @pytest.mark.unit + def test_init_with_parameters(self): + embedder = OpenAIDocumentEmbedder( + api_key="fake-api-key", + model_name="model", + organization="my-org", + prefix="prefix", + suffix="suffix", + batch_size=64, + progress_bar=False, + metadata_fields_to_embed=["test_field"], + embedding_separator=" | ", + ) + assert openai.api_key == "fake-api-key" + assert openai.organization == "my-org" + + assert embedder.organization == "my-org" + assert embedder.model_name == "model" + assert embedder.prefix == "prefix" + assert embedder.suffix == "suffix" + assert embedder.batch_size == 64 + assert embedder.progress_bar is False + assert embedder.metadata_fields_to_embed == ["test_field"] + assert embedder.embedding_separator == " | " + + @pytest.mark.unit + def test_init_fail_wo_api_key(self, monkeypatch): + openai.api_key = None + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with pytest.raises(ValueError, match="OpenAIDocumentEmbedder expects an OpenAI API key"): + OpenAIDocumentEmbedder() + + @pytest.mark.unit + def test_to_dict(self): + component = OpenAIDocumentEmbedder(api_key="fake-api-key") + data = component.to_dict() + assert data == { + "type": "haystack.preview.components.embedders.openai_document_embedder.OpenAIDocumentEmbedder", + "init_parameters": { + "model_name": "text-embedding-ada-002", + "organization": None, + "prefix": "", + "suffix": "", + "batch_size": 32, + "progress_bar": True, + "metadata_fields_to_embed": [], + "embedding_separator": "\n", + }, + } + + @pytest.mark.unit + def test_to_dict_with_custom_init_parameters(self): + component = OpenAIDocumentEmbedder( + api_key="fake-api-key", + model_name="model", + organization="my-org", + prefix="prefix", + suffix="suffix", + batch_size=64, + progress_bar=False, + metadata_fields_to_embed=["test_field"], + embedding_separator=" | ", + ) + data = component.to_dict() + assert data == { + "type": "haystack.preview.components.embedders.openai_document_embedder.OpenAIDocumentEmbedder", + "init_parameters": { + "model_name": "model", + "organization": "my-org", + "prefix": "prefix", + "suffix": "suffix", + "batch_size": 64, + "progress_bar": False, + "metadata_fields_to_embed": ["test_field"], + "embedding_separator": " | ", + }, + } + + @pytest.mark.unit + def test_prepare_texts_to_embed_w_metadata(self): + documents = [ + Document(content=f"document number {i}:\ncontent", meta={"meta_field": f"meta_value {i}"}) for i in range(5) + ] + + embedder = OpenAIDocumentEmbedder( + api_key="fake-api-key", metadata_fields_to_embed=["meta_field"], embedding_separator=" | " + ) + + prepared_texts = embedder._prepare_texts_to_embed(documents) + + # note that newline is replaced by space + assert prepared_texts == [ + "meta_value 0 | document number 0: content", + "meta_value 1 | document number 1: content", + "meta_value 2 | document number 2: content", + "meta_value 3 | document number 3: content", + "meta_value 4 | document number 4: content", + ] + + @pytest.mark.unit + def test_prepare_texts_to_embed_w_suffix(self): + documents = [Document(content=f"document number {i}") for i in range(5)] + + embedder = OpenAIDocumentEmbedder(api_key="fake-api-key", prefix="my_prefix ", suffix=" my_suffix") + + prepared_texts = embedder._prepare_texts_to_embed(documents) + + assert prepared_texts == [ + "my_prefix document number 0 my_suffix", + "my_prefix document number 1 my_suffix", + "my_prefix document number 2 my_suffix", + "my_prefix document number 3 my_suffix", + "my_prefix document number 4 my_suffix", + ] + + @pytest.mark.unit + def test_embed_batch(self): + texts = ["text 1", "text 2", "text 3", "text 4", "text 5"] + + with patch( + "haystack.preview.components.embedders.openai_document_embedder.openai.Embedding" + ) as openai_embedding_patch: + openai_embedding_patch.create.side_effect = mock_openai_response + embedder = OpenAIDocumentEmbedder(api_key="fake-api-key", model_name="model") + + embeddings, metadata = embedder._embed_batch(texts_to_embed=texts, batch_size=2) + + assert openai_embedding_patch.create.call_count == 3 + + assert isinstance(embeddings, list) + assert len(embeddings) == len(texts) + for embedding in embeddings: + assert isinstance(embedding, list) + assert len(embedding) == 1536 + assert all(isinstance(x, float) for x in embedding) + + # openai.Embedding.create is called 3 times + assert metadata == {"model": "model", "usage": {"prompt_tokens": 3 * 4, "total_tokens": 3 * 4}} + + @pytest.mark.unit + def test_run(self): + docs = [ + Document(content="I love cheese", meta={"topic": "Cuisine"}), + Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}), + ] + + model = "text-similarity-ada-001" + with patch( + "haystack.preview.components.embedders.openai_document_embedder.openai.Embedding" + ) as openai_embedding_patch: + openai_embedding_patch.create.side_effect = mock_openai_response + embedder = OpenAIDocumentEmbedder( + api_key="fake-api-key", + model_name=model, + prefix="prefix ", + suffix=" suffix", + metadata_fields_to_embed=["topic"], + embedding_separator=" | ", + ) + + result = embedder.run(documents=docs) + + openai_embedding_patch.create.assert_called_once_with( + model=model, + input=[ + "prefix Cuisine | I love cheese suffix", + "prefix ML | A transformer is a deep learning architecture suffix", + ], + ) + documents_with_embeddings = result["documents"] + metadata = result["metadata"] + + assert isinstance(documents_with_embeddings, list) + assert len(documents_with_embeddings) == len(docs) + for doc in documents_with_embeddings: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert len(doc.embedding) == 1536 + assert all(isinstance(x, float) for x in doc.embedding) + assert metadata == {"model": model, "usage": {"prompt_tokens": 4, "total_tokens": 4}} + + @pytest.mark.unit + def test_run_custom_batch_size(self): + docs = [ + Document(content="I love cheese", meta={"topic": "Cuisine"}), + Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}), + ] + + model = "text-similarity-ada-001" + with patch( + "haystack.preview.components.embedders.openai_document_embedder.openai.Embedding" + ) as openai_embedding_patch: + openai_embedding_patch.create.side_effect = mock_openai_response + embedder = OpenAIDocumentEmbedder( + api_key="fake-api-key", + model_name=model, + prefix="prefix ", + suffix=" suffix", + metadata_fields_to_embed=["topic"], + embedding_separator=" | ", + batch_size=1, + ) + + result = embedder.run(documents=docs) + + assert openai_embedding_patch.create.call_count == 2 + + documents_with_embeddings = result["documents"] + metadata = result["metadata"] + + assert isinstance(documents_with_embeddings, list) + assert len(documents_with_embeddings) == len(docs) + for doc in documents_with_embeddings: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert len(doc.embedding) == 1536 + assert all(isinstance(x, float) for x in doc.embedding) + + # openai.Embedding.create is called 2 times + assert metadata == {"model": model, "usage": {"prompt_tokens": 2 * 4, "total_tokens": 2 * 4}} + + @pytest.mark.unit + def test_run_wrong_input_format(self): + embedder = OpenAIDocumentEmbedder(api_key="fake-api-key") + + # wrong formats + string_input = "text" + list_integers_input = [1, 2, 3] + + with pytest.raises(TypeError, match="OpenAIDocumentEmbedder expects a list of Documents as input"): + embedder.run(documents=string_input) + + with pytest.raises(TypeError, match="OpenAIDocumentEmbedder expects a list of Documents as input"): + embedder.run(documents=list_integers_input) + + @pytest.mark.unit + def test_run_on_empty_list(self): + embedder = OpenAIDocumentEmbedder(api_key="fake-api-key") + + empty_list_input = [] + result = embedder.run(documents=empty_list_input) + + assert result["documents"] is not None + assert not result["documents"] # empty list diff --git a/test/components/embedders/test_openai_text_embedder.py b/test/components/embedders/test_openai_text_embedder.py new file mode 100644 index 0000000000..50be49ac5d --- /dev/null +++ b/test/components/embedders/test_openai_text_embedder.py @@ -0,0 +1,118 @@ +from unittest.mock import patch +import pytest +import openai +from openai.util import convert_to_openai_object +import numpy as np + +from haystack.preview.components.embedders.openai_text_embedder import OpenAITextEmbedder + + +def mock_openai_response(model: str = "text-embedding-ada-002", **kwargs) -> openai.openai_object.OpenAIObject: + dict_response = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": np.random.rand(1536).tolist()}], + "model": model, + "usage": {"prompt_tokens": 4, "total_tokens": 4}, + } + + return convert_to_openai_object(dict_response) + + +class TestOpenAITextEmbedder: + @pytest.mark.unit + def test_init_default(self, monkeypatch): + openai.api_key = None + monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") + embedder = OpenAITextEmbedder() + + assert openai.api_key == "fake-api-key" + assert embedder.model_name == "text-embedding-ada-002" + assert embedder.organization is None + assert embedder.prefix == "" + assert embedder.suffix == "" + + @pytest.mark.unit + def test_init_with_parameters(self): + embedder = OpenAITextEmbedder( + api_key="fake-api-key", + model_name="model", + organization="fake-organization", + prefix="prefix", + suffix="suffix", + ) + assert openai.api_key == "fake-api-key" + assert embedder.model_name == "model" + assert embedder.organization == "fake-organization" + assert openai.organization == "fake-organization" + assert embedder.prefix == "prefix" + assert embedder.suffix == "suffix" + + @pytest.mark.unit + def test_init_fail_wo_api_key(self, monkeypatch): + openai.api_key = None + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with pytest.raises(ValueError, match="OpenAITextEmbedder expects an OpenAI API key"): + OpenAITextEmbedder() + + @pytest.mark.unit + def test_to_dict(self): + component = OpenAITextEmbedder(api_key="fake-api-key") + data = component.to_dict() + assert data == { + "type": "haystack.preview.components.embedders.openai_text_embedder.OpenAITextEmbedder", + "init_parameters": { + "model_name": "text-embedding-ada-002", + "organization": None, + "prefix": "", + "suffix": "", + }, + } + + @pytest.mark.unit + def test_to_dict_with_custom_init_parameters(self): + component = OpenAITextEmbedder( + api_key="fake-api-key", + model_name="model", + organization="fake-organization", + prefix="prefix", + suffix="suffix", + ) + data = component.to_dict() + assert data == { + "type": "haystack.preview.components.embedders.openai_text_embedder.OpenAITextEmbedder", + "init_parameters": { + "model_name": "model", + "organization": "fake-organization", + "prefix": "prefix", + "suffix": "suffix", + }, + } + + @pytest.mark.unit + def test_run(self): + model = "text-similarity-ada-001" + + with patch( + "haystack.preview.components.embedders.openai_text_embedder.openai.Embedding" + ) as openai_embedding_patch: + openai_embedding_patch.create.side_effect = mock_openai_response + + embedder = OpenAITextEmbedder(api_key="fake-api-key", model_name=model, prefix="prefix ", suffix=" suffix") + result = embedder.run(text="The food was delicious") + + openai_embedding_patch.create.assert_called_once_with( + model=model, input="prefix The food was delicious suffix" + ) + + assert len(result["embedding"]) == 1536 + assert all(isinstance(x, float) for x in result["embedding"]) + assert result["metadata"] == {"model": model, "usage": {"prompt_tokens": 4, "total_tokens": 4}} + + @pytest.mark.unit + def test_run_wrong_input_format(self): + embedder = OpenAITextEmbedder(api_key="fake-api-key") + + list_integers_input = [1, 2, 3] + + with pytest.raises(TypeError, match="OpenAITextEmbedder expects a string as an input"): + embedder.run(text=list_integers_input) diff --git a/test/components/embedders/test_sentence_transformers_document_embedder.py b/test/components/embedders/test_sentence_transformers_document_embedder.py new file mode 100644 index 0000000000..2f5e5e667f --- /dev/null +++ b/test/components/embedders/test_sentence_transformers_document_embedder.py @@ -0,0 +1,210 @@ +from unittest.mock import patch, MagicMock +import pytest +import numpy as np + +from haystack.preview import Document +from haystack.preview.components.embedders.sentence_transformers_document_embedder import ( + SentenceTransformersDocumentEmbedder, +) + + +class TestSentenceTransformersDocumentEmbedder: + @pytest.mark.unit + def test_init_default(self): + embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") + assert embedder.model_name_or_path == "model" + assert embedder.device == "cpu" + assert embedder.token is None + assert embedder.prefix == "" + assert embedder.suffix == "" + assert embedder.batch_size == 32 + assert embedder.progress_bar is True + assert embedder.normalize_embeddings is False + assert embedder.metadata_fields_to_embed == [] + assert embedder.embedding_separator == "\n" + + @pytest.mark.unit + def test_init_with_parameters(self): + embedder = SentenceTransformersDocumentEmbedder( + model_name_or_path="model", + device="cuda", + token=True, + prefix="prefix", + suffix="suffix", + batch_size=64, + progress_bar=False, + normalize_embeddings=True, + metadata_fields_to_embed=["test_field"], + embedding_separator=" | ", + ) + assert embedder.model_name_or_path == "model" + assert embedder.device == "cuda" + assert embedder.token is True + assert embedder.prefix == "prefix" + assert embedder.suffix == "suffix" + assert embedder.batch_size == 64 + assert embedder.progress_bar is False + assert embedder.normalize_embeddings is True + assert embedder.metadata_fields_to_embed == ["test_field"] + assert embedder.embedding_separator == " | " + + @pytest.mark.unit + def test_to_dict(self): + component = SentenceTransformersDocumentEmbedder(model_name_or_path="model") + data = component.to_dict() + assert data == { + "type": "haystack.preview.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", + "init_parameters": { + "model_name_or_path": "model", + "device": "cpu", + "token": None, + "prefix": "", + "suffix": "", + "batch_size": 32, + "progress_bar": True, + "normalize_embeddings": False, + "embedding_separator": "\n", + "metadata_fields_to_embed": [], + }, + } + + @pytest.mark.unit + def test_to_dict_with_custom_init_parameters(self): + component = SentenceTransformersDocumentEmbedder( + model_name_or_path="model", + device="cuda", + token="the-token", + prefix="prefix", + suffix="suffix", + batch_size=64, + progress_bar=False, + normalize_embeddings=True, + metadata_fields_to_embed=["meta_field"], + embedding_separator=" - ", + ) + data = component.to_dict() + + assert data == { + "type": "haystack.preview.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", + "init_parameters": { + "model_name_or_path": "model", + "device": "cuda", + "token": None, # the token is not serialized + "prefix": "prefix", + "suffix": "suffix", + "batch_size": 64, + "progress_bar": False, + "normalize_embeddings": True, + "embedding_separator": " - ", + "metadata_fields_to_embed": ["meta_field"], + }, + } + + @pytest.mark.unit + @patch( + "haystack.preview.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory" + ) + def test_warmup(self, mocked_factory): + embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") + mocked_factory.get_embedding_backend.assert_not_called() + embedder.warm_up() + mocked_factory.get_embedding_backend.assert_called_once_with( + model_name_or_path="model", device="cpu", use_auth_token=None + ) + + @pytest.mark.unit + @patch( + "haystack.preview.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory" + ) + def test_warmup_doesnt_reload(self, mocked_factory): + embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") + mocked_factory.get_embedding_backend.assert_not_called() + embedder.warm_up() + embedder.warm_up() + mocked_factory.get_embedding_backend.assert_called_once() + + @pytest.mark.unit + def test_run(self): + embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") + embedder.embedding_backend = MagicMock() + embedder.embedding_backend.embed = lambda x, **kwargs: np.random.rand(len(x), 16).tolist() + + documents = [Document(content=f"document number {i}") for i in range(5)] + + result = embedder.run(documents=documents) + + assert isinstance(result["documents"], list) + assert len(result["documents"]) == len(documents) + for doc in result["documents"]: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert isinstance(doc.embedding[0], float) + + @pytest.mark.unit + def test_run_wrong_input_format(self): + embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") + + string_input = "text" + list_integers_input = [1, 2, 3] + + with pytest.raises( + TypeError, match="SentenceTransformersDocumentEmbedder expects a list of Documents as input" + ): + embedder.run(documents=string_input) + + with pytest.raises( + TypeError, match="SentenceTransformersDocumentEmbedder expects a list of Documents as input" + ): + embedder.run(documents=list_integers_input) + + @pytest.mark.unit + def test_embed_metadata(self): + embedder = SentenceTransformersDocumentEmbedder( + model_name_or_path="model", metadata_fields_to_embed=["meta_field"], embedding_separator="\n" + ) + embedder.embedding_backend = MagicMock() + + documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)] + + embedder.run(documents=documents) + + embedder.embedding_backend.embed.assert_called_once_with( + [ + "meta_value 0\ndocument number 0", + "meta_value 1\ndocument number 1", + "meta_value 2\ndocument number 2", + "meta_value 3\ndocument number 3", + "meta_value 4\ndocument number 4", + ], + batch_size=32, + show_progress_bar=True, + normalize_embeddings=False, + ) + + @pytest.mark.unit + def test_prefix_suffix(self): + embedder = SentenceTransformersDocumentEmbedder( + model_name_or_path="model", + prefix="my_prefix ", + suffix=" my_suffix", + metadata_fields_to_embed=["meta_field"], + embedding_separator="\n", + ) + embedder.embedding_backend = MagicMock() + + documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)] + + embedder.run(documents=documents) + + embedder.embedding_backend.embed.assert_called_once_with( + [ + "my_prefix meta_value 0\ndocument number 0 my_suffix", + "my_prefix meta_value 1\ndocument number 1 my_suffix", + "my_prefix meta_value 2\ndocument number 2 my_suffix", + "my_prefix meta_value 3\ndocument number 3 my_suffix", + "my_prefix meta_value 4\ndocument number 4 my_suffix", + ], + batch_size=32, + show_progress_bar=True, + normalize_embeddings=False, + ) diff --git a/test/components/embedders/test_sentence_transformers_embedding_backend.py b/test/components/embedders/test_sentence_transformers_embedding_backend.py new file mode 100644 index 0000000000..4ac8c55869 --- /dev/null +++ b/test/components/embedders/test_sentence_transformers_embedding_backend.py @@ -0,0 +1,42 @@ +from unittest.mock import patch +import pytest +from haystack.preview.components.embedders.backends.sentence_transformers_backend import ( + _SentenceTransformersEmbeddingBackendFactory, +) + + +@pytest.mark.unit +@patch("haystack.preview.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") +def test_factory_behavior(mock_sentence_transformer): + embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend( + model_name_or_path="my_model", device="cpu" + ) + same_embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend("my_model", "cpu") + another_embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend( + model_name_or_path="another_model", device="cpu" + ) + + assert same_embedding_backend is embedding_backend + assert another_embedding_backend is not embedding_backend + + +@pytest.mark.unit +@patch("haystack.preview.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") +def test_model_initialization(mock_sentence_transformer): + _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend( + model_name_or_path="model", device="cpu", use_auth_token="my_token" + ) + mock_sentence_transformer.assert_called_once_with( + model_name_or_path="model", device="cpu", use_auth_token="my_token" + ) + + +@pytest.mark.unit +@patch("haystack.preview.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") +def test_embedding_function_with_kwargs(mock_sentence_transformer): + embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(model_name_or_path="model") + + data = ["sentence1", "sentence2"] + embedding_backend.embed(data=data, normalize_embeddings=True) + + embedding_backend.model.encode.assert_called_once_with(data, normalize_embeddings=True) diff --git a/test/components/embedders/test_sentence_transformers_text_embedder.py b/test/components/embedders/test_sentence_transformers_text_embedder.py new file mode 100644 index 0000000000..d93e576ac8 --- /dev/null +++ b/test/components/embedders/test_sentence_transformers_text_embedder.py @@ -0,0 +1,151 @@ +from unittest.mock import patch, MagicMock +import pytest + +import numpy as np + +from haystack.preview.components.embedders.sentence_transformers_text_embedder import SentenceTransformersTextEmbedder + + +class TestSentenceTransformersTextEmbedder: + @pytest.mark.unit + def test_init_default(self): + embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") + assert embedder.model_name_or_path == "model" + assert embedder.device == "cpu" + assert embedder.token is None + assert embedder.prefix == "" + assert embedder.suffix == "" + assert embedder.batch_size == 32 + assert embedder.progress_bar is True + assert embedder.normalize_embeddings is False + + @pytest.mark.unit + def test_init_with_parameters(self): + embedder = SentenceTransformersTextEmbedder( + model_name_or_path="model", + device="cuda", + token=True, + prefix="prefix", + suffix="suffix", + batch_size=64, + progress_bar=False, + normalize_embeddings=True, + ) + assert embedder.model_name_or_path == "model" + assert embedder.device == "cuda" + assert embedder.token is True + assert embedder.prefix == "prefix" + assert embedder.suffix == "suffix" + assert embedder.batch_size == 64 + assert embedder.progress_bar is False + assert embedder.normalize_embeddings is True + + @pytest.mark.unit + def test_to_dict(self): + component = SentenceTransformersTextEmbedder(model_name_or_path="model") + data = component.to_dict() + assert data == { + "type": "haystack.preview.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", + "init_parameters": { + "model_name_or_path": "model", + "device": "cpu", + "token": None, + "prefix": "", + "suffix": "", + "batch_size": 32, + "progress_bar": True, + "normalize_embeddings": False, + }, + } + + @pytest.mark.unit + def test_to_dict_with_custom_init_parameters(self): + component = SentenceTransformersTextEmbedder( + model_name_or_path="model", + device="cuda", + token=True, + prefix="prefix", + suffix="suffix", + batch_size=64, + progress_bar=False, + normalize_embeddings=True, + ) + data = component.to_dict() + assert data == { + "type": "haystack.preview.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", + "init_parameters": { + "model_name_or_path": "model", + "device": "cuda", + "token": True, + "prefix": "prefix", + "suffix": "suffix", + "batch_size": 64, + "progress_bar": False, + "normalize_embeddings": True, + }, + } + + @pytest.mark.unit + def test_to_dict_not_serialize_token(self): + component = SentenceTransformersTextEmbedder(model_name_or_path="model", token="awesome-token") + data = component.to_dict() + assert data == { + "type": "haystack.preview.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", + "init_parameters": { + "model_name_or_path": "model", + "device": "cpu", + "token": None, + "prefix": "", + "suffix": "", + "batch_size": 32, + "progress_bar": True, + "normalize_embeddings": False, + }, + } + + @pytest.mark.unit + @patch( + "haystack.preview.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" + ) + def test_warmup(self, mocked_factory): + embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") + mocked_factory.get_embedding_backend.assert_not_called() + embedder.warm_up() + mocked_factory.get_embedding_backend.assert_called_once_with( + model_name_or_path="model", device="cpu", use_auth_token=None + ) + + @pytest.mark.unit + @patch( + "haystack.preview.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" + ) + def test_warmup_doesnt_reload(self, mocked_factory): + embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") + mocked_factory.get_embedding_backend.assert_not_called() + embedder.warm_up() + embedder.warm_up() + mocked_factory.get_embedding_backend.assert_called_once() + + @pytest.mark.unit + def test_run(self): + embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") + embedder.embedding_backend = MagicMock() + embedder.embedding_backend.embed = lambda x, **kwargs: np.random.rand(len(x), 16).tolist() + + text = "a nice text to embed" + + result = embedder.run(text=text) + embedding = result["embedding"] + + assert isinstance(embedding, list) + assert all(isinstance(el, float) for el in embedding) + + @pytest.mark.unit + def test_run_wrong_input_format(self): + embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") + embedder.embedding_backend = MagicMock() + + list_integers_input = [1, 2, 3] + + with pytest.raises(TypeError, match="SentenceTransformersTextEmbedder expects a string as input"): + embedder.run(text=list_integers_input) From 8f34fae99b50e8b1e36e71a1838bb4e7cfa0f879 Mon Sep 17 00:00:00 2001 From: Silvano Cerza Date: Fri, 24 Nov 2023 14:48:43 +0100 Subject: [PATCH 02/17] Fix all tests --- .../test_openai_document_embedder.py | 20 +++++++------------ .../embedders/test_openai_text_embedder.py | 10 ++++------ ...sentence_transformers_document_embedder.py | 14 ++++++------- ...sentence_transformers_embedding_backend.py | 8 ++++---- ...est_sentence_transformers_text_embedder.py | 12 +++++------ 5 files changed, 27 insertions(+), 37 deletions(-) diff --git a/test/components/embedders/test_openai_document_embedder.py b/test/components/embedders/test_openai_document_embedder.py index 954846c480..c98dd27653 100644 --- a/test/components/embedders/test_openai_document_embedder.py +++ b/test/components/embedders/test_openai_document_embedder.py @@ -7,8 +7,8 @@ from openai.util import convert_to_openai_object from openai.openai_object import OpenAIObject -from haystack.preview import Document -from haystack.preview.components.embedders.openai_document_embedder import OpenAIDocumentEmbedder +from haystack import Document +from haystack.components.embedders.openai_document_embedder import OpenAIDocumentEmbedder def mock_openai_response(input: List[str], model: str = "text-embedding-ada-002", **kwargs) -> OpenAIObject: @@ -79,7 +79,7 @@ def test_to_dict(self): component = OpenAIDocumentEmbedder(api_key="fake-api-key") data = component.to_dict() assert data == { - "type": "haystack.preview.components.embedders.openai_document_embedder.OpenAIDocumentEmbedder", + "type": "haystack.components.embedders.openai_document_embedder.OpenAIDocumentEmbedder", "init_parameters": { "model_name": "text-embedding-ada-002", "organization": None, @@ -107,7 +107,7 @@ def test_to_dict_with_custom_init_parameters(self): ) data = component.to_dict() assert data == { - "type": "haystack.preview.components.embedders.openai_document_embedder.OpenAIDocumentEmbedder", + "type": "haystack.components.embedders.openai_document_embedder.OpenAIDocumentEmbedder", "init_parameters": { "model_name": "model", "organization": "my-org", @@ -161,9 +161,7 @@ def test_prepare_texts_to_embed_w_suffix(self): def test_embed_batch(self): texts = ["text 1", "text 2", "text 3", "text 4", "text 5"] - with patch( - "haystack.preview.components.embedders.openai_document_embedder.openai.Embedding" - ) as openai_embedding_patch: + with patch("haystack.components.embedders.openai_document_embedder.openai.Embedding") as openai_embedding_patch: openai_embedding_patch.create.side_effect = mock_openai_response embedder = OpenAIDocumentEmbedder(api_key="fake-api-key", model_name="model") @@ -189,9 +187,7 @@ def test_run(self): ] model = "text-similarity-ada-001" - with patch( - "haystack.preview.components.embedders.openai_document_embedder.openai.Embedding" - ) as openai_embedding_patch: + with patch("haystack.components.embedders.openai_document_embedder.openai.Embedding") as openai_embedding_patch: openai_embedding_patch.create.side_effect = mock_openai_response embedder = OpenAIDocumentEmbedder( api_key="fake-api-key", @@ -231,9 +227,7 @@ def test_run_custom_batch_size(self): ] model = "text-similarity-ada-001" - with patch( - "haystack.preview.components.embedders.openai_document_embedder.openai.Embedding" - ) as openai_embedding_patch: + with patch("haystack.components.embedders.openai_document_embedder.openai.Embedding") as openai_embedding_patch: openai_embedding_patch.create.side_effect = mock_openai_response embedder = OpenAIDocumentEmbedder( api_key="fake-api-key", diff --git a/test/components/embedders/test_openai_text_embedder.py b/test/components/embedders/test_openai_text_embedder.py index 50be49ac5d..bbe4c8b31f 100644 --- a/test/components/embedders/test_openai_text_embedder.py +++ b/test/components/embedders/test_openai_text_embedder.py @@ -4,7 +4,7 @@ from openai.util import convert_to_openai_object import numpy as np -from haystack.preview.components.embedders.openai_text_embedder import OpenAITextEmbedder +from haystack.components.embedders.openai_text_embedder import OpenAITextEmbedder def mock_openai_response(model: str = "text-embedding-ada-002", **kwargs) -> openai.openai_object.OpenAIObject: @@ -59,7 +59,7 @@ def test_to_dict(self): component = OpenAITextEmbedder(api_key="fake-api-key") data = component.to_dict() assert data == { - "type": "haystack.preview.components.embedders.openai_text_embedder.OpenAITextEmbedder", + "type": "haystack.components.embedders.openai_text_embedder.OpenAITextEmbedder", "init_parameters": { "model_name": "text-embedding-ada-002", "organization": None, @@ -79,7 +79,7 @@ def test_to_dict_with_custom_init_parameters(self): ) data = component.to_dict() assert data == { - "type": "haystack.preview.components.embedders.openai_text_embedder.OpenAITextEmbedder", + "type": "haystack.components.embedders.openai_text_embedder.OpenAITextEmbedder", "init_parameters": { "model_name": "model", "organization": "fake-organization", @@ -92,9 +92,7 @@ def test_to_dict_with_custom_init_parameters(self): def test_run(self): model = "text-similarity-ada-001" - with patch( - "haystack.preview.components.embedders.openai_text_embedder.openai.Embedding" - ) as openai_embedding_patch: + with patch("haystack.components.embedders.openai_text_embedder.openai.Embedding") as openai_embedding_patch: openai_embedding_patch.create.side_effect = mock_openai_response embedder = OpenAITextEmbedder(api_key="fake-api-key", model_name=model, prefix="prefix ", suffix=" suffix") diff --git a/test/components/embedders/test_sentence_transformers_document_embedder.py b/test/components/embedders/test_sentence_transformers_document_embedder.py index 2f5e5e667f..44aba64560 100644 --- a/test/components/embedders/test_sentence_transformers_document_embedder.py +++ b/test/components/embedders/test_sentence_transformers_document_embedder.py @@ -2,10 +2,8 @@ import pytest import numpy as np -from haystack.preview import Document -from haystack.preview.components.embedders.sentence_transformers_document_embedder import ( - SentenceTransformersDocumentEmbedder, -) +from haystack import Document +from haystack.components.embedders.sentence_transformers_document_embedder import SentenceTransformersDocumentEmbedder class TestSentenceTransformersDocumentEmbedder: @@ -53,7 +51,7 @@ def test_to_dict(self): component = SentenceTransformersDocumentEmbedder(model_name_or_path="model") data = component.to_dict() assert data == { - "type": "haystack.preview.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", + "type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", "init_parameters": { "model_name_or_path": "model", "device": "cpu", @@ -85,7 +83,7 @@ def test_to_dict_with_custom_init_parameters(self): data = component.to_dict() assert data == { - "type": "haystack.preview.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", + "type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", "init_parameters": { "model_name_or_path": "model", "device": "cuda", @@ -102,7 +100,7 @@ def test_to_dict_with_custom_init_parameters(self): @pytest.mark.unit @patch( - "haystack.preview.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory" + "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory" ) def test_warmup(self, mocked_factory): embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") @@ -114,7 +112,7 @@ def test_warmup(self, mocked_factory): @pytest.mark.unit @patch( - "haystack.preview.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory" + "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory" ) def test_warmup_doesnt_reload(self, mocked_factory): embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") diff --git a/test/components/embedders/test_sentence_transformers_embedding_backend.py b/test/components/embedders/test_sentence_transformers_embedding_backend.py index 4ac8c55869..9aec7954b3 100644 --- a/test/components/embedders/test_sentence_transformers_embedding_backend.py +++ b/test/components/embedders/test_sentence_transformers_embedding_backend.py @@ -1,12 +1,12 @@ from unittest.mock import patch import pytest -from haystack.preview.components.embedders.backends.sentence_transformers_backend import ( +from haystack.components.embedders.backends.sentence_transformers_backend import ( _SentenceTransformersEmbeddingBackendFactory, ) @pytest.mark.unit -@patch("haystack.preview.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") +@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") def test_factory_behavior(mock_sentence_transformer): embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend( model_name_or_path="my_model", device="cpu" @@ -21,7 +21,7 @@ def test_factory_behavior(mock_sentence_transformer): @pytest.mark.unit -@patch("haystack.preview.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") +@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") def test_model_initialization(mock_sentence_transformer): _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend( model_name_or_path="model", device="cpu", use_auth_token="my_token" @@ -32,7 +32,7 @@ def test_model_initialization(mock_sentence_transformer): @pytest.mark.unit -@patch("haystack.preview.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") +@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") def test_embedding_function_with_kwargs(mock_sentence_transformer): embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(model_name_or_path="model") diff --git a/test/components/embedders/test_sentence_transformers_text_embedder.py b/test/components/embedders/test_sentence_transformers_text_embedder.py index d93e576ac8..cc0f96e830 100644 --- a/test/components/embedders/test_sentence_transformers_text_embedder.py +++ b/test/components/embedders/test_sentence_transformers_text_embedder.py @@ -3,7 +3,7 @@ import numpy as np -from haystack.preview.components.embedders.sentence_transformers_text_embedder import SentenceTransformersTextEmbedder +from haystack.components.embedders.sentence_transformers_text_embedder import SentenceTransformersTextEmbedder class TestSentenceTransformersTextEmbedder: @@ -45,7 +45,7 @@ def test_to_dict(self): component = SentenceTransformersTextEmbedder(model_name_or_path="model") data = component.to_dict() assert data == { - "type": "haystack.preview.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", + "type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", "init_parameters": { "model_name_or_path": "model", "device": "cpu", @@ -72,7 +72,7 @@ def test_to_dict_with_custom_init_parameters(self): ) data = component.to_dict() assert data == { - "type": "haystack.preview.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", + "type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", "init_parameters": { "model_name_or_path": "model", "device": "cuda", @@ -90,7 +90,7 @@ def test_to_dict_not_serialize_token(self): component = SentenceTransformersTextEmbedder(model_name_or_path="model", token="awesome-token") data = component.to_dict() assert data == { - "type": "haystack.preview.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", + "type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", "init_parameters": { "model_name_or_path": "model", "device": "cpu", @@ -105,7 +105,7 @@ def test_to_dict_not_serialize_token(self): @pytest.mark.unit @patch( - "haystack.preview.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" + "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" ) def test_warmup(self, mocked_factory): embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") @@ -117,7 +117,7 @@ def test_warmup(self, mocked_factory): @pytest.mark.unit @patch( - "haystack.preview.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" + "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" ) def test_warmup_doesnt_reload(self, mocked_factory): embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") From a9c028d8cb6e1c4f8581a576de2ae9c27d917791 Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Wed, 29 Nov 2023 19:24:25 +0100 Subject: [PATCH 03/17] remove unit marker (#6450) --- .../embedders/test_openai_document_embedder.py | 12 ------------ .../embedders/test_openai_text_embedder.py | 7 ------- .../test_sentence_transformers_document_embedder.py | 10 ---------- .../test_sentence_transformers_embedding_backend.py | 3 --- .../test_sentence_transformers_text_embedder.py | 9 --------- 5 files changed, 41 deletions(-) diff --git a/test/components/embedders/test_openai_document_embedder.py b/test/components/embedders/test_openai_document_embedder.py index c98dd27653..29b3074a23 100644 --- a/test/components/embedders/test_openai_document_embedder.py +++ b/test/components/embedders/test_openai_document_embedder.py @@ -25,7 +25,6 @@ def mock_openai_response(input: List[str], model: str = "text-embedding-ada-002" class TestOpenAIDocumentEmbedder: - @pytest.mark.unit def test_init_default(self, monkeypatch): openai.api_key = None monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") @@ -42,7 +41,6 @@ def test_init_default(self, monkeypatch): assert embedder.metadata_fields_to_embed == [] assert embedder.embedding_separator == "\n" - @pytest.mark.unit def test_init_with_parameters(self): embedder = OpenAIDocumentEmbedder( api_key="fake-api-key", @@ -67,14 +65,12 @@ def test_init_with_parameters(self): assert embedder.metadata_fields_to_embed == ["test_field"] assert embedder.embedding_separator == " | " - @pytest.mark.unit def test_init_fail_wo_api_key(self, monkeypatch): openai.api_key = None monkeypatch.delenv("OPENAI_API_KEY", raising=False) with pytest.raises(ValueError, match="OpenAIDocumentEmbedder expects an OpenAI API key"): OpenAIDocumentEmbedder() - @pytest.mark.unit def test_to_dict(self): component = OpenAIDocumentEmbedder(api_key="fake-api-key") data = component.to_dict() @@ -92,7 +88,6 @@ def test_to_dict(self): }, } - @pytest.mark.unit def test_to_dict_with_custom_init_parameters(self): component = OpenAIDocumentEmbedder( api_key="fake-api-key", @@ -120,7 +115,6 @@ def test_to_dict_with_custom_init_parameters(self): }, } - @pytest.mark.unit def test_prepare_texts_to_embed_w_metadata(self): documents = [ Document(content=f"document number {i}:\ncontent", meta={"meta_field": f"meta_value {i}"}) for i in range(5) @@ -141,7 +135,6 @@ def test_prepare_texts_to_embed_w_metadata(self): "meta_value 4 | document number 4: content", ] - @pytest.mark.unit def test_prepare_texts_to_embed_w_suffix(self): documents = [Document(content=f"document number {i}") for i in range(5)] @@ -157,7 +150,6 @@ def test_prepare_texts_to_embed_w_suffix(self): "my_prefix document number 4 my_suffix", ] - @pytest.mark.unit def test_embed_batch(self): texts = ["text 1", "text 2", "text 3", "text 4", "text 5"] @@ -179,7 +171,6 @@ def test_embed_batch(self): # openai.Embedding.create is called 3 times assert metadata == {"model": "model", "usage": {"prompt_tokens": 3 * 4, "total_tokens": 3 * 4}} - @pytest.mark.unit def test_run(self): docs = [ Document(content="I love cheese", meta={"topic": "Cuisine"}), @@ -219,7 +210,6 @@ def test_run(self): assert all(isinstance(x, float) for x in doc.embedding) assert metadata == {"model": model, "usage": {"prompt_tokens": 4, "total_tokens": 4}} - @pytest.mark.unit def test_run_custom_batch_size(self): docs = [ Document(content="I love cheese", meta={"topic": "Cuisine"}), @@ -257,7 +247,6 @@ def test_run_custom_batch_size(self): # openai.Embedding.create is called 2 times assert metadata == {"model": model, "usage": {"prompt_tokens": 2 * 4, "total_tokens": 2 * 4}} - @pytest.mark.unit def test_run_wrong_input_format(self): embedder = OpenAIDocumentEmbedder(api_key="fake-api-key") @@ -271,7 +260,6 @@ def test_run_wrong_input_format(self): with pytest.raises(TypeError, match="OpenAIDocumentEmbedder expects a list of Documents as input"): embedder.run(documents=list_integers_input) - @pytest.mark.unit def test_run_on_empty_list(self): embedder = OpenAIDocumentEmbedder(api_key="fake-api-key") diff --git a/test/components/embedders/test_openai_text_embedder.py b/test/components/embedders/test_openai_text_embedder.py index bbe4c8b31f..183fd45a1c 100644 --- a/test/components/embedders/test_openai_text_embedder.py +++ b/test/components/embedders/test_openai_text_embedder.py @@ -19,7 +19,6 @@ def mock_openai_response(model: str = "text-embedding-ada-002", **kwargs) -> ope class TestOpenAITextEmbedder: - @pytest.mark.unit def test_init_default(self, monkeypatch): openai.api_key = None monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") @@ -31,7 +30,6 @@ def test_init_default(self, monkeypatch): assert embedder.prefix == "" assert embedder.suffix == "" - @pytest.mark.unit def test_init_with_parameters(self): embedder = OpenAITextEmbedder( api_key="fake-api-key", @@ -47,14 +45,12 @@ def test_init_with_parameters(self): assert embedder.prefix == "prefix" assert embedder.suffix == "suffix" - @pytest.mark.unit def test_init_fail_wo_api_key(self, monkeypatch): openai.api_key = None monkeypatch.delenv("OPENAI_API_KEY", raising=False) with pytest.raises(ValueError, match="OpenAITextEmbedder expects an OpenAI API key"): OpenAITextEmbedder() - @pytest.mark.unit def test_to_dict(self): component = OpenAITextEmbedder(api_key="fake-api-key") data = component.to_dict() @@ -68,7 +64,6 @@ def test_to_dict(self): }, } - @pytest.mark.unit def test_to_dict_with_custom_init_parameters(self): component = OpenAITextEmbedder( api_key="fake-api-key", @@ -88,7 +83,6 @@ def test_to_dict_with_custom_init_parameters(self): }, } - @pytest.mark.unit def test_run(self): model = "text-similarity-ada-001" @@ -106,7 +100,6 @@ def test_run(self): assert all(isinstance(x, float) for x in result["embedding"]) assert result["metadata"] == {"model": model, "usage": {"prompt_tokens": 4, "total_tokens": 4}} - @pytest.mark.unit def test_run_wrong_input_format(self): embedder = OpenAITextEmbedder(api_key="fake-api-key") diff --git a/test/components/embedders/test_sentence_transformers_document_embedder.py b/test/components/embedders/test_sentence_transformers_document_embedder.py index 44aba64560..66c2304d9e 100644 --- a/test/components/embedders/test_sentence_transformers_document_embedder.py +++ b/test/components/embedders/test_sentence_transformers_document_embedder.py @@ -7,7 +7,6 @@ class TestSentenceTransformersDocumentEmbedder: - @pytest.mark.unit def test_init_default(self): embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") assert embedder.model_name_or_path == "model" @@ -21,7 +20,6 @@ def test_init_default(self): assert embedder.metadata_fields_to_embed == [] assert embedder.embedding_separator == "\n" - @pytest.mark.unit def test_init_with_parameters(self): embedder = SentenceTransformersDocumentEmbedder( model_name_or_path="model", @@ -46,7 +44,6 @@ def test_init_with_parameters(self): assert embedder.metadata_fields_to_embed == ["test_field"] assert embedder.embedding_separator == " | " - @pytest.mark.unit def test_to_dict(self): component = SentenceTransformersDocumentEmbedder(model_name_or_path="model") data = component.to_dict() @@ -66,7 +63,6 @@ def test_to_dict(self): }, } - @pytest.mark.unit def test_to_dict_with_custom_init_parameters(self): component = SentenceTransformersDocumentEmbedder( model_name_or_path="model", @@ -98,7 +94,6 @@ def test_to_dict_with_custom_init_parameters(self): }, } - @pytest.mark.unit @patch( "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory" ) @@ -110,7 +105,6 @@ def test_warmup(self, mocked_factory): model_name_or_path="model", device="cpu", use_auth_token=None ) - @pytest.mark.unit @patch( "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory" ) @@ -121,7 +115,6 @@ def test_warmup_doesnt_reload(self, mocked_factory): embedder.warm_up() mocked_factory.get_embedding_backend.assert_called_once() - @pytest.mark.unit def test_run(self): embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") embedder.embedding_backend = MagicMock() @@ -138,7 +131,6 @@ def test_run(self): assert isinstance(doc.embedding, list) assert isinstance(doc.embedding[0], float) - @pytest.mark.unit def test_run_wrong_input_format(self): embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") @@ -155,7 +147,6 @@ def test_run_wrong_input_format(self): ): embedder.run(documents=list_integers_input) - @pytest.mark.unit def test_embed_metadata(self): embedder = SentenceTransformersDocumentEmbedder( model_name_or_path="model", metadata_fields_to_embed=["meta_field"], embedding_separator="\n" @@ -179,7 +170,6 @@ def test_embed_metadata(self): normalize_embeddings=False, ) - @pytest.mark.unit def test_prefix_suffix(self): embedder = SentenceTransformersDocumentEmbedder( model_name_or_path="model", diff --git a/test/components/embedders/test_sentence_transformers_embedding_backend.py b/test/components/embedders/test_sentence_transformers_embedding_backend.py index 9aec7954b3..64f96cd734 100644 --- a/test/components/embedders/test_sentence_transformers_embedding_backend.py +++ b/test/components/embedders/test_sentence_transformers_embedding_backend.py @@ -5,7 +5,6 @@ ) -@pytest.mark.unit @patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") def test_factory_behavior(mock_sentence_transformer): embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend( @@ -20,7 +19,6 @@ def test_factory_behavior(mock_sentence_transformer): assert another_embedding_backend is not embedding_backend -@pytest.mark.unit @patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") def test_model_initialization(mock_sentence_transformer): _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend( @@ -31,7 +29,6 @@ def test_model_initialization(mock_sentence_transformer): ) -@pytest.mark.unit @patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") def test_embedding_function_with_kwargs(mock_sentence_transformer): embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(model_name_or_path="model") diff --git a/test/components/embedders/test_sentence_transformers_text_embedder.py b/test/components/embedders/test_sentence_transformers_text_embedder.py index cc0f96e830..ba763bfb91 100644 --- a/test/components/embedders/test_sentence_transformers_text_embedder.py +++ b/test/components/embedders/test_sentence_transformers_text_embedder.py @@ -7,7 +7,6 @@ class TestSentenceTransformersTextEmbedder: - @pytest.mark.unit def test_init_default(self): embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") assert embedder.model_name_or_path == "model" @@ -19,7 +18,6 @@ def test_init_default(self): assert embedder.progress_bar is True assert embedder.normalize_embeddings is False - @pytest.mark.unit def test_init_with_parameters(self): embedder = SentenceTransformersTextEmbedder( model_name_or_path="model", @@ -40,7 +38,6 @@ def test_init_with_parameters(self): assert embedder.progress_bar is False assert embedder.normalize_embeddings is True - @pytest.mark.unit def test_to_dict(self): component = SentenceTransformersTextEmbedder(model_name_or_path="model") data = component.to_dict() @@ -58,7 +55,6 @@ def test_to_dict(self): }, } - @pytest.mark.unit def test_to_dict_with_custom_init_parameters(self): component = SentenceTransformersTextEmbedder( model_name_or_path="model", @@ -85,7 +81,6 @@ def test_to_dict_with_custom_init_parameters(self): }, } - @pytest.mark.unit def test_to_dict_not_serialize_token(self): component = SentenceTransformersTextEmbedder(model_name_or_path="model", token="awesome-token") data = component.to_dict() @@ -103,7 +98,6 @@ def test_to_dict_not_serialize_token(self): }, } - @pytest.mark.unit @patch( "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" ) @@ -115,7 +109,6 @@ def test_warmup(self, mocked_factory): model_name_or_path="model", device="cpu", use_auth_token=None ) - @pytest.mark.unit @patch( "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" ) @@ -126,7 +119,6 @@ def test_warmup_doesnt_reload(self, mocked_factory): embedder.warm_up() mocked_factory.get_embedding_backend.assert_called_once() - @pytest.mark.unit def test_run(self): embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") embedder.embedding_backend = MagicMock() @@ -140,7 +132,6 @@ def test_run(self): assert isinstance(embedding, list) assert all(isinstance(el, float) for el in embedding) - @pytest.mark.unit def test_run_wrong_input_format(self): embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") embedder.embedding_backend = MagicMock() From cde33fa906ebdb2fdfb356649bfa68030b6a5308 Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Thu, 26 Oct 2023 14:05:29 -0700 Subject: [PATCH 04/17] Gradient embeddings (#2) --- .../test_gradient_document_embedder.py | 139 ++++++++++++++++++ .../embedders/test_gradient_text_embedder.py | 111 ++++++++++++++ 2 files changed, 250 insertions(+) create mode 100644 test/components/embedders/test_gradient_document_embedder.py create mode 100644 test/components/embedders/test_gradient_text_embedder.py diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py new file mode 100644 index 0000000000..ba9568270f --- /dev/null +++ b/test/components/embedders/test_gradient_document_embedder.py @@ -0,0 +1,139 @@ +import pytest +from haystack.preview.components.embedders.gradient_document_embedder import GradientDocumentEmbedder +from unittest.mock import MagicMock, NonCallableMagicMock +import numpy as np + +from haystack.preview import Document + + +access_token = "access_token" +workspace_id = "workspace_id" +model = "bge-large" + + +def has_gradient(): + try: + import gradientai + + return True + except ModuleNotFoundError: + return False + + +@pytest.mark.skipif(not has_gradient(), reason="Gradient is not installed") +class TestGradientDocumentEmbedder: + @pytest.mark.unit + def test_init_from_env(self, monkeypatch): + monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", access_token) + monkeypatch.setenv("GRADIENT_WORKSPACE_ID", workspace_id) + + embedder = GradientDocumentEmbedder() + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_init_without_access_token(self, monkeypatch): + monkeypatch.delenv("GRADIENT_ACCESS_TOKEN", raising=True) + + with pytest.raises(ValueError): + GradientDocumentEmbedder(workspace_id=workspace_id) + + @pytest.mark.unit + def test_init_without_workspace(self, monkeypatch): + monkeypatch.delenv("GRADIENT_WORKSPACE_ID", raising=True) + + with pytest.raises(ValueError): + GradientDocumentEmbedder(access_token=access_token) + + @pytest.mark.unit + def test_init_from_params(self): + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_init_from_params_precedence(self, monkeypatch): + monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", "env_access_token") + monkeypatch.setenv("GRADIENT_WORKSPACE_ID", "env_workspace_id") + + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_to_dict(self): + component = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + data = component.to_dict() + assert data == { + "type": "GradientDocumentEmbedder", + "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, + } + + @pytest.mark.unit + def test_warmup(self): + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._gradient.get_embeddings_model = MagicMock() + embedder.warm_up() + embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") + + @pytest.mark.unit + def test_warmup_doesnt_reload(self): + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._gradient.get_embeddings_model = MagicMock(default_return_value="fake model") + embedder.warm_up() + embedder.warm_up() + embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") + + @pytest.mark.unit + def test_run_fail_if_not_warmed_up(self): + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + + with pytest.raises(RuntimeError, match="warm_up()"): + embedder.run(documents=[Document(text=f"document number {i}") for i in range(5)]) + + @pytest.mark.unit + def test_run(self): + from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess + + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(5)] + ) + + documents = [Document(text=f"document number {i}") for i in range(5)] + + result = embedder.run(documents=documents) + + assert isinstance(result["documents"], list) + assert len(result["documents"]) == len(documents) + for doc in result["documents"]: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert isinstance(doc.embedding[0], float) + + @pytest.mark.unit + def test_run_batch(self): + from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess + + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(110)] + ) + + documents = [Document(text=f"document number {i}") for i in range(110)] + + result = embedder.run(documents=documents) + + assert embedder._embedding_model.generate_embeddings.call_count == 2 + assert isinstance(result["documents"], list) + assert len(result["documents"]) == len(documents) + for doc in result["documents"]: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert isinstance(doc.embedding[0], float) diff --git a/test/components/embedders/test_gradient_text_embedder.py b/test/components/embedders/test_gradient_text_embedder.py new file mode 100644 index 0000000000..3e2a6762f2 --- /dev/null +++ b/test/components/embedders/test_gradient_text_embedder.py @@ -0,0 +1,111 @@ +import pytest +from haystack.preview.components.embedders.gradient_text_embedder import GradientTextEmbedder +from unittest.mock import MagicMock, NonCallableMagicMock +import numpy as np + + +access_token = "access_token" +workspace_id = "workspace_id" +model = "bge-large" + + +def has_gradient(): + try: + import gradientai + + return True + except ModuleNotFoundError: + return False + + +@pytest.mark.skipif(not has_gradient(), reason="Gradient is not installed") +class TestGradientTextEmbedder: + @pytest.mark.unit + def test_init_from_env(self, monkeypatch): + monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", access_token) + monkeypatch.setenv("GRADIENT_WORKSPACE_ID", workspace_id) + + embedder = GradientTextEmbedder() + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_init_without_access_token(self, monkeypatch): + monkeypatch.delenv("GRADIENT_ACCESS_TOKEN", raising=True) + + with pytest.raises(ValueError): + GradientTextEmbedder(workspace_id=workspace_id) + + @pytest.mark.unit + def test_init_without_workspace(self, monkeypatch): + monkeypatch.delenv("GRADIENT_WORKSPACE_ID", raising=True) + + with pytest.raises(ValueError): + GradientTextEmbedder(access_token=access_token) + + @pytest.mark.unit + def test_init_from_params(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_init_from_params_precedence(self, monkeypatch): + monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", "env_access_token") + monkeypatch.setenv("GRADIENT_WORKSPACE_ID", "env_workspace_id") + + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_to_dict(self): + component = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + data = component.to_dict() + assert data == { + "type": "GradientTextEmbedder", + "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, + } + + @pytest.mark.unit + def test_warmup(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._gradient.get_embeddings_model = MagicMock() + embedder.warm_up() + embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") + + @pytest.mark.unit + def test_warmup_doesnt_reload(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._gradient.get_embeddings_model = MagicMock(default_return_value="fake model") + embedder.warm_up() + embedder.warm_up() + embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") + + @pytest.mark.unit + def test_run_fail_if_not_warmed_up(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + + with pytest.raises(RuntimeError, match="warm_up()"): + embedder.run(text="The food was delicious") + + @pytest.mark.unit + def test_run(self): + from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess + + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": 0}] + ) + + result = embedder.run(text="The food was delicious") + embedder._embedding_model.generate_embeddings.assert_called_once_with( + inputs=[{"input": "The food was delicious"}] + ) + + assert len(result["embedding"]) == 1024 # 1024 is the bge-large embedding size + assert all(isinstance(x, float) for x in result["embedding"]) From 78a4ac1f20bed9ca56796e1a544fc954bc95c532 Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Fri, 27 Oct 2023 13:31:11 -0700 Subject: [PATCH 05/17] apply pr requests --- .../embedders/test_gradient_document_embedder.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py index ba9568270f..e68a2e552e 100644 --- a/test/components/embedders/test_gradient_document_embedder.py +++ b/test/components/embedders/test_gradient_document_embedder.py @@ -137,3 +137,14 @@ def test_run_batch(self): assert isinstance(doc, Document) assert isinstance(doc.embedding, list) assert isinstance(doc.embedding[0], float) + + @pytest.mark.unit + def test_run_empty(self): + from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess + + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + + result = embedder.run(documents=[]) + + assert result["documents"] == [] From bfe0320ebfe555504592ffabeafb8fc6b7ea64bf Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Tue, 7 Nov 2023 19:20:24 -0800 Subject: [PATCH 06/17] apply the changes (#10) --- .../test_gradient_document_embedder.py | 16 +------- .../embedders/test_gradient_text_embedder.py | 39 +++++++++++++------ 2 files changed, 29 insertions(+), 26 deletions(-) diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py index e68a2e552e..297596df9e 100644 --- a/test/components/embedders/test_gradient_document_embedder.py +++ b/test/components/embedders/test_gradient_document_embedder.py @@ -1,4 +1,5 @@ import pytest +from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess from haystack.preview.components.embedders.gradient_document_embedder import GradientDocumentEmbedder from unittest.mock import MagicMock, NonCallableMagicMock import numpy as np @@ -11,16 +12,6 @@ model = "bge-large" -def has_gradient(): - try: - import gradientai - - return True - except ModuleNotFoundError: - return False - - -@pytest.mark.skipif(not has_gradient(), reason="Gradient is not installed") class TestGradientDocumentEmbedder: @pytest.mark.unit def test_init_from_env(self, monkeypatch): @@ -96,8 +87,6 @@ def test_run_fail_if_not_warmed_up(self): @pytest.mark.unit def test_run(self): - from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess - embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( @@ -108,6 +97,7 @@ def test_run(self): result = embedder.run(documents=documents) + assert embedder._embedding_model.generate_embeddings.call_count == 1 assert isinstance(result["documents"], list) assert len(result["documents"]) == len(documents) for doc in result["documents"]: @@ -140,8 +130,6 @@ def test_run_batch(self): @pytest.mark.unit def test_run_empty(self): - from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess - embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() diff --git a/test/components/embedders/test_gradient_text_embedder.py b/test/components/embedders/test_gradient_text_embedder.py index 3e2a6762f2..fa553e27f1 100644 --- a/test/components/embedders/test_gradient_text_embedder.py +++ b/test/components/embedders/test_gradient_text_embedder.py @@ -1,4 +1,5 @@ import pytest +from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess from haystack.preview.components.embedders.gradient_text_embedder import GradientTextEmbedder from unittest.mock import MagicMock, NonCallableMagicMock import numpy as np @@ -9,16 +10,6 @@ model = "bge-large" -def has_gradient(): - try: - import gradientai - - return True - except ModuleNotFoundError: - return False - - -@pytest.mark.skipif(not has_gradient(), reason="Gradient is not installed") class TestGradientTextEmbedder: @pytest.mark.unit def test_init_from_env(self, monkeypatch): @@ -93,9 +84,33 @@ def test_run_fail_if_not_warmed_up(self): embedder.run(text="The food was delicious") @pytest.mark.unit - def test_run(self): - from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess + def test_run_fail_when_no_embeddings_returned(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess(embeddings=[]) + + with pytest.raises(RuntimeError): + _result = embedder.run(text="The food was delicious") + embedder._embedding_model.generate_embeddings.assert_called_once_with( + inputs=[{"input": "The food was delicious"}] + ) + + @pytest.mark.unit + def test_run_empty_string(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": 0}] + ) + + result = embedder.run(text="") + embedder._embedding_model.generate_embeddings.assert_called_once_with(inputs=[{"input": ""}]) + assert len(result["embedding"]) == 1024 # 1024 is the bge-large embedding size + assert all(isinstance(x, float) for x in result["embedding"]) + + @pytest.mark.unit + def test_run(self): embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( From 2cea8ff8a1e26133981133d45cc23b676f5b1e8b Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Tue, 14 Nov 2023 18:19:14 -0800 Subject: [PATCH 07/17] update tests --- .../test_gradient_document_embedder.py | 30 +++++++++++++++---- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py index 297596df9e..75c57c61ff 100644 --- a/test/components/embedders/test_gradient_document_embedder.py +++ b/test/components/embedders/test_gradient_document_embedder.py @@ -83,7 +83,7 @@ def test_run_fail_if_not_warmed_up(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) with pytest.raises(RuntimeError, match="warm_up()"): - embedder.run(documents=[Document(text=f"document number {i}") for i in range(5)]) + embedder.run(documents=[Document(content=f"document number {i}") for i in range(5)]) @pytest.mark.unit def test_run(self): @@ -93,7 +93,7 @@ def test_run(self): embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(5)] ) - documents = [Document(text=f"document number {i}") for i in range(5)] + documents = [Document(content=f"document number {i}") for i in range(5)] result = embedder.run(documents=documents) @@ -107,8 +107,6 @@ def test_run(self): @pytest.mark.unit def test_run_batch(self): - from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess - embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() @@ -116,7 +114,7 @@ def test_run_batch(self): embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(110)] ) - documents = [Document(text=f"document number {i}") for i in range(110)] + documents = [Document(content=f"document number {i}") for i in range(110)] result = embedder.run(documents=documents) @@ -128,6 +126,28 @@ def test_run_batch(self): assert isinstance(doc.embedding, list) assert isinstance(doc.embedding[0], float) + @pytest.mark.unit + def test_run_custom_batch(self): + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id, batch_size=20) + embedder._embedding_model = NonCallableMagicMock() + + DOCUMENT_COUNT = 101 + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(DOCUMENT_COUNT)] + ) + + documents = [Document(content=f"document number {i}") for i in range(DOCUMENT_COUNT)] + + result = embedder.run(documents=documents) + + assert embedder._embedding_model.generate_embeddings.call_count == 6 + assert isinstance(result["documents"], list) + assert len(result["documents"]) == len(documents) + for doc in result["documents"]: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert isinstance(doc.embedding[0], float) + @pytest.mark.unit def test_run_empty(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) From 62a0578235b94ae088044dd5472e2560c0f8aaf3 Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Mon, 27 Nov 2023 19:55:59 -0800 Subject: [PATCH 08/17] more updates --- test/components/embedders/test_gradient_document_embedder.py | 4 ++-- test/components/embedders/test_gradient_text_embedder.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py index 75c57c61ff..e9cbfdd9e3 100644 --- a/test/components/embedders/test_gradient_document_embedder.py +++ b/test/components/embedders/test_gradient_document_embedder.py @@ -1,6 +1,6 @@ import pytest from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess -from haystack.preview.components.embedders.gradient_document_embedder import GradientDocumentEmbedder +from haystack.components.embedders.gradient_document_embedder import GradientDocumentEmbedder from unittest.mock import MagicMock, NonCallableMagicMock import numpy as np @@ -59,7 +59,7 @@ def test_to_dict(self): component = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) data = component.to_dict() assert data == { - "type": "GradientDocumentEmbedder", + "type": "haystack.component.embedders.gradient_document_embedder.GradientDocumentEmbedder", "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, } diff --git a/test/components/embedders/test_gradient_text_embedder.py b/test/components/embedders/test_gradient_text_embedder.py index fa553e27f1..779ed5c1ee 100644 --- a/test/components/embedders/test_gradient_text_embedder.py +++ b/test/components/embedders/test_gradient_text_embedder.py @@ -1,6 +1,6 @@ import pytest from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess -from haystack.preview.components.embedders.gradient_text_embedder import GradientTextEmbedder +from haystack.components.embedders.gradient_text_embedder import GradientTextEmbedder from unittest.mock import MagicMock, NonCallableMagicMock import numpy as np @@ -57,7 +57,7 @@ def test_to_dict(self): component = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) data = component.to_dict() assert data == { - "type": "GradientTextEmbedder", + "type": "haystack.component.embedders.gradient_document_embedder.GradientTextEmbedder", "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, } From 8c8ffd658b1e22b9d72b78fc578d9b74f6e1583e Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Mon, 27 Nov 2023 22:43:56 -0800 Subject: [PATCH 09/17] more updates --- .../embedders/test_gradient_document_embedder.py | 8 ++++---- test/components/embedders/test_gradient_text_embedder.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py index e9cbfdd9e3..ffb44519fd 100644 --- a/test/components/embedders/test_gradient_document_embedder.py +++ b/test/components/embedders/test_gradient_document_embedder.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock, NonCallableMagicMock import numpy as np -from haystack.preview import Document +from haystack import Document access_token = "access_token" @@ -25,14 +25,14 @@ def test_init_from_env(self, monkeypatch): @pytest.mark.unit def test_init_without_access_token(self, monkeypatch): - monkeypatch.delenv("GRADIENT_ACCESS_TOKEN", raising=True) + monkeypatch.delenv("GRADIENT_ACCESS_TOKEN", raising=False) with pytest.raises(ValueError): GradientDocumentEmbedder(workspace_id=workspace_id) @pytest.mark.unit def test_init_without_workspace(self, monkeypatch): - monkeypatch.delenv("GRADIENT_WORKSPACE_ID", raising=True) + monkeypatch.delenv("GRADIENT_WORKSPACE_ID", raising=False) with pytest.raises(ValueError): GradientDocumentEmbedder(access_token=access_token) @@ -59,7 +59,7 @@ def test_to_dict(self): component = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) data = component.to_dict() assert data == { - "type": "haystack.component.embedders.gradient_document_embedder.GradientDocumentEmbedder", + "type": "haystack.components.embedders.gradient_document_embedder.GradientDocumentEmbedder", "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, } diff --git a/test/components/embedders/test_gradient_text_embedder.py b/test/components/embedders/test_gradient_text_embedder.py index 779ed5c1ee..82c84e3904 100644 --- a/test/components/embedders/test_gradient_text_embedder.py +++ b/test/components/embedders/test_gradient_text_embedder.py @@ -23,14 +23,14 @@ def test_init_from_env(self, monkeypatch): @pytest.mark.unit def test_init_without_access_token(self, monkeypatch): - monkeypatch.delenv("GRADIENT_ACCESS_TOKEN", raising=True) + monkeypatch.delenv("GRADIENT_ACCESS_TOKEN", raising=False) with pytest.raises(ValueError): GradientTextEmbedder(workspace_id=workspace_id) @pytest.mark.unit def test_init_without_workspace(self, monkeypatch): - monkeypatch.delenv("GRADIENT_WORKSPACE_ID", raising=True) + monkeypatch.delenv("GRADIENT_WORKSPACE_ID", raising=False) with pytest.raises(ValueError): GradientTextEmbedder(access_token=access_token) @@ -57,7 +57,7 @@ def test_to_dict(self): component = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) data = component.to_dict() assert data == { - "type": "haystack.component.embedders.gradient_document_embedder.GradientTextEmbedder", + "type": "haystack.components.embedders.gradient_text_embedder.GradientTextEmbedder", "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, } From 01bd48991ac2bf150fe96bdac865ca063246fd60 Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Fri, 1 Dec 2023 19:20:25 +0100 Subject: [PATCH 10/17] prepare --- .../embedders/gradient_document_embedder.py | 112 +++++++++++++++ .../embedders/gradient_text_embedder.py | 88 ++++++++++++ integrations/gradient/generator/__init__.py | 0 integrations/gradient/generator/base.py | 128 ++++++++++++++++++ .../test}/test_gradient_document_embedder.py | 0 .../test}/test_gradient_text_embedder.py | 0 6 files changed, 328 insertions(+) create mode 100644 integrations/gradient/embedders/gradient_document_embedder.py create mode 100644 integrations/gradient/embedders/gradient_text_embedder.py create mode 100644 integrations/gradient/generator/__init__.py create mode 100644 integrations/gradient/generator/base.py rename {test/components/embedders => integrations/gradient/test}/test_gradient_document_embedder.py (100%) rename {test/components/embedders => integrations/gradient/test}/test_gradient_text_embedder.py (100%) diff --git a/integrations/gradient/embedders/gradient_document_embedder.py b/integrations/gradient/embedders/gradient_document_embedder.py new file mode 100644 index 0000000000..e2d3458c59 --- /dev/null +++ b/integrations/gradient/embedders/gradient_document_embedder.py @@ -0,0 +1,112 @@ +import logging +from typing import List, Optional, Dict, Any + +from haystack import component, Document, default_to_dict +from haystack.lazy_imports import LazyImport + +with LazyImport(message="Run 'pip install gradientai'") as gradientai_import: + from gradientai import Gradient + +logger = logging.getLogger(__name__) + + +@component +class GradientDocumentEmbedder: + """ + A component for computing Document embeddings using Gradient AI API.. + The embedding of each Document is stored in the `embedding` field of the Document. + + ```python + embedder = GradientDocumentEmbedder( + access_token=gradient_access_token, + workspace_id=gradient_workspace_id, + model_name="bge_large")) + p = Pipeline() + p.add_component(embedder, name="document_embedder") + p.add_component(instance=GradientDocumentEmbedder( + p.add_component(instance=DocumentWriter(document_store=InMemoryDocumentStore()), name="document_writer") + p.connect("document_embedder", "document_writer") + p.run({"document_embedder": {"documents": documents}}) + ``` + """ + + def __init__( + self, + *, + model_name: str = "bge-large", + batch_size: int = 100, + access_token: Optional[str] = None, + workspace_id: Optional[str] = None, + host: Optional[str] = None, + ) -> None: + """ + Create a GradientDocumentEmbedder component. + + :param model_name: The name of the model to use. + :param access_token: The Gradient access token. If not provided it's read from the environment + variable GRADIENT_ACCESS_TOKEN. + :param workspace_id: The Gradient workspace ID. If not provided it's read from the environment + variable GRADIENT_WORKSPACE_ID. + :param host: The Gradient host. By default it uses https://api.gradient.ai/. + """ + gradientai_import.check() + self._batch_size = batch_size + self._host = host + self._model_name = model_name + + self._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id) + + def _get_telemetry_data(self) -> Dict[str, Any]: + """ + Data that is sent to Posthog for usage analytics. + """ + return {"model": self._model_name} + + def to_dict(self) -> dict: + """ + Serialize the component to a Python dictionary. + """ + return default_to_dict(self, workspace_id=self._gradient.workspace_id, model_name=self._model_name) + + def warm_up(self) -> None: + """ + Load the embedding model. + """ + if not hasattr(self, "_embedding_model"): + self._embedding_model = self._gradient.get_embeddings_model(slug=self._model_name) + + def _generate_embeddings(self, documents: List[Document], batch_size: int) -> List[List[float]]: + """ + Batches the documents and generates the embeddings. + """ + batches = [documents[i : i + batch_size] for i in range(0, len(documents), batch_size)] + + embeddings = [] + for batch in batches: + response = self._embedding_model.generate_embeddings(inputs=[{"input": doc.content} for doc in batch]) + embeddings.extend([e.embedding for e in response.embeddings]) + + return embeddings + + @component.output_types(documents=List[Document]) + def run(self, documents: List[Document]): + """ + Embed a list of Documents. + The embedding of each Document is stored in the `embedding` field of the Document. + + :param documents: A list of Documents to embed. + """ + if not isinstance(documents, list) or documents and any(not isinstance(doc, Document) for doc in documents): + raise TypeError( + "GradientDocumentEmbedder expects a list of Documents as input." + "In case you want to embed a list of strings, please use the GradientTextEmbedder." + ) + + if not hasattr(self, "_embedding_model"): + raise RuntimeError("The embedding model has not been loaded. Please call warm_up() before running.") + + embeddings = self._generate_embeddings(documents=documents, batch_size=self._batch_size) + for doc, embedding in zip(documents, embeddings): + doc.embedding = embedding + + return {"documents": documents} diff --git a/integrations/gradient/embedders/gradient_text_embedder.py b/integrations/gradient/embedders/gradient_text_embedder.py new file mode 100644 index 0000000000..28a83a52e2 --- /dev/null +++ b/integrations/gradient/embedders/gradient_text_embedder.py @@ -0,0 +1,88 @@ +from typing import Any, Dict, List, Optional + +from haystack import component, default_to_dict +from haystack.lazy_imports import LazyImport + +with LazyImport(message="Run 'pip install gradientai'") as gradientai_import: + from gradientai import Gradient + + +@component +class GradientTextEmbedder: + """ + A component for embedding strings using models hosted on Gradient AI (https://gradient.ai). + + ```python + embedder = GradientTextEmbedder( + access_token=gradient_access_token, + workspace_id=gradient_workspace_id, + model_name="bge_large") + p = Pipeline() + p.add_component(instance=embedder, name="text_embedder") + p.add_component(instance=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()), name="retriever") + p.connect("text_embedder", "retriever") + p.run("embed me!!!") + ``` + """ + + def __init__( + self, + *, + model_name: str = "bge-large", + access_token: Optional[str] = None, + workspace_id: Optional[str] = None, + host: Optional[str] = None, + ) -> None: + """ + Create a GradientTextEmbedder component. + + :param model_name: The name of the model to use. + :param access_token: The Gradient access token. If not provided it's read from the environment + variable GRADIENT_ACCESS_TOKEN. + :param workspace_id: The Gradient workspace ID. If not provided it's read from the environment + variable GRADIENT_WORKSPACE_ID. + :param host: The Gradient host. By default it uses https://api.gradient.ai/. + """ + gradientai_import.check() + self._host = host + self._model_name = model_name + + self._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id) + + def _get_telemetry_data(self) -> Dict[str, Any]: + """ + Data that is sent to Posthog for usage analytics. + """ + return {"model": self._model_name} + + def to_dict(self) -> dict: + """ + Serialize the component to a Python dictionary. + """ + return default_to_dict(self, workspace_id=self._gradient.workspace_id, model_name=self._model_name) + + def warm_up(self) -> None: + """ + Load the embedding model. + """ + if not hasattr(self, "_embedding_model"): + self._embedding_model = self._gradient.get_embeddings_model(slug=self._model_name) + + @component.output_types(embedding=List[float]) + def run(self, text: str): + """Generates an embedding for a single text.""" + if not isinstance(text, str): + raise TypeError( + "GradientTextEmbedder expects a string as an input." + "In case you want to embed a list of Documents, please use the GradientDocumentEmbedder." + ) + + if not hasattr(self, "_embedding_model"): + raise RuntimeError("The embedding model has not been loaded. Please call warm_up() before running.") + + result = self._embedding_model.generate_embeddings(inputs=[{"input": text}]) + + if (not result) or (result.embeddings is None) or (len(result.embeddings) == 0): + raise RuntimeError("The embedding model did not return any embeddings.") + + return {"embedding": result.embeddings[0].embedding} diff --git a/integrations/gradient/generator/__init__.py b/integrations/gradient/generator/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/integrations/gradient/generator/base.py b/integrations/gradient/generator/base.py new file mode 100644 index 0000000000..4853a92291 --- /dev/null +++ b/integrations/gradient/generator/base.py @@ -0,0 +1,128 @@ +from typing import List, Optional, Dict, Any, overload + +import logging +from haystack.lazy_imports import LazyImport + +from haystack import component, default_to_dict + +with LazyImport(message="Run 'pip install gradientai'") as gradientai_import: + from gradientai import Gradient + +logger = logging.getLogger(__name__) + + +@component +class GradientGenerator: + """ + LLM Generator interfacing [Gradient AI](https://gradient.ai/). + + Queries the LLM using Gradient AI's SDK ('gradientai' package). + See [Gradient AI API](https://docs.gradient.ai/docs/sdk-quickstart) for more details. + + ```python + llm = GradientGenerator( + access_token=gradient_access_token, + workspace_id=gradient_workspace_id, + base_model_slug="llama2-7b-chat") + llm.warm_up() + print(llm.run(prompt="What is the meaning of life?")) + # Output: {'replies': ['42']} + ``` + """ + + def __init__( + self, + *, + access_token: Optional[str] = None, + base_model_slug: Optional[str] = None, + host: Optional[str] = None, + max_generated_token_count: Optional[int] = None, + model_adapter_id: Optional[str] = None, + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + workspace_id: Optional[str] = None, + ) -> None: + """ + Create a GradientGenerator component. + + :param access_token: The Gradient access token. If not provided it's read from the environment + variable GRADIENT_ACCESS_TOKEN. + :param base_model_slug: The base model slug to use. + :param host: The Gradient host. By default it uses https://api.gradient.ai/. + :param max_generated_token_count: The maximum number of tokens to generate. + :param model_adapter_id: The model adapter ID to use. + :param temperature: The temperature to use. + :param top_k: The top k to use. + :param top_p: The top p to use. + :param workspace_id: The Gradient workspace ID. If not provided it's read from the environment + variable GRADIENT_WORKSPACE_ID. + """ + gradientai_import.check() + + self._access_token = access_token + self._base_model_slug = base_model_slug + self._host = host + self._max_generated_token_count = max_generated_token_count + self._model_adapter_id = model_adapter_id + self._temperature = temperature + self._top_k = top_k + self._top_p = top_p + self._workspace_id = workspace_id + + has_base_model_slug = base_model_slug is not None and base_model_slug != "" + has_model_adapter_id = model_adapter_id is not None and model_adapter_id != "" + + if not has_base_model_slug and not has_model_adapter_id: + raise ValueError("Either base_model_slug or model_adapter_id must be provided.") + if has_base_model_slug and has_model_adapter_id: + raise ValueError("Only one of base_model_slug or model_adapter_id must be provided.") + + if has_base_model_slug: + self._base_model_slug = base_model_slug + if has_model_adapter_id: + self._model_adapter_id = model_adapter_id + + self._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id) + + def to_dict(self) -> Dict[str, Any]: + """ + Serialize this component to a dictionary. + """ + return default_to_dict( + self, + base_model_slug=self._base_model_slug, + host=self._host, + max_generated_token_count=self._max_generated_token_count, + model_adapter_id=self._model_adapter_id, + temperature=self._temperature, + top_k=self._top_k, + top_p=self._top_p, + workspace_id=self._workspace_id, + ) + + def warm_up(self): + """ + Initializes the LLM model instance if it doesn't exist. + """ + if not hasattr(self, "_model"): + if isinstance(self._base_model_slug, str): + self._model = self._gradient.get_base_model(base_model_slug=self._base_model_slug) + if isinstance(self._model_adapter_id, str): + self._model = self._gradient.get_model_adapter(model_adapter_id=self._model_adapter_id) + + @component.output_types(replies=List[str]) + def run(self, prompt: str): + """ + Queries the LLM with the prompt to produce replies. + + :param prompt: The prompt to be sent to the generative model. + """ + resp = self._model.complete( + query=prompt, + max_generated_token_count=self._max_generated_token_count, + temperature=self._temperature, + top_k=self._top_k, + top_p=self._top_p, + ) + return {"replies": [resp.generated_output]} diff --git a/test/components/embedders/test_gradient_document_embedder.py b/integrations/gradient/test/test_gradient_document_embedder.py similarity index 100% rename from test/components/embedders/test_gradient_document_embedder.py rename to integrations/gradient/test/test_gradient_document_embedder.py diff --git a/test/components/embedders/test_gradient_text_embedder.py b/integrations/gradient/test/test_gradient_text_embedder.py similarity index 100% rename from test/components/embedders/test_gradient_text_embedder.py rename to integrations/gradient/test/test_gradient_text_embedder.py From c0e18ddf925f8c1e0be1b5ac1a0f042cc4f7f774 Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Fri, 1 Dec 2023 19:31:54 +0100 Subject: [PATCH 11/17] remove commits migration artifacts --- test/components/embedders/__init__.py | 0 .../test_openai_document_embedder.py | 270 ------------------ .../embedders/test_openai_text_embedder.py | 109 ------- ...sentence_transformers_document_embedder.py | 198 ------------- ...sentence_transformers_embedding_backend.py | 39 --- ...est_sentence_transformers_text_embedder.py | 142 --------- 6 files changed, 758 deletions(-) delete mode 100644 test/components/embedders/__init__.py delete mode 100644 test/components/embedders/test_openai_document_embedder.py delete mode 100644 test/components/embedders/test_openai_text_embedder.py delete mode 100644 test/components/embedders/test_sentence_transformers_document_embedder.py delete mode 100644 test/components/embedders/test_sentence_transformers_embedding_backend.py delete mode 100644 test/components/embedders/test_sentence_transformers_text_embedder.py diff --git a/test/components/embedders/__init__.py b/test/components/embedders/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/test/components/embedders/test_openai_document_embedder.py b/test/components/embedders/test_openai_document_embedder.py deleted file mode 100644 index 29b3074a23..0000000000 --- a/test/components/embedders/test_openai_document_embedder.py +++ /dev/null @@ -1,270 +0,0 @@ -from unittest.mock import patch -from typing import List, cast - -import pytest -import numpy as np -import openai -from openai.util import convert_to_openai_object -from openai.openai_object import OpenAIObject - -from haystack import Document -from haystack.components.embedders.openai_document_embedder import OpenAIDocumentEmbedder - - -def mock_openai_response(input: List[str], model: str = "text-embedding-ada-002", **kwargs) -> OpenAIObject: - dict_response = { - "object": "list", - "data": [ - {"object": "embedding", "index": i, "embedding": np.random.rand(1536).tolist()} for i in range(len(input)) - ], - "model": model, - "usage": {"prompt_tokens": 4, "total_tokens": 4}, - } - - return cast(OpenAIObject, convert_to_openai_object(dict_response)) - - -class TestOpenAIDocumentEmbedder: - def test_init_default(self, monkeypatch): - openai.api_key = None - monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") - embedder = OpenAIDocumentEmbedder() - - assert openai.api_key == "fake-api-key" - - assert embedder.model_name == "text-embedding-ada-002" - assert embedder.organization is None - assert embedder.prefix == "" - assert embedder.suffix == "" - assert embedder.batch_size == 32 - assert embedder.progress_bar is True - assert embedder.metadata_fields_to_embed == [] - assert embedder.embedding_separator == "\n" - - def test_init_with_parameters(self): - embedder = OpenAIDocumentEmbedder( - api_key="fake-api-key", - model_name="model", - organization="my-org", - prefix="prefix", - suffix="suffix", - batch_size=64, - progress_bar=False, - metadata_fields_to_embed=["test_field"], - embedding_separator=" | ", - ) - assert openai.api_key == "fake-api-key" - assert openai.organization == "my-org" - - assert embedder.organization == "my-org" - assert embedder.model_name == "model" - assert embedder.prefix == "prefix" - assert embedder.suffix == "suffix" - assert embedder.batch_size == 64 - assert embedder.progress_bar is False - assert embedder.metadata_fields_to_embed == ["test_field"] - assert embedder.embedding_separator == " | " - - def test_init_fail_wo_api_key(self, monkeypatch): - openai.api_key = None - monkeypatch.delenv("OPENAI_API_KEY", raising=False) - with pytest.raises(ValueError, match="OpenAIDocumentEmbedder expects an OpenAI API key"): - OpenAIDocumentEmbedder() - - def test_to_dict(self): - component = OpenAIDocumentEmbedder(api_key="fake-api-key") - data = component.to_dict() - assert data == { - "type": "haystack.components.embedders.openai_document_embedder.OpenAIDocumentEmbedder", - "init_parameters": { - "model_name": "text-embedding-ada-002", - "organization": None, - "prefix": "", - "suffix": "", - "batch_size": 32, - "progress_bar": True, - "metadata_fields_to_embed": [], - "embedding_separator": "\n", - }, - } - - def test_to_dict_with_custom_init_parameters(self): - component = OpenAIDocumentEmbedder( - api_key="fake-api-key", - model_name="model", - organization="my-org", - prefix="prefix", - suffix="suffix", - batch_size=64, - progress_bar=False, - metadata_fields_to_embed=["test_field"], - embedding_separator=" | ", - ) - data = component.to_dict() - assert data == { - "type": "haystack.components.embedders.openai_document_embedder.OpenAIDocumentEmbedder", - "init_parameters": { - "model_name": "model", - "organization": "my-org", - "prefix": "prefix", - "suffix": "suffix", - "batch_size": 64, - "progress_bar": False, - "metadata_fields_to_embed": ["test_field"], - "embedding_separator": " | ", - }, - } - - def test_prepare_texts_to_embed_w_metadata(self): - documents = [ - Document(content=f"document number {i}:\ncontent", meta={"meta_field": f"meta_value {i}"}) for i in range(5) - ] - - embedder = OpenAIDocumentEmbedder( - api_key="fake-api-key", metadata_fields_to_embed=["meta_field"], embedding_separator=" | " - ) - - prepared_texts = embedder._prepare_texts_to_embed(documents) - - # note that newline is replaced by space - assert prepared_texts == [ - "meta_value 0 | document number 0: content", - "meta_value 1 | document number 1: content", - "meta_value 2 | document number 2: content", - "meta_value 3 | document number 3: content", - "meta_value 4 | document number 4: content", - ] - - def test_prepare_texts_to_embed_w_suffix(self): - documents = [Document(content=f"document number {i}") for i in range(5)] - - embedder = OpenAIDocumentEmbedder(api_key="fake-api-key", prefix="my_prefix ", suffix=" my_suffix") - - prepared_texts = embedder._prepare_texts_to_embed(documents) - - assert prepared_texts == [ - "my_prefix document number 0 my_suffix", - "my_prefix document number 1 my_suffix", - "my_prefix document number 2 my_suffix", - "my_prefix document number 3 my_suffix", - "my_prefix document number 4 my_suffix", - ] - - def test_embed_batch(self): - texts = ["text 1", "text 2", "text 3", "text 4", "text 5"] - - with patch("haystack.components.embedders.openai_document_embedder.openai.Embedding") as openai_embedding_patch: - openai_embedding_patch.create.side_effect = mock_openai_response - embedder = OpenAIDocumentEmbedder(api_key="fake-api-key", model_name="model") - - embeddings, metadata = embedder._embed_batch(texts_to_embed=texts, batch_size=2) - - assert openai_embedding_patch.create.call_count == 3 - - assert isinstance(embeddings, list) - assert len(embeddings) == len(texts) - for embedding in embeddings: - assert isinstance(embedding, list) - assert len(embedding) == 1536 - assert all(isinstance(x, float) for x in embedding) - - # openai.Embedding.create is called 3 times - assert metadata == {"model": "model", "usage": {"prompt_tokens": 3 * 4, "total_tokens": 3 * 4}} - - def test_run(self): - docs = [ - Document(content="I love cheese", meta={"topic": "Cuisine"}), - Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}), - ] - - model = "text-similarity-ada-001" - with patch("haystack.components.embedders.openai_document_embedder.openai.Embedding") as openai_embedding_patch: - openai_embedding_patch.create.side_effect = mock_openai_response - embedder = OpenAIDocumentEmbedder( - api_key="fake-api-key", - model_name=model, - prefix="prefix ", - suffix=" suffix", - metadata_fields_to_embed=["topic"], - embedding_separator=" | ", - ) - - result = embedder.run(documents=docs) - - openai_embedding_patch.create.assert_called_once_with( - model=model, - input=[ - "prefix Cuisine | I love cheese suffix", - "prefix ML | A transformer is a deep learning architecture suffix", - ], - ) - documents_with_embeddings = result["documents"] - metadata = result["metadata"] - - assert isinstance(documents_with_embeddings, list) - assert len(documents_with_embeddings) == len(docs) - for doc in documents_with_embeddings: - assert isinstance(doc, Document) - assert isinstance(doc.embedding, list) - assert len(doc.embedding) == 1536 - assert all(isinstance(x, float) for x in doc.embedding) - assert metadata == {"model": model, "usage": {"prompt_tokens": 4, "total_tokens": 4}} - - def test_run_custom_batch_size(self): - docs = [ - Document(content="I love cheese", meta={"topic": "Cuisine"}), - Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}), - ] - - model = "text-similarity-ada-001" - with patch("haystack.components.embedders.openai_document_embedder.openai.Embedding") as openai_embedding_patch: - openai_embedding_patch.create.side_effect = mock_openai_response - embedder = OpenAIDocumentEmbedder( - api_key="fake-api-key", - model_name=model, - prefix="prefix ", - suffix=" suffix", - metadata_fields_to_embed=["topic"], - embedding_separator=" | ", - batch_size=1, - ) - - result = embedder.run(documents=docs) - - assert openai_embedding_patch.create.call_count == 2 - - documents_with_embeddings = result["documents"] - metadata = result["metadata"] - - assert isinstance(documents_with_embeddings, list) - assert len(documents_with_embeddings) == len(docs) - for doc in documents_with_embeddings: - assert isinstance(doc, Document) - assert isinstance(doc.embedding, list) - assert len(doc.embedding) == 1536 - assert all(isinstance(x, float) for x in doc.embedding) - - # openai.Embedding.create is called 2 times - assert metadata == {"model": model, "usage": {"prompt_tokens": 2 * 4, "total_tokens": 2 * 4}} - - def test_run_wrong_input_format(self): - embedder = OpenAIDocumentEmbedder(api_key="fake-api-key") - - # wrong formats - string_input = "text" - list_integers_input = [1, 2, 3] - - with pytest.raises(TypeError, match="OpenAIDocumentEmbedder expects a list of Documents as input"): - embedder.run(documents=string_input) - - with pytest.raises(TypeError, match="OpenAIDocumentEmbedder expects a list of Documents as input"): - embedder.run(documents=list_integers_input) - - def test_run_on_empty_list(self): - embedder = OpenAIDocumentEmbedder(api_key="fake-api-key") - - empty_list_input = [] - result = embedder.run(documents=empty_list_input) - - assert result["documents"] is not None - assert not result["documents"] # empty list diff --git a/test/components/embedders/test_openai_text_embedder.py b/test/components/embedders/test_openai_text_embedder.py deleted file mode 100644 index 183fd45a1c..0000000000 --- a/test/components/embedders/test_openai_text_embedder.py +++ /dev/null @@ -1,109 +0,0 @@ -from unittest.mock import patch -import pytest -import openai -from openai.util import convert_to_openai_object -import numpy as np - -from haystack.components.embedders.openai_text_embedder import OpenAITextEmbedder - - -def mock_openai_response(model: str = "text-embedding-ada-002", **kwargs) -> openai.openai_object.OpenAIObject: - dict_response = { - "object": "list", - "data": [{"object": "embedding", "index": 0, "embedding": np.random.rand(1536).tolist()}], - "model": model, - "usage": {"prompt_tokens": 4, "total_tokens": 4}, - } - - return convert_to_openai_object(dict_response) - - -class TestOpenAITextEmbedder: - def test_init_default(self, monkeypatch): - openai.api_key = None - monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") - embedder = OpenAITextEmbedder() - - assert openai.api_key == "fake-api-key" - assert embedder.model_name == "text-embedding-ada-002" - assert embedder.organization is None - assert embedder.prefix == "" - assert embedder.suffix == "" - - def test_init_with_parameters(self): - embedder = OpenAITextEmbedder( - api_key="fake-api-key", - model_name="model", - organization="fake-organization", - prefix="prefix", - suffix="suffix", - ) - assert openai.api_key == "fake-api-key" - assert embedder.model_name == "model" - assert embedder.organization == "fake-organization" - assert openai.organization == "fake-organization" - assert embedder.prefix == "prefix" - assert embedder.suffix == "suffix" - - def test_init_fail_wo_api_key(self, monkeypatch): - openai.api_key = None - monkeypatch.delenv("OPENAI_API_KEY", raising=False) - with pytest.raises(ValueError, match="OpenAITextEmbedder expects an OpenAI API key"): - OpenAITextEmbedder() - - def test_to_dict(self): - component = OpenAITextEmbedder(api_key="fake-api-key") - data = component.to_dict() - assert data == { - "type": "haystack.components.embedders.openai_text_embedder.OpenAITextEmbedder", - "init_parameters": { - "model_name": "text-embedding-ada-002", - "organization": None, - "prefix": "", - "suffix": "", - }, - } - - def test_to_dict_with_custom_init_parameters(self): - component = OpenAITextEmbedder( - api_key="fake-api-key", - model_name="model", - organization="fake-organization", - prefix="prefix", - suffix="suffix", - ) - data = component.to_dict() - assert data == { - "type": "haystack.components.embedders.openai_text_embedder.OpenAITextEmbedder", - "init_parameters": { - "model_name": "model", - "organization": "fake-organization", - "prefix": "prefix", - "suffix": "suffix", - }, - } - - def test_run(self): - model = "text-similarity-ada-001" - - with patch("haystack.components.embedders.openai_text_embedder.openai.Embedding") as openai_embedding_patch: - openai_embedding_patch.create.side_effect = mock_openai_response - - embedder = OpenAITextEmbedder(api_key="fake-api-key", model_name=model, prefix="prefix ", suffix=" suffix") - result = embedder.run(text="The food was delicious") - - openai_embedding_patch.create.assert_called_once_with( - model=model, input="prefix The food was delicious suffix" - ) - - assert len(result["embedding"]) == 1536 - assert all(isinstance(x, float) for x in result["embedding"]) - assert result["metadata"] == {"model": model, "usage": {"prompt_tokens": 4, "total_tokens": 4}} - - def test_run_wrong_input_format(self): - embedder = OpenAITextEmbedder(api_key="fake-api-key") - - list_integers_input = [1, 2, 3] - - with pytest.raises(TypeError, match="OpenAITextEmbedder expects a string as an input"): - embedder.run(text=list_integers_input) diff --git a/test/components/embedders/test_sentence_transformers_document_embedder.py b/test/components/embedders/test_sentence_transformers_document_embedder.py deleted file mode 100644 index 66c2304d9e..0000000000 --- a/test/components/embedders/test_sentence_transformers_document_embedder.py +++ /dev/null @@ -1,198 +0,0 @@ -from unittest.mock import patch, MagicMock -import pytest -import numpy as np - -from haystack import Document -from haystack.components.embedders.sentence_transformers_document_embedder import SentenceTransformersDocumentEmbedder - - -class TestSentenceTransformersDocumentEmbedder: - def test_init_default(self): - embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") - assert embedder.model_name_or_path == "model" - assert embedder.device == "cpu" - assert embedder.token is None - assert embedder.prefix == "" - assert embedder.suffix == "" - assert embedder.batch_size == 32 - assert embedder.progress_bar is True - assert embedder.normalize_embeddings is False - assert embedder.metadata_fields_to_embed == [] - assert embedder.embedding_separator == "\n" - - def test_init_with_parameters(self): - embedder = SentenceTransformersDocumentEmbedder( - model_name_or_path="model", - device="cuda", - token=True, - prefix="prefix", - suffix="suffix", - batch_size=64, - progress_bar=False, - normalize_embeddings=True, - metadata_fields_to_embed=["test_field"], - embedding_separator=" | ", - ) - assert embedder.model_name_or_path == "model" - assert embedder.device == "cuda" - assert embedder.token is True - assert embedder.prefix == "prefix" - assert embedder.suffix == "suffix" - assert embedder.batch_size == 64 - assert embedder.progress_bar is False - assert embedder.normalize_embeddings is True - assert embedder.metadata_fields_to_embed == ["test_field"] - assert embedder.embedding_separator == " | " - - def test_to_dict(self): - component = SentenceTransformersDocumentEmbedder(model_name_or_path="model") - data = component.to_dict() - assert data == { - "type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", - "init_parameters": { - "model_name_or_path": "model", - "device": "cpu", - "token": None, - "prefix": "", - "suffix": "", - "batch_size": 32, - "progress_bar": True, - "normalize_embeddings": False, - "embedding_separator": "\n", - "metadata_fields_to_embed": [], - }, - } - - def test_to_dict_with_custom_init_parameters(self): - component = SentenceTransformersDocumentEmbedder( - model_name_or_path="model", - device="cuda", - token="the-token", - prefix="prefix", - suffix="suffix", - batch_size=64, - progress_bar=False, - normalize_embeddings=True, - metadata_fields_to_embed=["meta_field"], - embedding_separator=" - ", - ) - data = component.to_dict() - - assert data == { - "type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", - "init_parameters": { - "model_name_or_path": "model", - "device": "cuda", - "token": None, # the token is not serialized - "prefix": "prefix", - "suffix": "suffix", - "batch_size": 64, - "progress_bar": False, - "normalize_embeddings": True, - "embedding_separator": " - ", - "metadata_fields_to_embed": ["meta_field"], - }, - } - - @patch( - "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory" - ) - def test_warmup(self, mocked_factory): - embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") - mocked_factory.get_embedding_backend.assert_not_called() - embedder.warm_up() - mocked_factory.get_embedding_backend.assert_called_once_with( - model_name_or_path="model", device="cpu", use_auth_token=None - ) - - @patch( - "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory" - ) - def test_warmup_doesnt_reload(self, mocked_factory): - embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") - mocked_factory.get_embedding_backend.assert_not_called() - embedder.warm_up() - embedder.warm_up() - mocked_factory.get_embedding_backend.assert_called_once() - - def test_run(self): - embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") - embedder.embedding_backend = MagicMock() - embedder.embedding_backend.embed = lambda x, **kwargs: np.random.rand(len(x), 16).tolist() - - documents = [Document(content=f"document number {i}") for i in range(5)] - - result = embedder.run(documents=documents) - - assert isinstance(result["documents"], list) - assert len(result["documents"]) == len(documents) - for doc in result["documents"]: - assert isinstance(doc, Document) - assert isinstance(doc.embedding, list) - assert isinstance(doc.embedding[0], float) - - def test_run_wrong_input_format(self): - embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model") - - string_input = "text" - list_integers_input = [1, 2, 3] - - with pytest.raises( - TypeError, match="SentenceTransformersDocumentEmbedder expects a list of Documents as input" - ): - embedder.run(documents=string_input) - - with pytest.raises( - TypeError, match="SentenceTransformersDocumentEmbedder expects a list of Documents as input" - ): - embedder.run(documents=list_integers_input) - - def test_embed_metadata(self): - embedder = SentenceTransformersDocumentEmbedder( - model_name_or_path="model", metadata_fields_to_embed=["meta_field"], embedding_separator="\n" - ) - embedder.embedding_backend = MagicMock() - - documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)] - - embedder.run(documents=documents) - - embedder.embedding_backend.embed.assert_called_once_with( - [ - "meta_value 0\ndocument number 0", - "meta_value 1\ndocument number 1", - "meta_value 2\ndocument number 2", - "meta_value 3\ndocument number 3", - "meta_value 4\ndocument number 4", - ], - batch_size=32, - show_progress_bar=True, - normalize_embeddings=False, - ) - - def test_prefix_suffix(self): - embedder = SentenceTransformersDocumentEmbedder( - model_name_or_path="model", - prefix="my_prefix ", - suffix=" my_suffix", - metadata_fields_to_embed=["meta_field"], - embedding_separator="\n", - ) - embedder.embedding_backend = MagicMock() - - documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)] - - embedder.run(documents=documents) - - embedder.embedding_backend.embed.assert_called_once_with( - [ - "my_prefix meta_value 0\ndocument number 0 my_suffix", - "my_prefix meta_value 1\ndocument number 1 my_suffix", - "my_prefix meta_value 2\ndocument number 2 my_suffix", - "my_prefix meta_value 3\ndocument number 3 my_suffix", - "my_prefix meta_value 4\ndocument number 4 my_suffix", - ], - batch_size=32, - show_progress_bar=True, - normalize_embeddings=False, - ) diff --git a/test/components/embedders/test_sentence_transformers_embedding_backend.py b/test/components/embedders/test_sentence_transformers_embedding_backend.py deleted file mode 100644 index 64f96cd734..0000000000 --- a/test/components/embedders/test_sentence_transformers_embedding_backend.py +++ /dev/null @@ -1,39 +0,0 @@ -from unittest.mock import patch -import pytest -from haystack.components.embedders.backends.sentence_transformers_backend import ( - _SentenceTransformersEmbeddingBackendFactory, -) - - -@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") -def test_factory_behavior(mock_sentence_transformer): - embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend( - model_name_or_path="my_model", device="cpu" - ) - same_embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend("my_model", "cpu") - another_embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend( - model_name_or_path="another_model", device="cpu" - ) - - assert same_embedding_backend is embedding_backend - assert another_embedding_backend is not embedding_backend - - -@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") -def test_model_initialization(mock_sentence_transformer): - _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend( - model_name_or_path="model", device="cpu", use_auth_token="my_token" - ) - mock_sentence_transformer.assert_called_once_with( - model_name_or_path="model", device="cpu", use_auth_token="my_token" - ) - - -@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer") -def test_embedding_function_with_kwargs(mock_sentence_transformer): - embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(model_name_or_path="model") - - data = ["sentence1", "sentence2"] - embedding_backend.embed(data=data, normalize_embeddings=True) - - embedding_backend.model.encode.assert_called_once_with(data, normalize_embeddings=True) diff --git a/test/components/embedders/test_sentence_transformers_text_embedder.py b/test/components/embedders/test_sentence_transformers_text_embedder.py deleted file mode 100644 index ba763bfb91..0000000000 --- a/test/components/embedders/test_sentence_transformers_text_embedder.py +++ /dev/null @@ -1,142 +0,0 @@ -from unittest.mock import patch, MagicMock -import pytest - -import numpy as np - -from haystack.components.embedders.sentence_transformers_text_embedder import SentenceTransformersTextEmbedder - - -class TestSentenceTransformersTextEmbedder: - def test_init_default(self): - embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") - assert embedder.model_name_or_path == "model" - assert embedder.device == "cpu" - assert embedder.token is None - assert embedder.prefix == "" - assert embedder.suffix == "" - assert embedder.batch_size == 32 - assert embedder.progress_bar is True - assert embedder.normalize_embeddings is False - - def test_init_with_parameters(self): - embedder = SentenceTransformersTextEmbedder( - model_name_or_path="model", - device="cuda", - token=True, - prefix="prefix", - suffix="suffix", - batch_size=64, - progress_bar=False, - normalize_embeddings=True, - ) - assert embedder.model_name_or_path == "model" - assert embedder.device == "cuda" - assert embedder.token is True - assert embedder.prefix == "prefix" - assert embedder.suffix == "suffix" - assert embedder.batch_size == 64 - assert embedder.progress_bar is False - assert embedder.normalize_embeddings is True - - def test_to_dict(self): - component = SentenceTransformersTextEmbedder(model_name_or_path="model") - data = component.to_dict() - assert data == { - "type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", - "init_parameters": { - "model_name_or_path": "model", - "device": "cpu", - "token": None, - "prefix": "", - "suffix": "", - "batch_size": 32, - "progress_bar": True, - "normalize_embeddings": False, - }, - } - - def test_to_dict_with_custom_init_parameters(self): - component = SentenceTransformersTextEmbedder( - model_name_or_path="model", - device="cuda", - token=True, - prefix="prefix", - suffix="suffix", - batch_size=64, - progress_bar=False, - normalize_embeddings=True, - ) - data = component.to_dict() - assert data == { - "type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", - "init_parameters": { - "model_name_or_path": "model", - "device": "cuda", - "token": True, - "prefix": "prefix", - "suffix": "suffix", - "batch_size": 64, - "progress_bar": False, - "normalize_embeddings": True, - }, - } - - def test_to_dict_not_serialize_token(self): - component = SentenceTransformersTextEmbedder(model_name_or_path="model", token="awesome-token") - data = component.to_dict() - assert data == { - "type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", - "init_parameters": { - "model_name_or_path": "model", - "device": "cpu", - "token": None, - "prefix": "", - "suffix": "", - "batch_size": 32, - "progress_bar": True, - "normalize_embeddings": False, - }, - } - - @patch( - "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" - ) - def test_warmup(self, mocked_factory): - embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") - mocked_factory.get_embedding_backend.assert_not_called() - embedder.warm_up() - mocked_factory.get_embedding_backend.assert_called_once_with( - model_name_or_path="model", device="cpu", use_auth_token=None - ) - - @patch( - "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" - ) - def test_warmup_doesnt_reload(self, mocked_factory): - embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") - mocked_factory.get_embedding_backend.assert_not_called() - embedder.warm_up() - embedder.warm_up() - mocked_factory.get_embedding_backend.assert_called_once() - - def test_run(self): - embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") - embedder.embedding_backend = MagicMock() - embedder.embedding_backend.embed = lambda x, **kwargs: np.random.rand(len(x), 16).tolist() - - text = "a nice text to embed" - - result = embedder.run(text=text) - embedding = result["embedding"] - - assert isinstance(embedding, list) - assert all(isinstance(el, float) for el in embedding) - - def test_run_wrong_input_format(self): - embedder = SentenceTransformersTextEmbedder(model_name_or_path="model") - embedder.embedding_backend = MagicMock() - - list_integers_input = [1, 2, 3] - - with pytest.raises(TypeError, match="SentenceTransformersTextEmbedder expects a string as input"): - embedder.run(text=list_integers_input) From f2df5acb333f017c9d2217e4c23913d5308bf9df Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Fri, 1 Dec 2023 20:04:20 +0100 Subject: [PATCH 12/17] create package structure --- integrations/gradient/LICENSE.txt | 9 + integrations/gradient/README.md | 21 +++ integrations/gradient/generator/__init__.py | 0 integrations/gradient/pyproject.toml | 166 ++++++++++++++++++ .../src/gradient_haystack/__about__.py | 4 + .../src/gradient_haystack/__init__.py | 3 + .../gradient_haystack/embedders/__init__.py | 3 + .../embedders/gradient_document_embedder.py | 14 +- .../embedders/gradient_text_embedder.py | 13 +- .../gradient_haystack/generator/__init__.py | 3 + .../gradient_haystack}/generator/base.py | 11 +- integrations/gradient/tests/__init__.py | 3 + .../test_gradient_document_embedder.py | 16 +- .../test_gradient_text_embedder.py | 9 +- 14 files changed, 245 insertions(+), 30 deletions(-) create mode 100644 integrations/gradient/LICENSE.txt create mode 100644 integrations/gradient/README.md delete mode 100644 integrations/gradient/generator/__init__.py create mode 100644 integrations/gradient/pyproject.toml create mode 100644 integrations/gradient/src/gradient_haystack/__about__.py create mode 100644 integrations/gradient/src/gradient_haystack/__init__.py create mode 100644 integrations/gradient/src/gradient_haystack/embedders/__init__.py rename integrations/gradient/{ => src/gradient_haystack}/embedders/gradient_document_embedder.py (89%) rename integrations/gradient/{ => src/gradient_haystack}/embedders/gradient_text_embedder.py (86%) create mode 100644 integrations/gradient/src/gradient_haystack/generator/__init__.py rename integrations/gradient/{ => src/gradient_haystack}/generator/base.py (94%) create mode 100644 integrations/gradient/tests/__init__.py rename integrations/gradient/{test => tests}/test_gradient_document_embedder.py (95%) rename integrations/gradient/{test => tests}/test_gradient_text_embedder.py (96%) diff --git a/integrations/gradient/LICENSE.txt b/integrations/gradient/LICENSE.txt new file mode 100644 index 0000000000..cf4129e2b3 --- /dev/null +++ b/integrations/gradient/LICENSE.txt @@ -0,0 +1,9 @@ +MIT License + +Copyright (c) 2023-present Massimiliano Pippi + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/integrations/gradient/README.md b/integrations/gradient/README.md new file mode 100644 index 0000000000..853f5b4c33 --- /dev/null +++ b/integrations/gradient/README.md @@ -0,0 +1,21 @@ +# gradient-haystack + +[![PyPI - Version](https://img.shields.io/pypi/v/gradient-haystack.svg)](https://pypi.org/project/gradient-haystack) +[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/gradient-haystack.svg)](https://pypi.org/project/gradient-haystack) + +----- + +**Table of Contents** + +- [Installation](#installation) +- [License](#license) + +## Installation + +```console +pip install gradient-haystack +``` + +## License + +`gradient-haystack` is distributed under the terms of the [MIT](https://spdx.org/licenses/MIT.html) license. diff --git a/integrations/gradient/generator/__init__.py b/integrations/gradient/generator/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/integrations/gradient/pyproject.toml b/integrations/gradient/pyproject.toml new file mode 100644 index 0000000000..c05a6ad9cf --- /dev/null +++ b/integrations/gradient/pyproject.toml @@ -0,0 +1,166 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "gradient-haystack" +dynamic = ["version"] +description = '' +readme = "README.md" +requires-python = ">=3.7" +license = "MIT" +keywords = [] +authors = [ + { name = "Massimiliano Pippi", email = "mpippi@gmail.com" }, +] +classifiers = [ + "Development Status :: 4 - Beta", + "Programming Language :: Python", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: Implementation :: CPython", + "Programming Language :: Python :: Implementation :: PyPy", +] +dependencies = [] + +[project.urls] +Documentation = "https://github.com/unknown/gradient-haystack#readme" +Issues = "https://github.com/unknown/gradient-haystack/issues" +Source = "https://github.com/unknown/gradient-haystack" + +[tool.hatch.version] +path = "src/gradient_haystack/__about__.py" + +[tool.hatch.envs.default] +dependencies = [ + "coverage[toml]>=6.5", + "pytest", +] +[tool.hatch.envs.default.scripts] +test = "pytest {args:tests}" +test-cov = "coverage run -m pytest {args:tests}" +cov-report = [ + "- coverage combine", + "coverage report", +] +cov = [ + "test-cov", + "cov-report", +] + +[[tool.hatch.envs.all.matrix]] +python = ["3.8", "3.9", "3.10", "3.11"] + +[tool.hatch.envs.lint] +detached = true +dependencies = [ + "black>=23.1.0", + "mypy>=1.0.0", + "ruff>=0.0.243", +] +[tool.hatch.envs.lint.scripts] +typing = "mypy --install-types --non-interactive {args:src/gradient_haystack tests}" +style = [ + "ruff {args:.}", + "black --check --diff {args:.}", +] +fmt = [ + "black {args:.}", + "ruff --fix {args:.}", + "style", +] +all = [ + "style", + "typing", +] + +[tool.black] +target-version = ["py38"] +line-length = 120 +skip-string-normalization = true + +[tool.ruff] +target-version = "py38" +line-length = 120 +select = [ + "A", + "ARG", + "B", + "C", + "DTZ", + "E", + "EM", + "F", + "FBT", + "I", + "ICN", + "ISC", + "N", + "PLC", + "PLE", + "PLR", + "PLW", + "Q", + "RUF", + "S", + "T", + "TID", + "UP", + "W", + "YTT", +] +ignore = [ + # Allow non-abstract empty methods in abstract base classes + "B027", + # Allow boolean positional values in function calls, like `dict.get(... True)` + "FBT003", + # Ignore checks for possible passwords + "S105", "S106", "S107", + # Ignore complexity + "C901", "PLR0911", "PLR0912", "PLR0913", "PLR0915", +] +unfixable = [ + # Don't touch unused imports + "F401", +] + +[tool.ruff.isort] +known-first-party = ["gradient_haystack"] + +[tool.ruff.flake8-tidy-imports] +ban-relative-imports = "all" + +[tool.ruff.per-file-ignores] +# Tests can use magic values, assertions, and relative imports +"tests/**/*" = ["PLR2004", "S101", "TID252"] + +[tool.coverage.run] +source_pkgs = ["gradient_haystack", "tests"] +branch = true +parallel = true +omit = [ + "src/gradient_haystack/__about__.py", +] + +[tool.coverage.paths] +gradient_haystack = ["src/gradient_haystack", "*/gradient-haystack/src/gradient_haystack"] +tests = ["tests", "*/gradient-haystack/tests"] + +[tool.coverage.report] +exclude_lines = [ + "no cov", + "if __name__ == .__main__.:", + "if TYPE_CHECKING:", +] + +[[tool.mypy.overrides]] +module = [ + "gradientai.*", + "haystack.*", + "pytest.*", + "numpy.*", +] +ignore_missing_imports = true \ No newline at end of file diff --git a/integrations/gradient/src/gradient_haystack/__about__.py b/integrations/gradient/src/gradient_haystack/__about__.py new file mode 100644 index 0000000000..132530b415 --- /dev/null +++ b/integrations/gradient/src/gradient_haystack/__about__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi +# +# SPDX-License-Identifier: MIT +__version__ = "0.0.1" diff --git a/integrations/gradient/src/gradient_haystack/__init__.py b/integrations/gradient/src/gradient_haystack/__init__.py new file mode 100644 index 0000000000..bd78f6a287 --- /dev/null +++ b/integrations/gradient/src/gradient_haystack/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi +# +# SPDX-License-Identifier: MIT diff --git a/integrations/gradient/src/gradient_haystack/embedders/__init__.py b/integrations/gradient/src/gradient_haystack/embedders/__init__.py new file mode 100644 index 0000000000..bd78f6a287 --- /dev/null +++ b/integrations/gradient/src/gradient_haystack/embedders/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi +# +# SPDX-License-Identifier: MIT diff --git a/integrations/gradient/embedders/gradient_document_embedder.py b/integrations/gradient/src/gradient_haystack/embedders/gradient_document_embedder.py similarity index 89% rename from integrations/gradient/embedders/gradient_document_embedder.py rename to integrations/gradient/src/gradient_haystack/embedders/gradient_document_embedder.py index e2d3458c59..81a93ad2b1 100644 --- a/integrations/gradient/embedders/gradient_document_embedder.py +++ b/integrations/gradient/src/gradient_haystack/embedders/gradient_document_embedder.py @@ -1,7 +1,7 @@ import logging -from typing import List, Optional, Dict, Any +from typing import Any, Dict, List, Optional -from haystack import component, Document, default_to_dict +from haystack import Document, component, default_to_dict from haystack.lazy_imports import LazyImport with LazyImport(message="Run 'pip install gradientai'") as gradientai_import: @@ -97,13 +97,13 @@ def run(self, documents: List[Document]): :param documents: A list of Documents to embed. """ if not isinstance(documents, list) or documents and any(not isinstance(doc, Document) for doc in documents): - raise TypeError( - "GradientDocumentEmbedder expects a list of Documents as input." - "In case you want to embed a list of strings, please use the GradientTextEmbedder." - ) + msg = "GradientDocumentEmbedder expects a list of Documents as input.\ + In case you want to embed a list of strings, please use the GradientTextEmbedder." + raise TypeError(msg) if not hasattr(self, "_embedding_model"): - raise RuntimeError("The embedding model has not been loaded. Please call warm_up() before running.") + msg = "The embedding model has not been loaded. Please call warm_up() before running." + raise RuntimeError(msg) embeddings = self._generate_embeddings(documents=documents, batch_size=self._batch_size) for doc, embedding in zip(documents, embeddings): diff --git a/integrations/gradient/embedders/gradient_text_embedder.py b/integrations/gradient/src/gradient_haystack/embedders/gradient_text_embedder.py similarity index 86% rename from integrations/gradient/embedders/gradient_text_embedder.py rename to integrations/gradient/src/gradient_haystack/embedders/gradient_text_embedder.py index 28a83a52e2..53996b785e 100644 --- a/integrations/gradient/embedders/gradient_text_embedder.py +++ b/integrations/gradient/src/gradient_haystack/embedders/gradient_text_embedder.py @@ -72,17 +72,18 @@ def warm_up(self) -> None: def run(self, text: str): """Generates an embedding for a single text.""" if not isinstance(text, str): - raise TypeError( - "GradientTextEmbedder expects a string as an input." - "In case you want to embed a list of Documents, please use the GradientDocumentEmbedder." - ) + msg = "GradientTextEmbedder expects a string as an input.\ + In case you want to embed a list of Documents, please use the GradientDocumentEmbedder." + raise TypeError(msg) if not hasattr(self, "_embedding_model"): - raise RuntimeError("The embedding model has not been loaded. Please call warm_up() before running.") + msg = "The embedding model has not been loaded. Please call warm_up() before running." + raise RuntimeError(msg) result = self._embedding_model.generate_embeddings(inputs=[{"input": text}]) if (not result) or (result.embeddings is None) or (len(result.embeddings) == 0): - raise RuntimeError("The embedding model did not return any embeddings.") + msg = "The embedding model did not return any embeddings." + raise RuntimeError(msg) return {"embedding": result.embeddings[0].embedding} diff --git a/integrations/gradient/src/gradient_haystack/generator/__init__.py b/integrations/gradient/src/gradient_haystack/generator/__init__.py new file mode 100644 index 0000000000..bd78f6a287 --- /dev/null +++ b/integrations/gradient/src/gradient_haystack/generator/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi +# +# SPDX-License-Identifier: MIT diff --git a/integrations/gradient/generator/base.py b/integrations/gradient/src/gradient_haystack/generator/base.py similarity index 94% rename from integrations/gradient/generator/base.py rename to integrations/gradient/src/gradient_haystack/generator/base.py index 4853a92291..536525377d 100644 --- a/integrations/gradient/generator/base.py +++ b/integrations/gradient/src/gradient_haystack/generator/base.py @@ -1,9 +1,8 @@ -from typing import List, Optional, Dict, Any, overload - import logging -from haystack.lazy_imports import LazyImport +from typing import Any, Dict, List, Optional from haystack import component, default_to_dict +from haystack.lazy_imports import LazyImport with LazyImport(message="Run 'pip install gradientai'") as gradientai_import: from gradientai import Gradient @@ -74,9 +73,11 @@ def __init__( has_model_adapter_id = model_adapter_id is not None and model_adapter_id != "" if not has_base_model_slug and not has_model_adapter_id: - raise ValueError("Either base_model_slug or model_adapter_id must be provided.") + msg = "Either base_model_slug or model_adapter_id must be provided." + raise ValueError(msg) if has_base_model_slug and has_model_adapter_id: - raise ValueError("Only one of base_model_slug or model_adapter_id must be provided.") + msg = "Only one of base_model_slug or model_adapter_id must be provided." + raise ValueError(msg) if has_base_model_slug: self._base_model_slug = base_model_slug diff --git a/integrations/gradient/tests/__init__.py b/integrations/gradient/tests/__init__.py new file mode 100644 index 0000000000..bd78f6a287 --- /dev/null +++ b/integrations/gradient/tests/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi +# +# SPDX-License-Identifier: MIT diff --git a/integrations/gradient/test/test_gradient_document_embedder.py b/integrations/gradient/tests/test_gradient_document_embedder.py similarity index 95% rename from integrations/gradient/test/test_gradient_document_embedder.py rename to integrations/gradient/tests/test_gradient_document_embedder.py index ffb44519fd..bac02df5e0 100644 --- a/integrations/gradient/test/test_gradient_document_embedder.py +++ b/integrations/gradient/tests/test_gradient_document_embedder.py @@ -1,11 +1,11 @@ -import pytest -from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess -from haystack.components.embedders.gradient_document_embedder import GradientDocumentEmbedder from unittest.mock import MagicMock, NonCallableMagicMock -import numpy as np +import numpy as np +import pytest +from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess from haystack import Document +from gradient_haystack.embedders.gradient_document_embedder import GradientDocumentEmbedder access_token = "access_token" workspace_id = "workspace_id" @@ -59,7 +59,7 @@ def test_to_dict(self): component = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) data = component.to_dict() assert data == { - "type": "haystack.components.embedders.gradient_document_embedder.GradientDocumentEmbedder", + "type": "gradient_haystack.embedders.gradient_document_embedder.GradientDocumentEmbedder", "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, } @@ -131,12 +131,12 @@ def test_run_custom_batch(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id, batch_size=20) embedder._embedding_model = NonCallableMagicMock() - DOCUMENT_COUNT = 101 + document_count = 101 embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( - embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(DOCUMENT_COUNT)] + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(document_count)] ) - documents = [Document(content=f"document number {i}") for i in range(DOCUMENT_COUNT)] + documents = [Document(content=f"document number {i}") for i in range(document_count)] result = embedder.run(documents=documents) diff --git a/integrations/gradient/test/test_gradient_text_embedder.py b/integrations/gradient/tests/test_gradient_text_embedder.py similarity index 96% rename from integrations/gradient/test/test_gradient_text_embedder.py rename to integrations/gradient/tests/test_gradient_text_embedder.py index 82c84e3904..9623db5d44 100644 --- a/integrations/gradient/test/test_gradient_text_embedder.py +++ b/integrations/gradient/tests/test_gradient_text_embedder.py @@ -1,9 +1,10 @@ -import pytest -from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess -from haystack.components.embedders.gradient_text_embedder import GradientTextEmbedder from unittest.mock import MagicMock, NonCallableMagicMock + import numpy as np +import pytest +from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess +from gradient_haystack.embedders.gradient_text_embedder import GradientTextEmbedder access_token = "access_token" workspace_id = "workspace_id" @@ -57,7 +58,7 @@ def test_to_dict(self): component = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) data = component.to_dict() assert data == { - "type": "haystack.components.embedders.gradient_text_embedder.GradientTextEmbedder", + "type": "gradient_haystack.embedders.gradient_text_embedder.GradientTextEmbedder", "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, } From c793e6c5f604bc7922be021508b37277cc3302fc Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Fri, 1 Dec 2023 20:06:20 +0100 Subject: [PATCH 13/17] run tests in CI --- .github/workflows/gradient.yml | 56 ++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 .github/workflows/gradient.yml diff --git a/.github/workflows/gradient.yml b/.github/workflows/gradient.yml new file mode 100644 index 0000000000..f717ba2c98 --- /dev/null +++ b/.github/workflows/gradient.yml @@ -0,0 +1,56 @@ +# This workflow comes from https://github.com/ofek/hatch-mypyc +# https://github.com/ofek/hatch-mypyc/blob/5a198c0ba8660494d02716cfc9d79ce4adfb1442/.github/workflows/test.yml +name: Test / gradient + +on: + schedule: + - cron: "0 0 * * *" + pull_request: + paths: + - 'integrations/gradient/**' + - '.github/workflows/gradient.yml' + +defaults: + run: + working-directory: integrations/gradient + +concurrency: + group: gradient-${{ github.head_ref }} + cancel-in-progress: true + +env: + PYTHONUNBUFFERED: "1" + FORCE_COLOR: "1" + +jobs: + run: + name: Python ${{ matrix.python-version }} on ${{ startsWith(matrix.os, 'macos-') && 'macOS' || startsWith(matrix.os, 'windows-') && 'Windows' || 'Linux' }} + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + python-version: ['3.9', '3.10'] + + steps: + - name: Support longpaths + if: matrix.os == 'windows-latest' + working-directory: . + run: git config --system core.longpaths true + + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + + - name: Install Hatch + run: pip install --upgrade hatch + + - name: Lint + if: matrix.python-version == '3.9' && runner.os == 'Linux' + run: hatch run lint:all + + - name: Run tests + run: hatch run cov \ No newline at end of file From 755888246290105bffa33f33c6f9abd26f8ecc99 Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Fri, 1 Dec 2023 20:11:14 +0100 Subject: [PATCH 14/17] fix project file --- integrations/gradient/pyproject.toml | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/integrations/gradient/pyproject.toml b/integrations/gradient/pyproject.toml index c05a6ad9cf..afdd8ecb8e 100644 --- a/integrations/gradient/pyproject.toml +++ b/integrations/gradient/pyproject.toml @@ -11,7 +11,7 @@ requires-python = ">=3.7" license = "MIT" keywords = [] authors = [ - { name = "Massimiliano Pippi", email = "mpippi@gmail.com" }, + { name = "Mateusz Haligowski", email = "mhaligowski@gmail.com" }, ] classifiers = [ "Development Status :: 4 - Beta", @@ -24,7 +24,10 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] -dependencies = [] +dependencies = [ + "haystack-ai", + "gradientai", +] [project.urls] Documentation = "https://github.com/unknown/gradient-haystack#readme" From d45acd424d68c5a6aef1dff0f3f0582952f9ed10 Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Fri, 1 Dec 2023 20:19:46 +0100 Subject: [PATCH 15/17] add integration test --- .../tests/test_gradient_rag_pipelines.py | 92 +++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 integrations/gradient/tests/test_gradient_rag_pipelines.py diff --git a/integrations/gradient/tests/test_gradient_rag_pipelines.py b/integrations/gradient/tests/test_gradient_rag_pipelines.py new file mode 100644 index 0000000000..f0a4db9901 --- /dev/null +++ b/integrations/gradient/tests/test_gradient_rag_pipelines.py @@ -0,0 +1,92 @@ +import os +import json + +import pytest +from haystack import Pipeline, Document +from haystack.document_stores import InMemoryDocumentStore +from haystack.components.writers import DocumentWriter +from haystack.components.retrievers import InMemoryEmbeddingRetriever +from haystack.components.builders.answer_builder import AnswerBuilder +from haystack.components.builders.prompt_builder import PromptBuilder + +from gradient_haystack.embedders.gradient_document_embedder import GradientDocumentEmbedder +from gradient_haystack.embedders.gradient_text_embedder import GradientTextEmbedder +from gradient_haystack.generator.base import GradientGenerator + + +@pytest.mark.skipif( + not os.environ.get("GRADIENT_ACCESS_TOKEN", None) or not os.environ.get("GRADIENT_WORKSPACE_ID", None), + reason="Export env variables called GRADIENT_ACCESS_TOKEN and GRADIENT_WORKSPACE_ID containing the Gradient configuration settings to run this test.", +) +def test_gradient_embedding_retrieval_rag_pipeline(tmp_path): + # Create the RAG pipeline + prompt_template = """ + Given these documents, answer the question.\nDocuments: + {% for doc in documents %} + {{ doc.content }} + {% endfor %} + \nQuestion: {{question}} + \nAnswer: + """ + + gradient_access_token = os.environ.get("GRADIENT_ACCESS_TOKEN") + rag_pipeline = Pipeline() + embedder = GradientTextEmbedder(access_token=gradient_access_token) + rag_pipeline.add_component(instance=embedder, name="text_embedder") + rag_pipeline.add_component( + instance=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()), name="retriever" + ) + rag_pipeline.add_component(instance=PromptBuilder(template=prompt_template), name="prompt_builder") + rag_pipeline.add_component( + instance=GradientGenerator(access_token=gradient_access_token, base_model_slug="llama2-7b-chat"), name="llm" + ) + rag_pipeline.add_component(instance=AnswerBuilder(), name="answer_builder") + rag_pipeline.connect("text_embedder", "retriever") + rag_pipeline.connect("retriever", "prompt_builder.documents") + rag_pipeline.connect("prompt_builder", "llm") + rag_pipeline.connect("llm.replies", "answer_builder.replies") + rag_pipeline.connect("retriever", "answer_builder.documents") + + # Draw the pipeline + rag_pipeline.draw(tmp_path / "test_gradient_embedding_rag_pipeline.png") + + # Serialize the pipeline to JSON + with open(tmp_path / "test_bm25_rag_pipeline.json", "w") as f: + json.dump(rag_pipeline.to_dict(), f) + + # Load the pipeline back + with open(tmp_path / "test_bm25_rag_pipeline.json", "r") as f: + rag_pipeline = Pipeline.from_dict(json.load(f)) + + # Populate the document store + documents = [ + Document(content="My name is Jean and I live in Paris."), + Document(content="My name is Mark and I live in Berlin."), + Document(content="My name is Giorgio and I live in Rome."), + ] + document_store = rag_pipeline.get_component("retriever").document_store + indexing_pipeline = Pipeline() + indexing_pipeline.add_component(instance=GradientDocumentEmbedder(), name="document_embedder") + indexing_pipeline.add_component(instance=DocumentWriter(document_store=document_store), name="document_writer") + indexing_pipeline.connect("document_embedder", "document_writer") + indexing_pipeline.run({"document_embedder": {"documents": documents}}) + + # Query and assert + questions = ["Who lives in Paris?", "Who lives in Berlin?", "Who lives in Rome?"] + answers_spywords = ["Jean", "Mark", "Giorgio"] + + for question, spyword in zip(questions, answers_spywords): + result = rag_pipeline.run( + { + "text_embedder": {"text": question}, + "prompt_builder": {"question": question}, + "answer_builder": {"query": question}, + } + ) + + assert len(result["answer_builder"]["answers"]) == 1 + generated_answer = result["answer_builder"]["answers"][0] + assert spyword in generated_answer.data + assert generated_answer.query == question + assert hasattr(generated_answer, "documents") + assert hasattr(generated_answer, "metadata") From 36e799176e0c2edcb1b3b6a54435434843c929ab Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Fri, 1 Dec 2023 13:48:42 -0800 Subject: [PATCH 16/17] fix complaining linter --- .../tests/test_gradient_rag_pipelines.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/integrations/gradient/tests/test_gradient_rag_pipelines.py b/integrations/gradient/tests/test_gradient_rag_pipelines.py index f0a4db9901..8173c09c40 100644 --- a/integrations/gradient/tests/test_gradient_rag_pipelines.py +++ b/integrations/gradient/tests/test_gradient_rag_pipelines.py @@ -1,13 +1,13 @@ -import os import json +import os import pytest -from haystack import Pipeline, Document -from haystack.document_stores import InMemoryDocumentStore -from haystack.components.writers import DocumentWriter -from haystack.components.retrievers import InMemoryEmbeddingRetriever +from haystack import Document, Pipeline from haystack.components.builders.answer_builder import AnswerBuilder from haystack.components.builders.prompt_builder import PromptBuilder +from haystack.components.retrievers import InMemoryEmbeddingRetriever +from haystack.components.writers import DocumentWriter +from haystack.document_stores import InMemoryDocumentStore from gradient_haystack.embedders.gradient_document_embedder import GradientDocumentEmbedder from gradient_haystack.embedders.gradient_text_embedder import GradientTextEmbedder @@ -16,7 +16,10 @@ @pytest.mark.skipif( not os.environ.get("GRADIENT_ACCESS_TOKEN", None) or not os.environ.get("GRADIENT_WORKSPACE_ID", None), - reason="Export env variables called GRADIENT_ACCESS_TOKEN and GRADIENT_WORKSPACE_ID containing the Gradient configuration settings to run this test.", + reason=( + "Export env variables called GRADIENT_ACCESS_TOKEN and GRADIENT_WORKSPACE_ID" + "containing the Gradient configuration settings to run this test.", + ), ) def test_gradient_embedding_retrieval_rag_pipeline(tmp_path): # Create the RAG pipeline @@ -55,7 +58,7 @@ def test_gradient_embedding_retrieval_rag_pipeline(tmp_path): json.dump(rag_pipeline.to_dict(), f) # Load the pipeline back - with open(tmp_path / "test_bm25_rag_pipeline.json", "r") as f: + with open(tmp_path / "test_bm25_rag_pipeline.json") as f: rag_pipeline = Pipeline.from_dict(json.load(f)) # Populate the document store From 7a28f8b70eab8ea2d7bdb15b83d04c3b1988da88 Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Sat, 2 Dec 2023 16:13:14 +0100 Subject: [PATCH 17/17] fix skip call --- integrations/gradient/tests/test_gradient_rag_pipelines.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/integrations/gradient/tests/test_gradient_rag_pipelines.py b/integrations/gradient/tests/test_gradient_rag_pipelines.py index 8173c09c40..5835944a88 100644 --- a/integrations/gradient/tests/test_gradient_rag_pipelines.py +++ b/integrations/gradient/tests/test_gradient_rag_pipelines.py @@ -16,10 +16,8 @@ @pytest.mark.skipif( not os.environ.get("GRADIENT_ACCESS_TOKEN", None) or not os.environ.get("GRADIENT_WORKSPACE_ID", None), - reason=( - "Export env variables called GRADIENT_ACCESS_TOKEN and GRADIENT_WORKSPACE_ID" - "containing the Gradient configuration settings to run this test.", - ), + reason="Export env variables called GRADIENT_ACCESS_TOKEN and GRADIENT_WORKSPACE_ID \ + containing the Gradient configuration settings to run this test.", ) def test_gradient_embedding_retrieval_rag_pipeline(tmp_path): # Create the RAG pipeline