Skip to content

Commit 02d3250

Browse files
feat(drivers-embedding-twelvelabs): add TwelveLabsEmbeddingDriver
Adds an opt-in embedding driver backed by TwelveLabs' Marengo multimodal model. Marengo maps text, images, and video into the same 512-dimensional vector space, enabling text-to-visual retrieval. Gated behind the `drivers-embedding-twelvelabs` extra; no existing behavior changes.
1 parent ebf422b commit 02d3250

8 files changed

Lines changed: 146 additions & 2 deletions

File tree

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,18 @@ The [CohereEmbeddingDriver](../../reference/griptape/drivers/embedding/cohere_em
182182
--8<-- "docs/griptape-framework/drivers/logs/embedding_drivers_9.txt"
183183
```
184184

185+
### TwelveLabs
186+
187+
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.
188+
189+
!!! info
190+
191+
This driver requires the `drivers-embedding-twelvelabs` [extra](../index.md#extras).
192+
193+
```python
194+
--8<-- "docs/griptape-framework/drivers/src/embedding_drivers_11.py"
195+
```
196+
185197
### Nvidia NIM
186198

187199
The [NvidiaNimEmbeddingDriver](../../reference/griptape/drivers/embedding/nvidia_nim_embedding_driver.md) uses the [Nvidia NIM API](https://developer.nvidia.com/nim).
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
import os
2+
3+
from griptape.drivers.embedding.twelvelabs import TwelveLabsEmbeddingDriver
4+
5+
embedding_driver = TwelveLabsEmbeddingDriver(
6+
model="marengo3.0",
7+
api_key=os.environ["TWELVELABS_API_KEY"],
8+
)
9+
10+
embeddings = embedding_driver.embed("Hello world!")
11+
12+
# display the first 3 embeddings
13+
print(embeddings[:3])

griptape/drivers/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
from .embedding.dummy import DummyEmbeddingDriver
3535
from .embedding.cohere import CohereEmbeddingDriver
3636
from .embedding.ollama import OllamaEmbeddingDriver
37+
from .embedding.twelvelabs import TwelveLabsEmbeddingDriver
3738

3839
from .vector import BaseVectorStoreDriver
3940
from .vector.local import LocalVectorStoreDriver
@@ -263,6 +264,7 @@
263264
"TavilyWebSearchDriver",
264265
"TrafilaturaWebScraperDriver",
265266
"VoyageAiEmbeddingDriver",
267+
"TwelveLabsEmbeddingDriver",
266268
"WebhookEventListenerDriver",
267269
]
268270

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from griptape.drivers.embedding.twelvelabs_embedding_driver import TwelveLabsEmbeddingDriver
2+
3+
__all__ = ["TwelveLabsEmbeddingDriver"]
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
from __future__ import annotations
2+
3+
from typing import TYPE_CHECKING, Any
4+
5+
from attrs import define, field
6+
7+
from griptape.artifacts import ImageArtifact, TextArtifact
8+
from griptape.drivers.embedding import BaseEmbeddingDriver
9+
from griptape.utils import import_optional_dependency
10+
from griptape.utils.decorators import lazy_property
11+
12+
if TYPE_CHECKING:
13+
from twelvelabs import TwelveLabs
14+
15+
16+
@define
17+
class TwelveLabsEmbeddingDriver(BaseEmbeddingDriver):
18+
"""TwelveLabs Marengo Embedding Driver.
19+
20+
Generates multimodal embeddings with TwelveLabs' Marengo model. Marengo maps
21+
text, images, and video into the same 512-dimensional vector space, so text
22+
queries can be matched against visual content stored in a vector store.
23+
24+
Attributes:
25+
model: TwelveLabs Marengo model name. Defaults to `marengo3.0`.
26+
api_key: TwelveLabs API key. Defaults to the `TWELVELABS_API_KEY` environment variable.
27+
client: Optionally provide a custom `twelvelabs.TwelveLabs` client.
28+
"""
29+
30+
DEFAULT_MODEL = "marengo3.0"
31+
32+
model: str = field(default=DEFAULT_MODEL, kw_only=True, metadata={"serializable": True})
33+
api_key: str | None = field(default=None, kw_only=True, metadata={"serializable": False})
34+
_client: TwelveLabs | None = field(default=None, kw_only=True, alias="client", metadata={"serializable": False})
35+
36+
@lazy_property()
37+
def client(self) -> TwelveLabs:
38+
return import_optional_dependency("twelvelabs").TwelveLabs(api_key=self.api_key)
39+
40+
def try_embed_artifact(self, artifact: TextArtifact | ImageArtifact, **kwargs) -> list[float]:
41+
if isinstance(artifact, TextArtifact):
42+
return self.try_embed_chunk(artifact.value, **kwargs)
43+
response = self.client.embed.create(
44+
model_name=self.model,
45+
image_file=(artifact.name, artifact.value, artifact.mime_type),
46+
)
47+
return self._extract_vector(response.image_embedding)
48+
49+
def try_embed_chunk(self, chunk: str, **kwargs) -> list[float]:
50+
response = self.client.embed.create(model_name=self.model, text=chunk)
51+
return self._extract_vector(response.text_embedding)
52+
53+
def _extract_vector(self, result: Any) -> list[float]:
54+
if result is None or not result.segments:
55+
raise ValueError("TwelveLabs returned no embedding segments.")
56+
vector = result.segments[0].float_
57+
if vector is None:
58+
raise ValueError("TwelveLabs returned an empty embedding.")
59+
return vector

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ drivers-embedding-voyageai = ["voyageai>=0.2.1"]
7272
drivers-embedding-google = ["google-genai>=1.73.1"]
7373
drivers-embedding-cohere = ["cohere>=5.11.2"]
7474
drivers-embedding-ollama = ["ollama>=0.4.1"]
75+
drivers-embedding-twelvelabs = ["twelvelabs>=1.2.8"]
7576
drivers-web-scraper-trafilatura = ["trafilatura>=2.0"]
7677
drivers-web-scraper-markdownify = [
7778
"playwright>=1.42",
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
from unittest.mock import Mock
2+
3+
import pytest
4+
5+
from griptape.artifacts import ImageArtifact, TextArtifact
6+
from griptape.drivers.embedding.twelvelabs import TwelveLabsEmbeddingDriver
7+
8+
9+
class TestTwelveLabsEmbeddingDriver:
10+
@pytest.fixture(autouse=True)
11+
def mock_client(self, mocker):
12+
mock_client = mocker.patch("twelvelabs.TwelveLabs")
13+
mock_client.return_value.embed.create.return_value = Mock(
14+
text_embedding=Mock(segments=[Mock(float_=[0, 1, 0])]),
15+
image_embedding=Mock(segments=[Mock(float_=[1, 0, 1])]),
16+
)
17+
return mock_client
18+
19+
def test_init(self):
20+
assert TwelveLabsEmbeddingDriver()
21+
22+
def test_embed_string(self):
23+
assert TwelveLabsEmbeddingDriver().embed("foobar") == [0, 1, 0]
24+
25+
def test_embed_text_artifact(self):
26+
assert TwelveLabsEmbeddingDriver().embed(TextArtifact("foobar")) == [0, 1, 0]
27+
28+
def test_embed_image_artifact(self):
29+
artifact = ImageArtifact(b"foobar", format="jpeg", width=1, height=1)
30+
assert TwelveLabsEmbeddingDriver().embed(artifact) == [1, 0, 1]
31+
32+
def test_embed_no_segments_raises(self, mock_client):
33+
mock_client.return_value.embed.create.return_value = Mock(text_embedding=Mock(segments=[]))
34+
with pytest.raises(ValueError, match="no embedding segments"):
35+
TwelveLabsEmbeddingDriver().embed("foobar")

uv.lock

Lines changed: 21 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)