feat(drivers-embedding-twelvelabs): add TwelveLabsEmbeddingDriver

docs/griptape-framework/drivers/embedding-drivers.md

@@ -182,6 +182,18 @@ The [CohereEmbeddingDriver](../../reference/griptape/drivers/embedding/cohere_em
     --8<-- "docs/griptape-framework/drivers/logs/embedding_drivers_9.txt"
     ```
 
+### TwelveLabs
+
+The [TwelveLabsEmbeddingDriver](../../reference/griptape/drivers/embedding/twelvelabs_embedding_driver.md) uses the [TwelveLabs Marengo](https://docs.twelvelabs.io/) multimodal embedding model. Marengo maps text, images, and video into the same 512-dimensional vector space, so text queries can be matched against visual content.
+
+!!! info
+
+    This driver requires the `drivers-embedding-twelvelabs` [extra](../index.md#extras).
+
+```python
+--8<-- "docs/griptape-framework/drivers/src/embedding_drivers_11.py"
+```
+
 ### Nvidia NIM
 
 The [NvidiaNimEmbeddingDriver](../../reference/griptape/drivers/embedding/nvidia_nim_embedding_driver.md) uses the [Nvidia NIM API](https://developer.nvidia.com/nim).

docs/griptape-framework/drivers/src/embedding_drivers_11.py

@@ -0,0 +1,13 @@
+import os
+
+from griptape.drivers.embedding.twelvelabs import TwelveLabsEmbeddingDriver
+
+embedding_driver = TwelveLabsEmbeddingDriver(
+    model="marengo3.0",
+    api_key=os.environ["TWELVELABS_API_KEY"],
+)
+
+embeddings = embedding_driver.embed("Hello world!")
+
+# display the first 3 embeddings
+print(embeddings[:3])

griptape/drivers/__init__.py

@@ -34,6 +34,7 @@
 from .embedding.dummy import DummyEmbeddingDriver
 from .embedding.cohere import CohereEmbeddingDriver
 from .embedding.ollama import OllamaEmbeddingDriver
+from .embedding.twelvelabs import TwelveLabsEmbeddingDriver
 
 from .vector import BaseVectorStoreDriver
 from .vector.local import LocalVectorStoreDriver
@@ -263,6 +264,7 @@
     "TavilyWebSearchDriver",
     "TrafilaturaWebScraperDriver",
     "VoyageAiEmbeddingDriver",
+    "TwelveLabsEmbeddingDriver",
     "WebhookEventListenerDriver",
 ]
 

griptape/drivers/embedding/twelvelabs/__init__.py

@@ -0,0 +1,3 @@
+from griptape.drivers.embedding.twelvelabs_embedding_driver import TwelveLabsEmbeddingDriver
+
+__all__ = ["TwelveLabsEmbeddingDriver"]

griptape/drivers/embedding/twelvelabs_embedding_driver.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from attrs import define, field
+
+from griptape.artifacts import ImageArtifact, TextArtifact
+from griptape.drivers.embedding import BaseEmbeddingDriver
+from griptape.utils import import_optional_dependency
+from griptape.utils.decorators import lazy_property
+
+if TYPE_CHECKING:
+    from twelvelabs import TwelveLabs
+
+
+@define
+class TwelveLabsEmbeddingDriver(BaseEmbeddingDriver):
+    """TwelveLabs Marengo Embedding Driver.
+
+    Generates multimodal embeddings with TwelveLabs' Marengo model. Marengo maps
+    text, images, and video into the same 512-dimensional vector space, so text
+    queries can be matched against visual content stored in a vector store.
+
+    Attributes:
+        model: TwelveLabs Marengo model name. Defaults to `marengo3.0`.
+        api_key: TwelveLabs API key. Defaults to the `TWELVELABS_API_KEY` environment variable.
+        client: Optionally provide a custom `twelvelabs.TwelveLabs` client.
+    """
+
+    DEFAULT_MODEL = "marengo3.0"
+
+    model: str = field(default=DEFAULT_MODEL, kw_only=True, metadata={"serializable": True})
+    api_key: str | None = field(default=None, kw_only=True, metadata={"serializable": False})
+    _client: TwelveLabs | None = field(default=None, kw_only=True, alias="client", metadata={"serializable": False})
+
+    @lazy_property()
+    def client(self) -> TwelveLabs:
+        return import_optional_dependency("twelvelabs").TwelveLabs(api_key=self.api_key)
+
+    def try_embed_artifact(self, artifact: TextArtifact | ImageArtifact, **kwargs) -> list[float]:
+        if isinstance(artifact, TextArtifact):
+            return self.try_embed_chunk(artifact.value, **kwargs)
+        response = self.client.embed.create(
+            model_name=self.model,
+            image_file=(artifact.name, artifact.value, artifact.mime_type),
+        )
+        return self._extract_vector(response.image_embedding)
+
+    def try_embed_chunk(self, chunk: str, **kwargs) -> list[float]:
+        response = self.client.embed.create(model_name=self.model, text=chunk)
+        return self._extract_vector(response.text_embedding)
+
+    def _extract_vector(self, result: Any) -> list[float]:
+        if result is None or not result.segments:
+            raise ValueError("TwelveLabs returned no embedding segments.")
+        vector = result.segments[0].float_
+        if vector is None:
+            raise ValueError("TwelveLabs returned an empty embedding.")
+        return vector

pyproject.toml

@@ -72,6 +72,7 @@ drivers-embedding-voyageai = ["voyageai>=0.2.1"]
 drivers-embedding-google = ["google-genai>=1.73.1"]
 drivers-embedding-cohere = ["cohere>=5.11.2"]
 drivers-embedding-ollama = ["ollama>=0.4.1"]
+drivers-embedding-twelvelabs = ["twelvelabs>=1.2.8"]
 drivers-web-scraper-trafilatura = ["trafilatura>=2.0"]
 drivers-web-scraper-markdownify = [
   "playwright>=1.42",

tests/unit/drivers/embedding/test_twelvelabs_embedding_driver.py

@@ -0,0 +1,35 @@
+from unittest.mock import Mock
+
+import pytest
+
+from griptape.artifacts import ImageArtifact, TextArtifact
+from griptape.drivers.embedding.twelvelabs import TwelveLabsEmbeddingDriver
+
+
+class TestTwelveLabsEmbeddingDriver:
+    @pytest.fixture(autouse=True)
+    def mock_client(self, mocker):
+        mock_client = mocker.patch("twelvelabs.TwelveLabs")
+        mock_client.return_value.embed.create.return_value = Mock(
+            text_embedding=Mock(segments=[Mock(float_=[0, 1, 0])]),
+            image_embedding=Mock(segments=[Mock(float_=[1, 0, 1])]),
+        )
+        return mock_client
+
+    def test_init(self):
+        assert TwelveLabsEmbeddingDriver()
+
+    def test_embed_string(self):
+        assert TwelveLabsEmbeddingDriver().embed("foobar") == [0, 1, 0]
+
+    def test_embed_text_artifact(self):
+        assert TwelveLabsEmbeddingDriver().embed(TextArtifact("foobar")) == [0, 1, 0]
+
+    def test_embed_image_artifact(self):
+        artifact = ImageArtifact(b"foobar", format="jpeg", width=1, height=1)
+        assert TwelveLabsEmbeddingDriver().embed(artifact) == [1, 0, 1]
+
+    def test_embed_no_segments_raises(self, mock_client):
+        mock_client.return_value.embed.create.return_value = Mock(text_embedding=Mock(segments=[]))
+        with pytest.raises(ValueError, match="no embedding segments"):
+            TwelveLabsEmbeddingDriver().embed("foobar")