From aa05b97255868c1b9d9ed6c3e027a933259db56b Mon Sep 17 00:00:00 2001
From: lilwilsond2 <lilwilsond2@gmail.com>
Date: Fri, 13 Oct 2023 03:25:26 -0500
Subject: [PATCH 01/13] init generator

---
 .../generators/gradient/__init__.py           |   0
 .../components/generators/gradient/base.py    | 103 ++++++++++++++++++
 2 files changed, 103 insertions(+)
 create mode 100644 haystack/preview/components/generators/gradient/__init__.py
 create mode 100644 haystack/preview/components/generators/gradient/base.py

diff --git a/haystack/preview/components/generators/gradient/__init__.py b/haystack/preview/components/generators/gradient/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/preview/components/generators/gradient/base.py
new file mode 100644
index 00000000000..41084d04c04
--- /dev/null
+++ b/haystack/preview/components/generators/gradient/base.py
@@ -0,0 +1,103 @@
+from typing import 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:
+    @overload
+    def __init__(
+        self,
+        *,
+        access_token: Optional[str] = None,
+        base_model_slug: str,
+        host: Optional[str] = None,
+        max_generated_token_count: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        workspace_id: Optional[str] = None,
+    ) -> None:
+        ...
+
+    @overload
+    def __init__(
+        self,
+        *,
+        access_token: Optional[str] = None,
+        host: Optional[str] = None,
+        max_generated_token_count: Optional[int] = None,
+        model_adapter_id: str,
+        temperature: Optional[float] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        workspace_id: Optional[str] = None,
+    ) -> None:
+        ...
+
+    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:
+        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
+
+        if (base_model_slug is None and model_adapter_id is None) or (
+            isinstance(base_model_slug, str) and isinstance(model_adapter_id, str)
+        ):
+            raise ValueError("expected be provided exactly one of base_model_slug or model_adapter_id")
+
+        self._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id)
+        if isinstance(base_model_slug, str):
+            self._model = self._gradient.get_base_model(base_model_slug=base_model_slug)
+        if isinstance(model_adapter_id, str):
+            self._model = self._gradient.get_model_adapter(model_adapter_id=model_adapter_id)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return default_to_dict(
+            self,
+            access_token=self._access_token,
+            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,
+        )
+
+    @component.output_types(response=str)
+    def run(self, prompt: str):
+        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 {"response": resp.generated_output}

From 86299e94d013e660987203edb242ce69dc723b03 Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <mhaligowski@gmail.com>
Date: Thu, 26 Oct 2023 14:05:29 -0700
Subject: [PATCH 02/13] Gradient embeddings (#2)

---
 e2e/pipelines/test_gradient_rag_pipelines.py  |  92 ++++++++++++
 .../embedders/gradient_document_embedder.py   |  96 ++++++++++++
 .../embedders/gradient_text_embedder.py       |  71 +++++++++
 .../components/generators/gradient/base.py    |  10 +-
 .../test_gradient_document_embedder.py        | 139 ++++++++++++++++++
 .../embedders/test_gradient_text_embedder.py  | 111 ++++++++++++++
 6 files changed, 514 insertions(+), 5 deletions(-)
 create mode 100644 e2e/pipelines/test_gradient_rag_pipelines.py
 create mode 100644 haystack/components/embedders/gradient_document_embedder.py
 create mode 100644 haystack/components/embedders/gradient_text_embedder.py
 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/e2e/pipelines/test_gradient_rag_pipelines.py b/e2e/pipelines/test_gradient_rag_pipelines.py
new file mode 100644
index 00000000000..ca2113bbb6d
--- /dev/null
+++ b/e2e/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),
+    reason="Export an env var called GRADIENT_ACCESS_TOKEN containing the Gradient access token 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.text }}
+    {% 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(text="My name is Jean and I live in Paris."),
+        Document(text="My name is Mark and I live in Berlin."),
+        Document(text="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/components/embedders/gradient_document_embedder.py b/haystack/components/embedders/gradient_document_embedder.py
new file mode 100644
index 00000000000..7fd6084ca51
--- /dev/null
+++ b/haystack/components/embedders/gradient_document_embedder.py
@@ -0,0 +1,96 @@
+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.
+    """
+
+    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 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/.
+        """
+        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=100) -> 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.text} 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 not isinstance(documents[0], Document):
+            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)
+        for doc, embedding in zip(documents, embeddings):
+            doc.embedding = embedding
+
+        return {"documents": documents}
diff --git a/haystack/components/embedders/gradient_text_embedder.py b/haystack/components/embedders/gradient_text_embedder.py
new file mode 100644
index 00000000000..39f28f51952
--- /dev/null
+++ b/haystack/components/embedders/gradient_text_embedder.py
@@ -0,0 +1,71 @@
+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).
+    """
+
+    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/.
+        """
+        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}])
+        return {"embedding": result.embeddings[0].embedding}
diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/preview/components/generators/gradient/base.py
index 41084d04c04..199742096a6 100644
--- a/haystack/preview/components/generators/gradient/base.py
+++ b/haystack/preview/components/generators/gradient/base.py
@@ -1,4 +1,4 @@
-from typing import Optional, Dict, Any, overload
+from typing import List, Optional, Dict, Any, overload
 
 import logging
 from haystack.lazy_imports import LazyImport
