diff --git a/.github/workflows/e2e_preview.yml b/.github/workflows/e2e_preview.yml index 2e9763d280..a695c8c34b 100644 --- a/.github/workflows/e2e_preview.yml +++ b/.github/workflows/e2e_preview.yml @@ -36,7 +36,7 @@ jobs: sudo apt install ffmpeg # for local Whisper tests - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Run tests run: pytest e2e/preview diff --git a/.github/workflows/linting_preview.yml b/.github/workflows/linting_preview.yml index 1c29984436..08ecc2645e 100644 --- a/.github/workflows/linting_preview.yml +++ b/.github/workflows/linting_preview.yml @@ -38,7 +38,7 @@ jobs: python-version: ${{ env.PYTHON_VERSION }} - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Mypy if: steps.files.outputs.any_changed == 'true' @@ -69,7 +69,7 @@ jobs: - name: Install Haystack run: | - pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' + pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai pip install ./haystack-linter - name: Pylint diff --git a/.github/workflows/tests_preview.yml b/.github/workflows/tests_preview.yml index a69bb369eb..076fc13744 100644 --- a/.github/workflows/tests_preview.yml +++ b/.github/workflows/tests_preview.yml @@ -116,7 +116,7 @@ jobs: python-version: ${{ env.PYTHON_VERSION }} - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Run run: pytest -m "not integration" test/preview @@ -174,7 +174,7 @@ jobs: sudo apt install ffmpeg # for local Whisper tests - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Run run: pytest --maxfail=5 -m "integration" test/preview @@ -230,7 +230,7 @@ jobs: colima start - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Run Tika run: docker run -d -p 9998:9998 apache/tika:2.9.0.0 @@ -281,7 +281,7 @@ jobs: python-version: ${{ env.PYTHON_VERSION }} - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Run run: pytest --maxfail=5 -m "integration" test/preview -k 'not tika' diff --git a/e2e/preview/pipelines/test_gradient_rag_pipelines.py b/e2e/preview/pipelines/test_gradient_rag_pipelines.py new file mode 100644 index 0000000000..147e7fb8e6 --- /dev/null +++ b/e2e/preview/pipelines/test_gradient_rag_pipelines.py @@ -0,0 +1,92 @@ +import os +import json +import pytest + +from haystack.preview import Pipeline, Document +from haystack.preview.components.embedders.gradient_document_embedder import GradientDocumentEmbedder +from haystack.preview.components.embedders.gradient_text_embedder import GradientTextEmbedder +from haystack.preview.document_stores import InMemoryDocumentStore +from haystack.preview.components.writers import DocumentWriter +from haystack.preview.components.retrievers import InMemoryEmbeddingRetriever +from haystack.preview.components.generators.gradient.base import GradientGenerator +from haystack.preview.components.builders.answer_builder import AnswerBuilder +from haystack.preview.components.builders.prompt_builder import PromptBuilder + + +@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") diff --git a/haystack/preview/components/embedders/__init__.py b/haystack/preview/components/embedders/__init__.py index a0840d7e0a..4af150d9ce 100644 --- a/haystack/preview/components/embedders/__init__.py +++ b/haystack/preview/components/embedders/__init__.py @@ -1,3 +1,6 @@ +from haystack.preview.components.embedders.gradient_text_embedder import GradientTextEmbedder +from haystack.preview.components.embedders.gradient_document_embedder import GradientDocumentEmbedder + from haystack.preview.components.embedders.sentence_transformers_text_embedder import SentenceTransformersTextEmbedder from haystack.preview.components.embedders.sentence_transformers_document_embedder import ( SentenceTransformersDocumentEmbedder, @@ -6,6 +9,8 @@ from haystack.preview.components.embedders.openai_text_embedder import OpenAITextEmbedder __all__ = [ + "GradientTextEmbedder", + "GradientDocumentEmbedder", "SentenceTransformersTextEmbedder", "SentenceTransformersDocumentEmbedder", "OpenAITextEmbedder", diff --git a/haystack/preview/components/embedders/gradient_document_embedder.py b/haystack/preview/components/embedders/gradient_document_embedder.py new file mode 100644 index 0000000000..c1e2fa71a6 --- /dev/null +++ b/haystack/preview/components/embedders/gradient_document_embedder.py @@ -0,0 +1,112 @@ +import logging +from typing import List, Optional, Dict, Any + +from haystack.preview import component, Document, default_to_dict +from haystack.preview.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/haystack/preview/components/embedders/gradient_text_embedder.py b/haystack/preview/components/embedders/gradient_text_embedder.py new file mode 100644 index 0000000000..8c89da08b4 --- /dev/null +++ b/haystack/preview/components/embedders/gradient_text_embedder.py @@ -0,0 +1,88 @@ +from typing import Any, Dict, List, Optional + +from haystack.preview import component, default_to_dict +from haystack.preview.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/haystack/preview/components/generators/__init__.py b/haystack/preview/components/generators/__init__.py index bc81975a3f..8230c46bf9 100644 --- a/haystack/preview/components/generators/__init__.py +++ b/haystack/preview/components/generators/__init__.py @@ -1,5 +1,6 @@ +from haystack.preview.components.generators.gradient.base import GradientGenerator from haystack.preview.components.generators.hugging_face_local import HuggingFaceLocalGenerator from haystack.preview.components.generators.hugging_face_tgi import HuggingFaceTGIGenerator from haystack.preview.components.generators.openai import GPTGenerator -__all__ = ["HuggingFaceLocalGenerator", "HuggingFaceTGIGenerator", "GPTGenerator"] +__all__ = ["GPTGenerator", "GradientGenerator", "HuggingFaceTGIGenerator", "HuggingFaceLocalGenerator"] diff --git a/haystack/preview/components/generators/gradient/__init__.py b/haystack/preview/components/generators/gradient/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/preview/components/generators/gradient/base.py new file mode 100644 index 0000000000..dd420a43f4 --- /dev/null +++ b/haystack/preview/components/generators/gradient/base.py @@ -0,0 +1,128 @@ +from typing import List, Optional, Dict, Any, overload + +import logging +from haystack.lazy_imports import LazyImport + +from haystack.preview 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/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml new file mode 100644 index 0000000000..355a3874a5 --- /dev/null +++ b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml @@ -0,0 +1,4 @@ +--- +preview: + - | + Adds integration with [Gradient AI](https://gradient.ai). diff --git a/test/preview/components/embedders/test_gradient_document_embedder.py b/test/preview/components/embedders/test_gradient_document_embedder.py new file mode 100644 index 0000000000..b114d382e2 --- /dev/null +++ b/test/preview/components/embedders/test_gradient_document_embedder.py @@ -0,0 +1,158 @@ +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 + +from haystack.preview import Document + + +access_token = "access_token" +workspace_id = "workspace_id" +model = "bge-large" + + +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": "haystack.preview.components.embedders.gradient_document_embedder.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(content=f"document number {i}") for i in range(5)]) + + @pytest.mark.unit + def test_run(self): + 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(content=f"document number {i}") for i in range(5)] + + 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"]: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert isinstance(doc.embedding[0], float) + + @pytest.mark.unit + def test_run_batch(self): + 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(content=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) + + @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) + embedder._embedding_model = NonCallableMagicMock() + + result = embedder.run(documents=[]) + + assert result["documents"] == [] diff --git a/test/preview/components/embedders/test_gradient_text_embedder.py b/test/preview/components/embedders/test_gradient_text_embedder.py new file mode 100644 index 0000000000..9eddff7eef --- /dev/null +++ b/test/preview/components/embedders/test_gradient_text_embedder.py @@ -0,0 +1,126 @@ +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 + + +access_token = "access_token" +workspace_id = "workspace_id" +model = "bge-large" + + +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": "haystack.preview.components.embedders.gradient_text_embedder.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_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( + 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"])