Add Gradient Integration (#72)

* Remove preview folder in test/ --------- Co-authored-by: Silvano Cerza <[email protected]> * Fix all tests * remove unit marker (#6450) * Gradient embeddings (#2) * apply pr requests * apply the changes (#10) * update tests * more updates * more updates * prepare * remove commits migration artifacts * create package structure * run tests in CI * fix project file * add integration test * fix complaining linter * fix skip call --------- Co-authored-by: Silvano Cerza <[email protected]> Co-authored-by: Mateusz Haligowski <[email protected]> Co-authored-by: Mateusz Haligowski <[email protected]>
deepset-ai · Dec 2, 2023 · d3aa338 · d3aa338
1 parent cbc3bfc
commit d3aa338
Show file tree

Hide file tree

Showing 15 changed files with 979 additions and 0 deletions.
diff --git a/.github/workflows/gradient.yml b/.github/workflows/gradient.yml
@@ -0,0 +1,56 @@
+# This workflow comes from https://github.com/ofek/hatch-mypyc
+# https://github.com/ofek/hatch-mypyc/blob/5a198c0ba8660494d02716cfc9d79ce4adfb1442/.github/workflows/test.yml
+name: Test / gradient
+
+on:
+  schedule:
+    - cron: "0 0 * * *"
+  pull_request:
+    paths:
+      - 'integrations/gradient/**'
+      - '.github/workflows/gradient.yml'
+
+defaults:
+  run:
+    working-directory: integrations/gradient
+
+concurrency:
+  group: gradient-${{ github.head_ref }}
+  cancel-in-progress: true
+
+env:
+  PYTHONUNBUFFERED: "1"
+  FORCE_COLOR: "1"
+
+jobs:
+  run:
+    name: Python ${{ matrix.python-version }} on ${{ startsWith(matrix.os, 'macos-') && 'macOS' || startsWith(matrix.os, 'windows-') && 'Windows' || 'Linux' }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+        python-version: ['3.9', '3.10']
+
+    steps:
+    - name: Support longpaths
+      if: matrix.os == 'windows-latest'
+      working-directory: .
+      run: git config --system core.longpaths true
+
+    - uses: actions/checkout@v4
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install Hatch
+      run: pip install --upgrade hatch
+
+    - name: Lint
+      if: matrix.python-version == '3.9' && runner.os == 'Linux'
+      run: hatch run lint:all
+
+    - name: Run tests
+      run: hatch run cov
diff --git a/integrations/gradient/LICENSE.txt b/integrations/gradient/LICENSE.txt
@@ -0,0 +1,9 @@
+MIT License
+
+Copyright (c) 2023-present Massimiliano Pippi <[email protected]>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
diff --git a/integrations/gradient/README.md b/integrations/gradient/README.md
@@ -0,0 +1,21 @@
+# gradient-haystack
+
+[![PyPI - Version](https://img.shields.io/pypi/v/gradient-haystack.svg)](https://pypi.org/project/gradient-haystack)
+[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/gradient-haystack.svg)](https://pypi.org/project/gradient-haystack)
+
+-----
+
+**Table of Contents**
+
+- [Installation](#installation)
+- [License](#license)
+
+## Installation
+
+```console
+pip install gradient-haystack
+```
+
+## License
+
+`gradient-haystack` is distributed under the terms of the [MIT](https://spdx.org/licenses/MIT.html) license.
diff --git a/integrations/gradient/pyproject.toml b/integrations/gradient/pyproject.toml
@@ -0,0 +1,169 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "gradient-haystack"
+dynamic = ["version"]
+description = ''
+readme = "README.md"
+requires-python = ">=3.7"
+license = "MIT"
+keywords = []
+authors = [
+  { name = "Mateusz Haligowski", email = "[email protected]" },
+]
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3.7",
+  "Programming Language :: Python :: 3.8",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: Implementation :: CPython",
+  "Programming Language :: Python :: Implementation :: PyPy",
+]
+dependencies = [
+  "haystack-ai",
+  "gradientai",
+]
+
+[project.urls]
+Documentation = "https://github.com/unknown/gradient-haystack#readme"
+Issues = "https://github.com/unknown/gradient-haystack/issues"
+Source = "https://github.com/unknown/gradient-haystack"
+
+[tool.hatch.version]
+path = "src/gradient_haystack/__about__.py"
+
+[tool.hatch.envs.default]
+dependencies = [
+  "coverage[toml]>=6.5",
+  "pytest",
+]
+[tool.hatch.envs.default.scripts]
+test = "pytest {args:tests}"
+test-cov = "coverage run -m pytest {args:tests}"
+cov-report = [
+  "- coverage combine",
+  "coverage report",
+]
+cov = [
+  "test-cov",
+  "cov-report",
+]
+
+[[tool.hatch.envs.all.matrix]]
+python = ["3.8", "3.9", "3.10", "3.11"]
+
+[tool.hatch.envs.lint]
+detached = true
+dependencies = [
+  "black>=23.1.0",
+  "mypy>=1.0.0",
+  "ruff>=0.0.243",
+]
+[tool.hatch.envs.lint.scripts]
+typing = "mypy --install-types --non-interactive {args:src/gradient_haystack tests}"
+style = [
+  "ruff {args:.}",
+  "black --check --diff {args:.}",
+]
+fmt = [
+  "black {args:.}",
+  "ruff --fix {args:.}",
+  "style",
+]
+all = [
+  "style",
+  "typing",
+]
+
+[tool.black]
+target-version = ["py38"]
+line-length = 120
+skip-string-normalization = true
+
+[tool.ruff]
+target-version = "py38"
+line-length = 120
+select = [
+  "A",
+  "ARG",
+  "B",
+  "C",
+  "DTZ",
+  "E",
+  "EM",
+  "F",
+  "FBT",
+  "I",
+  "ICN",
+  "ISC",
+  "N",
+  "PLC",
+  "PLE",
+  "PLR",
+  "PLW",
+  "Q",
+  "RUF",
+  "S",
+  "T",
+  "TID",
+  "UP",
+  "W",
+  "YTT",
+]
+ignore = [
+  # Allow non-abstract empty methods in abstract base classes
+  "B027",
+  # Allow boolean positional values in function calls, like `dict.get(... True)`
+  "FBT003",
+  # Ignore checks for possible passwords
+  "S105", "S106", "S107",
+  # Ignore complexity
+  "C901", "PLR0911", "PLR0912", "PLR0913", "PLR0915",
+]
+unfixable = [
+  # Don't touch unused imports
+  "F401",
+]
+
+[tool.ruff.isort]
+known-first-party = ["gradient_haystack"]
+
+[tool.ruff.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+[tool.ruff.per-file-ignores]
+# Tests can use magic values, assertions, and relative imports
+"tests/**/*" = ["PLR2004", "S101", "TID252"]
+
+[tool.coverage.run]
+source_pkgs = ["gradient_haystack", "tests"]
+branch = true
+parallel = true
+omit = [
+  "src/gradient_haystack/__about__.py",
+]
+
+[tool.coverage.paths]
+gradient_haystack = ["src/gradient_haystack", "*/gradient-haystack/src/gradient_haystack"]
+tests = ["tests", "*/gradient-haystack/tests"]
+
+[tool.coverage.report]
+exclude_lines = [
+  "no cov",
+  "if __name__ == .__main__.:",
+  "if TYPE_CHECKING:",
+]
+
+[[tool.mypy.overrides]]
+module = [
+  "gradientai.*",
+  "haystack.*",
+  "pytest.*",
+  "numpy.*",
+]
+ignore_missing_imports = true
diff --git a/integrations/gradient/src/gradient_haystack/__about__.py b/integrations/gradient/src/gradient_haystack/__about__.py
@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <[email protected]>
+#
+# SPDX-License-Identifier: MIT
+__version__ = "0.0.1"
diff --git a/integrations/gradient/src/gradient_haystack/__init__.py b/integrations/gradient/src/gradient_haystack/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <[email protected]>
+#
+# SPDX-License-Identifier: MIT
diff --git a/integrations/gradient/src/gradient_haystack/embedders/__init__.py b/integrations/gradient/src/gradient_haystack/embedders/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <[email protected]>
+#
+# SPDX-License-Identifier: MIT
diff --git a/integrations/gradient/src/gradient_haystack/embedders/gradient_document_embedder.py b/integrations/gradient/src/gradient_haystack/embedders/gradient_document_embedder.py
@@ -0,0 +1,112 @@
+import logging
+from typing import Any, Dict, List, Optional
+
+from haystack import Document, component, default_to_dict
+from haystack.lazy_imports import LazyImport
+
+with LazyImport(message="Run 'pip install gradientai'") as gradientai_import:
+    from gradientai import Gradient
+
+logger = logging.getLogger(__name__)
+
+
+@component
+class GradientDocumentEmbedder:
+    """
+    A component for computing Document embeddings using Gradient AI API..
+    The embedding of each Document is stored in the `embedding` field of the Document.
+
+    ```python
+    embedder = GradientDocumentEmbedder(
+        access_token=gradient_access_token,
+        workspace_id=gradient_workspace_id,
+        model_name="bge_large"))
+    p = Pipeline()
+    p.add_component(embedder, name="document_embedder")
+    p.add_component(instance=GradientDocumentEmbedder(
+    p.add_component(instance=DocumentWriter(document_store=InMemoryDocumentStore()), name="document_writer")
+    p.connect("document_embedder", "document_writer")
+    p.run({"document_embedder": {"documents": documents}})
+    ```
+    """
+
+    def __init__(
+        self,
+        *,
+        model_name: str = "bge-large",
+        batch_size: int = 100,
+        access_token: Optional[str] = None,
+        workspace_id: Optional[str] = None,
+        host: Optional[str] = None,
+    ) -> None:
+        """
+        Create a GradientDocumentEmbedder component.
+
+        :param model_name: The name of the model to use.
+        :param access_token: The Gradient access token. If not provided it's read from the environment
+                             variable GRADIENT_ACCESS_TOKEN.
+        :param workspace_id: The Gradient workspace ID. If not provided it's read from the environment
+                             variable GRADIENT_WORKSPACE_ID.
+        :param host: The Gradient host. By default it uses https://api.gradient.ai/.
+        """
+        gradientai_import.check()
+        self._batch_size = batch_size
+        self._host = host
+        self._model_name = model_name
+
+        self._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id)
+
+    def _get_telemetry_data(self) -> Dict[str, Any]:
+        """
+        Data that is sent to Posthog for usage analytics.
+        """
+        return {"model": self._model_name}
+
+    def to_dict(self) -> dict:
+        """
+        Serialize the component to a Python dictionary.
+        """
+        return default_to_dict(self, workspace_id=self._gradient.workspace_id, model_name=self._model_name)
+
+    def warm_up(self) -> None:
+        """
+        Load the embedding model.
+        """
+        if not hasattr(self, "_embedding_model"):
+            self._embedding_model = self._gradient.get_embeddings_model(slug=self._model_name)
+
+    def _generate_embeddings(self, documents: List[Document], batch_size: int) -> List[List[float]]:
+        """
+        Batches the documents and generates the embeddings.
+        """
+        batches = [documents[i : i + batch_size] for i in range(0, len(documents), batch_size)]
+
+        embeddings = []
+        for batch in batches:
+            response = self._embedding_model.generate_embeddings(inputs=[{"input": doc.content} for doc in batch])
+            embeddings.extend([e.embedding for e in response.embeddings])
+
+        return embeddings
+
+    @component.output_types(documents=List[Document])
+    def run(self, documents: List[Document]):
+        """
+        Embed a list of Documents.
+        The embedding of each Document is stored in the `embedding` field of the Document.
+
+        :param documents: A list of Documents to embed.
+        """
+        if not isinstance(documents, list) or documents and any(not isinstance(doc, Document) for doc in documents):
+            msg = "GradientDocumentEmbedder expects a list of Documents as input.\
+                  In case you want to embed a list of strings, please use the GradientTextEmbedder."
+            raise TypeError(msg)
+
+        if not hasattr(self, "_embedding_model"):
+            msg = "The embedding model has not been loaded. Please call warm_up() before running."
+            raise RuntimeError(msg)
+
+        embeddings = self._generate_embeddings(documents=documents, batch_size=self._batch_size)
+        for doc, embedding in zip(documents, embeddings):
+            doc.embedding = embedding
+
+        return {"documents": documents}