@@ -57,10 +57,10 @@ def __init__(
         workspace_id: Optional[str] = None,
     ) -> None:
         self._access_token = access_token
-        self._base_model_slug = (base_model_slug,)
+        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._model_adapter_id = model_adapter_id
         self._temperature = temperature
         self._top_k = top_k
         self._top_p = top_p
@@ -91,7 +91,7 @@ def to_dict(self) -> Dict[str, Any]:
             workspace_id=self._workspace_id,
         )
 
-    @component.output_types(response=str)
+    @component.output_types(replies=List[str])
     def run(self, prompt: str):
         resp = self._model.complete(
             query=prompt,
@@ -100,4 +100,4 @@ def run(self, prompt: str):
             top_k=self._top_k,
             top_p=self._top_p,
         )
-        return {"response": resp.generated_output}
+        return {"replies": [resp.generated_output]}
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 00000000000..ba9568270f6
--- /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 00000000000..3e2a6762f27
--- /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 e302ccb8bb1e0b22bc6f8e401912f28d94135689 Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <mhaligowski@gmail.com>
Date: Fri, 27 Oct 2023 13:31:11 -0700
Subject: [PATCH 03/13] apply pr requests

---
 e2e/pipelines/test_gradient_rag_pipelines.py  |  4 +--
 .../embedders/gradient_document_embedder.py   |  8 +++++
 .../embedders/gradient_text_embedder.py       |  9 ++++++
 .../components/generators/gradient/base.py    | 31 ++++++++++++++++++-
 .../test_gradient_document_embedder.py        | 11 +++++++
 5 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/e2e/pipelines/test_gradient_rag_pipelines.py b/e2e/pipelines/test_gradient_rag_pipelines.py
index ca2113bbb6d..84f7f438b5d 100644
--- a/e2e/pipelines/test_gradient_rag_pipelines.py
+++ b/e2e/pipelines/test_gradient_rag_pipelines.py
@@ -14,8 +14,8 @@
 
 
 @pytest.mark.skipif(
-    not os.environ.get("GRADIENT_ACCESS_TOKEN", None),
-    reason="Export an env var called GRADIENT_ACCESS_TOKEN containing the Gradient access token to run this test.",
+    not os.environ.get("GRADIENT_ACCESS_TOKEN", None) or not os.environ.get("GRADIENT_WORKSPACE_ID", None),
+    reason="Export an env variables called GRADIENT_ACCESS_TOKEN and GRADIENT_WORKSPACE_ID containing the Gradient access token to run this test.",
 )
 def test_gradient_embedding_retrieval_rag_pipeline(tmp_path):
     # Create the RAG pipeline
diff --git a/haystack/components/embedders/gradient_document_embedder.py b/haystack/components/embedders/gradient_document_embedder.py
index 7fd6084ca51..772d16cd018 100644
--- a/haystack/components/embedders/gradient_document_embedder.py
+++ b/haystack/components/embedders/gradient_document_embedder.py
@@ -15,6 +15,14 @@ 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
+    p = Pipeline()
+    p.add_component(instance=GradientDocumentEmbedder(), name="document_embedder")
+    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__(
diff --git a/haystack/components/embedders/gradient_text_embedder.py b/haystack/components/embedders/gradient_text_embedder.py
index 39f28f51952..3e88d554a85 100644
--- a/haystack/components/embedders/gradient_text_embedder.py
+++ b/haystack/components/embedders/gradient_text_embedder.py
@@ -24,6 +24,15 @@ def __init__(
         """
         Create a GradientTextEmbedder component.
 
+        ```python
+        p = Pipeline()
+        embedder = GradientTextEmbedder(access_token=gradient_access_token)
+        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!!!")
+        ```
+
         :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.
diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/preview/components/generators/gradient/base.py
index 199742096a6..da594493206 100644
--- a/haystack/preview/components/generators/gradient/base.py
+++ b/haystack/preview/components/generators/gradient/base.py
@@ -13,6 +13,13 @@
 
 @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.
+    """
+
     @overload
     def __init__(
         self,
@@ -56,6 +63,21 @@ def __init__(
         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.
+        """
         self._access_token = access_token
         self._base_model_slug = base_model_slug
         self._host = host
@@ -78,9 +100,11 @@ def __init__(
             self._model = self._gradient.get_model_adapter(model_adapter_id=model_adapter_id)
 
     def to_dict(self) -> Dict[str, Any]:
+        """
+        Serialize this component to a dictionary.
+        """
         return default_to_dict(
             self,
-            access_token=self._access_token,
             base_model_slug=self._base_model_slug,
             host=self._host,
             max_generated_token_count=self._max_generated_token_count,
@@ -93,6 +117,11 @@ def to_dict(self) -> Dict[str, Any]:
 
     @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,
diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py
index ba9568270f6..e68a2e552e9 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 1db17e72e3f55ba07047f8c83dc77222f12a4ce1 Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <matehal@gradient.ai>
Date: Tue, 31 Oct 2023 11:57:53 -0700
Subject: [PATCH 04/13] add release notes

---
 .../notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml   | 4 ++++
 1 file changed, 4 insertions(+)
 create mode 100644 releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml

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 00000000000..44dcb7f1d2e
--- /dev/null
+++ b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml
@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Adds integration with [Gradient AI](https://gradient.ai).

From d7838d6d54916a3102fc3551e8a83608aae582cf Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <matehal@gradient.ai>
Date: Tue, 31 Oct 2023 12:20:57 -0700
Subject: [PATCH 05/13] update texts (#5)

---
 e2e/pipelines/test_gradient_rag_pipelines.py                 | 2 +-
 .../notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml  | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/e2e/pipelines/test_gradient_rag_pipelines.py b/e2e/pipelines/test_gradient_rag_pipelines.py
index 84f7f438b5d..c7b7cbeed98 100644
--- a/e2e/pipelines/test_gradient_rag_pipelines.py
+++ b/e2e/pipelines/test_gradient_rag_pipelines.py
@@ -15,7 +15,7 @@
 
 @pytest.mark.skipif(
     not os.environ.get("GRADIENT_ACCESS_TOKEN", None) or not os.environ.get("GRADIENT_WORKSPACE_ID", None),
-    reason="Export an env variables called GRADIENT_ACCESS_TOKEN and GRADIENT_WORKSPACE_ID containing the Gradient access token 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
diff --git a/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml
index 44dcb7f1d2e..bb3917b0d39 100644
--- a/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml
+++ b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml
@@ -1,4 +1,5 @@
 ---
-features:
+preview:
   - |
-    Adds integration with [Gradient AI](https://gradient.ai).
+    Add changes to Haystack version 2, or remove this section.
+    Haystack version 2 can be found under haystack/preview.

From d3ae646755f5256d67f5f5bbc34a9337f1d610b3 Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <matehal@gradient.ai>
Date: Mon, 6 Nov 2023 15:34:13 -0800
Subject: [PATCH 06/13] fix tests after rebase

---
 e2e/pipelines/test_gradient_rag_pipelines.py              | 8 ++++----
 .../components/embedders/gradient_document_embedder.py    | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/e2e/pipelines/test_gradient_rag_pipelines.py b/e2e/pipelines/test_gradient_rag_pipelines.py
index c7b7cbeed98..147e7fb8e6a 100644
--- a/e2e/pipelines/test_gradient_rag_pipelines.py
+++ b/e2e/pipelines/test_gradient_rag_pipelines.py
@@ -22,7 +22,7 @@ def test_gradient_embedding_retrieval_rag_pipeline(tmp_path):
     prompt_template = """
     Given these documents, answer the question.\nDocuments:
     {% for doc in documents %}
-        {{ doc.text }}
+        {{ doc.content }}
     {% endfor %}
 
     \nQuestion: {{question}}
@@ -60,9 +60,9 @@ def test_gradient_embedding_retrieval_rag_pipeline(tmp_path):
 
     # Populate the document store
     documents = [
-        Document(text="My name is Jean and I live in Paris."),
-        Document(text="My name is Mark and I live in Berlin."),
-        Document(text="My name is Giorgio and I live in Rome."),
+        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()
diff --git a/haystack/components/embedders/gradient_document_embedder.py b/haystack/components/embedders/gradient_document_embedder.py
index 772d16cd018..1ddf0642a8c 100644
--- a/haystack/components/embedders/gradient_document_embedder.py
+++ b/haystack/components/embedders/gradient_document_embedder.py
@@ -75,7 +75,7 @@ def _generate_embeddings(self, documents: List[Document], batch_size=100) -> Lis
 
         embeddings = []
         for batch in batches:
-            response = self._embedding_model.generate_embeddings(inputs=[{"input": doc.text} for doc in batch])
+            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

From 8db2b35994fe361933600db425382d7cd4c0b4b1 Mon Sep 17 00:00:00 2001
From: hayden <lilwilsond2@gmail.com>
Date: Tue, 7 Nov 2023 16:44:16 -0600
Subject: [PATCH 07/13] Update
 add-gradient-ai-integration-045fd476e7d3aa6a.yaml (#9)

---
 .../notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml    | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml
index bb3917b0d39..355a3874a57 100644
--- a/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml
+++ b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml
@@ -1,5 +1,4 @@
 ---
 preview:
   - |
-    Add changes to Haystack version 2, or remove this section.
-    Haystack version 2 can be found under haystack/preview.
+    Adds integration with [Gradient AI](https://gradient.ai).

From 593f9b5e2859f0eba3d2de9a0acb3d4cd3d0f736 Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <matehal@gradient.ai>
Date: Tue, 7 Nov 2023 19:20:24 -0800
Subject: [PATCH 08/13] apply the changes (#10)

---
 .../embedders/gradient_document_embedder.py   | 10 ++-
 .../embedders/gradient_text_embedder.py       | 26 ++++---
 .../components/generators/gradient/base.py    | 72 +++++++++----------
 .../test_gradient_document_embedder.py        | 16 +----
 .../embedders/test_gradient_text_embedder.py  | 39 ++++++----
 5 files changed, 88 insertions(+), 75 deletions(-)

diff --git a/haystack/components/embedders/gradient_document_embedder.py b/haystack/components/embedders/gradient_document_embedder.py
index 1ddf0642a8c..bfd34f4bb37 100644
--- a/haystack/components/embedders/gradient_document_embedder.py
+++ b/haystack/components/embedders/gradient_document_embedder.py
@@ -17,8 +17,13 @@ class GradientDocumentEmbedder:
     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(instance=GradientDocumentEmbedder(), name="document_embedder")
+    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}})
@@ -43,6 +48,7 @@ def __init__(
                              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
 
@@ -88,7 +94,7 @@ def run(self, documents: List[Document]):
 
         :param documents: A list of Documents to embed.
         """
-        if not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
+        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."
diff --git a/haystack/components/embedders/gradient_text_embedder.py b/haystack/components/embedders/gradient_text_embedder.py
index 3e88d554a85..8c89da08b4d 100644
--- a/haystack/components/embedders/gradient_text_embedder.py
+++ b/haystack/components/embedders/gradient_text_embedder.py
@@ -11,6 +11,18 @@
 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__(
@@ -24,15 +36,6 @@ def __init__(
         """
         Create a GradientTextEmbedder component.
 
-        ```python
-        p = Pipeline()
-        embedder = GradientTextEmbedder(access_token=gradient_access_token)
-        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!!!")
-        ```
-
         :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.
@@ -40,6 +43,7 @@ def __init__(
                              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
 
@@ -77,4 +81,8 @@ def run(self, text: str):
             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/gradient/base.py b/haystack/preview/components/generators/gradient/base.py
index da594493206..3a9fc63566e 100644
--- a/haystack/preview/components/generators/gradient/base.py
+++ b/haystack/preview/components/generators/gradient/base.py
@@ -18,37 +18,17 @@ class GradientGenerator:
 
     Queries the LLM using Gradient AI's SDK ('gradientai' package).
     See [Gradient AI API](https://docs.gradient.ai/docs/sdk-quickstart) for more details.
-    """
-
-    @overload
-    def __init__(
-        self,
-        *,
-        access_token: Optional[str] = None,
-        base_model_slug: str,
-        host: Optional[str] = None,
-        max_generated_token_count: Optional[int] = None,
-        temperature: Optional[float] = None,
-        top_k: Optional[int] = None,
-        top_p: Optional[float] = None,
-        workspace_id: Optional[str] = None,
-    ) -> None:
-        ...
 
-    @overload
-    def __init__(
-        self,
-        *,
-        access_token: Optional[str] = None,
-        host: Optional[str] = None,
-        max_generated_token_count: Optional[int] = None,
-        model_adapter_id: str,
-        temperature: Optional[float] = None,
-        top_k: Optional[int] = None,
-        top_p: Optional[float] = None,
-        workspace_id: Optional[str] = None,
-    ) -> None:
-        ...
+    ```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,
@@ -78,6 +58,8 @@ def __init__(
         :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
@@ -88,16 +70,20 @@ def __init__(
         self._top_p = top_p
         self._workspace_id = workspace_id
 
-        if (base_model_slug is None and model_adapter_id is None) or (
-            isinstance(base_model_slug, str) and isinstance(model_adapter_id, str)
-        ):
-            raise ValueError("expected be provided exactly one of base_model_slug or model_adapter_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)
-        if isinstance(base_model_slug, str):
-            self._model = self._gradient.get_base_model(base_model_slug=base_model_slug)
-        if isinstance(model_adapter_id, str):
-            self._model = self._gradient.get_model_adapter(model_adapter_id=model_adapter_id)
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -115,6 +101,16 @@ def to_dict(self) -> Dict[str, Any]:
             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):
         """
diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py
index e68a2e552e9..297596df9e7 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 3e2a6762f27..fa553e27f1d 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 e21ffa3e209d804e9b0e32acd9cfef122f2efa56 Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <matehal@gradient.ai>
Date: Tue, 14 Nov 2023 18:19:14 -0800
Subject: [PATCH 09/13] update tests

---
 .../embedders/gradient_document_embedder.py   |  6 ++--
 .../components/generators/gradient/base.py    |  4 +--
 .../test_gradient_document_embedder.py        | 30 +++++++++++++++----
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/haystack/components/embedders/gradient_document_embedder.py b/haystack/components/embedders/gradient_document_embedder.py
index bfd34f4bb37..c1e2fa71a60 100644
--- a/haystack/components/embedders/gradient_document_embedder.py
+++ b/haystack/components/embedders/gradient_document_embedder.py
@@ -34,6 +34,7 @@ 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,
@@ -49,6 +50,7 @@ def __init__(
         :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
 
@@ -73,7 +75,7 @@ def warm_up(self) -> None:
         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=100) -> List[List[float]]:
+    def _generate_embeddings(self, documents: List[Document], batch_size: int) -> List[List[float]]:
         """
         Batches the documents and generates the embeddings.
         """
@@ -103,7 +105,7 @@ def run(self, documents: List[Document]):
         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)
+        embeddings = self._generate_embeddings(documents=documents, batch_size=self._batch_size)
         for doc, embedding in zip(documents, embeddings):
             doc.embedding = embedding
 
diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/preview/components/generators/gradient/base.py
index 3a9fc63566e..dd420a43f4a 100644
--- a/haystack/preview/components/generators/gradient/base.py
+++ b/haystack/preview/components/generators/gradient/base.py
@@ -107,9 +107,9 @@ def warm_up(self):
         """
         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)
+                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)
+                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):
diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py
index 297596df9e7..75c57c61ff2 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 4d29ee0b41354613efca07454316c42b4e6d664a Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <mhaligowski@gmail.com>
Date: Mon, 27 Nov 2023 19:16:07 -0800
Subject: [PATCH 10/13] install gradient

---
 .github/workflows/e2e.yml     | 2 +-
 .github/workflows/linting.yml | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml
index b3a0f656010..77fb8d7f68a 100644
--- a/.github/workflows/e2e.yml
+++ b/.github/workflows/e2e.yml
@@ -36,7 +36,7 @@ jobs:
         sudo apt install ffmpeg  # for local Whisper tests
 
     - name: Install Haystack
-      run: pip install .[dev,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,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
diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml
index b117a522cb1..096d25a01dd 100644
--- a/.github/workflows/linting.yml
+++ b/.github/workflows/linting.yml
@@ -38,7 +38,7 @@ jobs:
           python-version: ${{ env.PYTHON_VERSION }}
 
       - name: Install Haystack
-        run: pip install .[dev,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf tika 'azure-ai-formrecognizer>=3.2.0b2' cohere
+        run: pip install .[dev,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf tika 'azure-ai-formrecognizer>=3.2.0b2' cohere gradientai
 
       - name: Mypy
         if: steps.files.outputs.any_changed == 'true'
@@ -69,7 +69,7 @@ jobs:
 
       - name: Install Haystack
         run: |
-          pip install .[dev,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' cohere
+          pip install .[dev,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' cohere gradientai
 
       - name: Pylint
         if: steps.files.outputs.any_changed == 'true'

From a468cd33780bc1eb1917d4631dcda428988dc597 Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <mhaligowski@gmail.com>
Date: Mon, 27 Nov 2023 19:55:59 -0800
Subject: [PATCH 11/13] more updates

---
 .github/workflows/tests.yml                               | 8 ++++----
 e2e/pipelines/test_gradient_rag_pipelines.py              | 8 ++++----
 haystack/components/embedders/__init__.py                 | 4 ++++
 .../components/embedders/gradient_document_embedder.py    | 2 +-
 haystack/components/embedders/gradient_text_embedder.py   | 2 +-
 .../embedders/test_gradient_document_embedder.py          | 4 ++--
 test/components/embedders/test_gradient_text_embedder.py  | 4 ++--
 7 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 52a2eaf3533..60e75407199 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -99,7 +99,7 @@ jobs:
           python-version: ${{ env.PYTHON_VERSION }}
 
       - name: Install Haystack
-        run: pip install .[dev,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' cohere boilerpy3
+        run: pip install .[dev,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' cohere boilerpy3 gradientai
 
       - name: Run
         run: pytest -m "not integration" test
@@ -157,7 +157,7 @@ jobs:
           sudo apt install ffmpeg  # for local Whisper tests
 
       - name: Install Haystack
-        run: pip install .[dev,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' cohere boilerpy3
+        run: pip install .[dev,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' cohere boilerpy3 gradientai
 
       - name: Run
         run: pytest --maxfail=5 -m "integration" test
@@ -213,7 +213,7 @@ jobs:
           colima start
 
       - name: Install Haystack
-        run: pip install .[dev,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' cohere boilerpy3
+        run: pip install .[dev,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' cohere boilerpy3 gradientai
 
       - name: Run Tika
         run: docker run -d -p 9998:9998 apache/tika:2.9.0.0
@@ -264,7 +264,7 @@ jobs:
           python-version: ${{ env.PYTHON_VERSION }}
 
       - name: Install Haystack
-        run: pip install .[dev,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' cohere boilerpy3
+        run: pip install .[dev,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' cohere boilerpy3 gradientai
 
       - name: Run
         run: pytest --maxfail=5 -m "integration" test -k 'not tika'
diff --git a/e2e/pipelines/test_gradient_rag_pipelines.py b/e2e/pipelines/test_gradient_rag_pipelines.py
index 147e7fb8e6a..d5807f434f7 100644
--- a/e2e/pipelines/test_gradient_rag_pipelines.py
+++ b/e2e/pipelines/test_gradient_rag_pipelines.py
@@ -3,14 +3,14 @@
 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.components.embedders.gradient_document_embedder import GradientDocumentEmbedder
+from haystack.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
+from haystack.components.builders.answer_builder import AnswerBuilder
+from haystack.components.builders.prompt_builder import PromptBuilder
 
 
 @pytest.mark.skipif(
diff --git a/haystack/components/embedders/__init__.py b/haystack/components/embedders/__init__.py
index 538da678afd..1f78b6a794d 100644
--- a/haystack/components/embedders/__init__.py
+++ b/haystack/components/embedders/__init__.py
@@ -2,10 +2,14 @@
 from haystack.components.embedders.sentence_transformers_document_embedder import SentenceTransformersDocumentEmbedder
 from haystack.components.embedders.openai_document_embedder import OpenAIDocumentEmbedder
 from haystack.components.embedders.openai_text_embedder import OpenAITextEmbedder
+from haystack.components.embedders.gradient_text_embedder import GradientTextEmbedder
+from haystack.components.embedders.gradient_document_embedder import GradientDocumentEmbedder
 
 __all__ = [
     "SentenceTransformersTextEmbedder",
     "SentenceTransformersDocumentEmbedder",
     "OpenAITextEmbedder",
     "OpenAIDocumentEmbedder",
+    "GradientTextEmbedder",
+    "GradientDocumentEmbedder",
 ]
diff --git a/haystack/components/embedders/gradient_document_embedder.py b/haystack/components/embedders/gradient_document_embedder.py
index c1e2fa71a60..3b91698057c 100644
--- a/haystack/components/embedders/gradient_document_embedder.py
+++ b/haystack/components/embedders/gradient_document_embedder.py
@@ -2,7 +2,7 @@
 from typing import List, Optional, Dict, Any
 
 from haystack.preview import component, Document, default_to_dict
-from haystack.preview.lazy_imports import LazyImport
+from haystack.lazy_imports import LazyImport
 
 with LazyImport(message="Run 'pip install gradientai'") as gradientai_import:
     from gradientai import Gradient
diff --git a/haystack/components/embedders/gradient_text_embedder.py b/haystack/components/embedders/gradient_text_embedder.py
index 8c89da08b4d..cd8aefd1ca5 100644
--- a/haystack/components/embedders/gradient_text_embedder.py
+++ b/haystack/components/embedders/gradient_text_embedder.py
@@ -1,7 +1,7 @@
 from typing import Any, Dict, List, Optional
 
 from haystack.preview import component, default_to_dict
-from haystack.preview.lazy_imports import LazyImport
+from haystack.lazy_imports import LazyImport
 
 with LazyImport(message="Run 'pip install gradientai'") as gradientai_import:
     from gradientai import Gradient
diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py
index 75c57c61ff2..e9cbfdd9e32 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 fa553e27f1d..779ed5c1eef 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 14581ed7f764596e3ec2a09ecdd19852b17f9845 Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <mhaligowski@gmail.com>
Date: Mon, 27 Nov 2023 22:43:56 -0800
Subject: [PATCH 12/13] more updates

---
 .../components/embedders/gradient_document_embedder.py    | 2 +-
 haystack/components/embedders/gradient_text_embedder.py   | 2 +-
 .../embedders/test_gradient_document_embedder.py          | 8 ++++----
 test/components/embedders/test_gradient_text_embedder.py  | 6 +++---
 4 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/haystack/components/embedders/gradient_document_embedder.py b/haystack/components/embedders/gradient_document_embedder.py
index 3b91698057c..e2d3458c598 100644
--- a/haystack/components/embedders/gradient_document_embedder.py
+++ b/haystack/components/embedders/gradient_document_embedder.py
@@ -1,7 +1,7 @@
 import logging
 from typing import List, Optional, Dict, Any
 
-from haystack.preview import component, Document, default_to_dict
+from haystack import component, Document, default_to_dict
 from haystack.lazy_imports import LazyImport
 
 with LazyImport(message="Run 'pip install gradientai'") as gradientai_import:
diff --git a/haystack/components/embedders/gradient_text_embedder.py b/haystack/components/embedders/gradient_text_embedder.py
index cd8aefd1ca5..28a83a52e22 100644
--- a/haystack/components/embedders/gradient_text_embedder.py
+++ b/haystack/components/embedders/gradient_text_embedder.py
@@ -1,6 +1,6 @@
 from typing import Any, Dict, List, Optional
 
-from haystack.preview import component, default_to_dict
+from haystack import component, default_to_dict
 from haystack.lazy_imports import LazyImport
 
 with LazyImport(message="Run 'pip install gradientai'") as gradientai_import:
diff --git a/test/components/embedders/test_gradient_document_embedder.py b/test/components/embedders/test_gradient_document_embedder.py
index e9cbfdd9e32..ffb44519fdb 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 779ed5c1eef..82c84e39043 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 6b3f10750114ae1558c933231cb6aec510936450 Mon Sep 17 00:00:00 2001
From: Mateusz Haligowski <mhaligowski@gmail.com>
Date: Wed, 29 Nov 2023 11:37:00 -0800
Subject: [PATCH 13/13] fix e2e test

---
 e2e/pipelines/test_gradient_rag_pipelines.py           | 10 +++++-----
 haystack/components/generators/__init__.py             |  3 ++-
 .../components/generators/gradient/__init__.py         |  0
 .../components/generators/gradient/base.py             |  2 +-
 4 files changed, 8 insertions(+), 7 deletions(-)
 rename haystack/{preview => }/components/generators/gradient/__init__.py (100%)
 rename haystack/{preview => }/components/generators/gradient/base.py (98%)

diff --git a/e2e/pipelines/test_gradient_rag_pipelines.py b/e2e/pipelines/test_gradient_rag_pipelines.py
index d5807f434f7..ff655495004 100644
--- a/e2e/pipelines/test_gradient_rag_pipelines.py
+++ b/e2e/pipelines/test_gradient_rag_pipelines.py
@@ -2,13 +2,13 @@
 import json
 import pytest
 
-from haystack.preview import Pipeline, Document
+from haystack import Pipeline, Document
 from haystack.components.embedders.gradient_document_embedder import GradientDocumentEmbedder
 from haystack.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.document_stores import InMemoryDocumentStore
+from haystack.components.writers import DocumentWriter
+from haystack.components.retrievers import InMemoryEmbeddingRetriever
+from haystack.components.generators.gradient.base import GradientGenerator
 from haystack.components.builders.answer_builder import AnswerBuilder
 from haystack.components.builders.prompt_builder import PromptBuilder
 
diff --git a/haystack/components/generators/__init__.py b/haystack/components/generators/__init__.py
index 5466bce2c11..dfd7360ef7a 100644
--- a/haystack/components/generators/__init__.py
+++ b/haystack/components/generators/__init__.py
@@ -1,6 +1,7 @@
 from haystack.components.generators.cohere import CohereGenerator
+from haystack.components.generators.gradient.base import GradientGenerator
 from haystack.components.generators.hugging_face_local import HuggingFaceLocalGenerator
 from haystack.components.generators.hugging_face_tgi import HuggingFaceTGIGenerator
 from haystack.components.generators.openai import GPTGenerator
 
-__all__ = ["HuggingFaceLocalGenerator", "HuggingFaceTGIGenerator", "GPTGenerator", "CohereGenerator"]
+__all__ = ["HuggingFaceLocalGenerator", "HuggingFaceTGIGenerator", "GPTGenerator", "CohereGenerator", "GradientGenerator"]
diff --git a/haystack/preview/components/generators/gradient/__init__.py b/haystack/components/generators/gradient/__init__.py
similarity index 100%
rename from haystack/preview/components/generators/gradient/__init__.py
rename to haystack/components/generators/gradient/__init__.py
diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/components/generators/gradient/base.py
similarity index 98%
rename from haystack/preview/components/generators/gradient/base.py
rename to haystack/components/generators/gradient/base.py
index dd420a43f4a..4853a922916 100644
--- a/haystack/preview/components/generators/gradient/base.py
+++ b/haystack/components/generators/gradient/base.py
@@ -3,7 +3,7 @@
 import logging
 from haystack.lazy_imports import LazyImport
 
-from haystack.preview import component, default_to_dict
+from haystack import component, default_to_dict
 
 with LazyImport(message="Run 'pip install gradientai'") as gradientai_import:
     from gradientai import Gradient