diff --git a/docs/griptape-framework/drivers/embedding-drivers.md b/docs/griptape-framework/drivers/embedding-drivers.md index 39b76115d..9d4351572 100644 --- a/docs/griptape-framework/drivers/embedding-drivers.md +++ b/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). diff --git a/docs/griptape-framework/drivers/src/embedding_drivers_11.py b/docs/griptape-framework/drivers/src/embedding_drivers_11.py new file mode 100644 index 000000000..6363d77f2 --- /dev/null +++ b/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]) diff --git a/griptape/drivers/__init__.py b/griptape/drivers/__init__.py index 4fba24367..df99d992d 100644 --- a/griptape/drivers/__init__.py +++ b/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", ] diff --git a/griptape/drivers/embedding/twelvelabs/__init__.py b/griptape/drivers/embedding/twelvelabs/__init__.py new file mode 100644 index 000000000..71311c5bf --- /dev/null +++ b/griptape/drivers/embedding/twelvelabs/__init__.py @@ -0,0 +1,3 @@ +from griptape.drivers.embedding.twelvelabs_embedding_driver import TwelveLabsEmbeddingDriver + +__all__ = ["TwelveLabsEmbeddingDriver"] diff --git a/griptape/drivers/embedding/twelvelabs_embedding_driver.py b/griptape/drivers/embedding/twelvelabs_embedding_driver.py new file mode 100644 index 000000000..31212a9af --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 50b68c5f3..9e96db4fc 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/unit/drivers/embedding/test_twelvelabs_embedding_driver.py b/tests/unit/drivers/embedding/test_twelvelabs_embedding_driver.py new file mode 100644 index 000000000..a32f6faed --- /dev/null +++ b/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") diff --git a/uv.lock b/uv.lock index 617760633..b45c2c562 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.10, <4" resolution-markers = [ "python_full_version >= '3.13'", @@ -1465,6 +1465,9 @@ drivers-embedding-huggingface = [ drivers-embedding-ollama = [ { name = "ollama" }, ] +drivers-embedding-twelvelabs = [ + { name = "twelvelabs" }, +] drivers-embedding-voyageai = [ { name = "voyageai" }, ] @@ -1788,12 +1791,13 @@ requires-dist = [ { name = "transformers", marker = "extra == 'drivers-prompt-amazon-sagemaker'", specifier = ">=4.41.1" }, { name = "transformers", marker = "extra == 'drivers-prompt-huggingface-hub'", specifier = ">=4.41.1" }, { name = "transformers", marker = "extra == 'drivers-prompt-huggingface-pipeline'", specifier = ">=4.41.1" }, + { name = "twelvelabs", marker = "extra == 'drivers-embedding-twelvelabs'", specifier = ">=1.2.8" }, { name = "urllib3", specifier = ">=1.25.4,!=2.2.0,<3" }, { name = "voyageai", marker = "extra == 'all'", specifier = ">=0.2.1" }, { name = "voyageai", marker = "extra == 'drivers-embedding-voyageai'", specifier = ">=0.2.1" }, { name = "wrapt", specifier = ">=1.16.0" }, ] -provides-extras = ["drivers-prompt-cohere", "drivers-prompt-anthropic", "drivers-prompt-huggingface-hub", "drivers-prompt-huggingface-pipeline", "drivers-prompt-amazon-bedrock", "drivers-prompt-amazon-sagemaker", "drivers-prompt-google", "drivers-prompt-ollama", "drivers-sql", "drivers-sql-amazon-redshift", "drivers-sql-snowflake", "drivers-memory-conversation-amazon-dynamodb", "drivers-memory-conversation-redis", "drivers-vector-marqo", "drivers-vector-pinecone", "drivers-vector-mongodb", "drivers-vector-redis", "drivers-vector-opensearch", "drivers-vector-amazon-opensearch", "drivers-vector-pgvector", "drivers-vector-qdrant", "drivers-vector-astra-db", "drivers-vector-pgai", "drivers-embedding-amazon-bedrock", "drivers-embedding-amazon-sagemaker", "drivers-embedding-huggingface", "drivers-embedding-voyageai", "drivers-embedding-google", "drivers-embedding-cohere", "drivers-embedding-ollama", "drivers-web-scraper-trafilatura", "drivers-web-scraper-markdownify", "drivers-web-search-duckduckgo", "drivers-web-search-tavily", "drivers-web-search-exa", "drivers-event-listener-amazon-sqs", "drivers-event-listener-amazon-iot", "drivers-event-listener-pusher", "drivers-text-to-speech-elevenlabs", "drivers-rerank-cohere", "drivers-rerank-amazon-bedrock", "drivers-observability-opentelemetry", "drivers-observability-griptape-cloud", "drivers-observability-datadog", "drivers-image-generation-huggingface", "drivers-file-manager-amazon-s3", "loaders-pdf", "loaders-image", "loaders-email", "loaders-sql", "all"] +provides-extras = ["drivers-prompt-cohere", "drivers-prompt-anthropic", "drivers-prompt-huggingface-hub", "drivers-prompt-huggingface-pipeline", "drivers-prompt-amazon-bedrock", "drivers-prompt-amazon-sagemaker", "drivers-prompt-google", "drivers-prompt-ollama", "drivers-sql", "drivers-sql-amazon-redshift", "drivers-sql-snowflake", "drivers-memory-conversation-amazon-dynamodb", "drivers-memory-conversation-redis", "drivers-vector-marqo", "drivers-vector-pinecone", "drivers-vector-mongodb", "drivers-vector-redis", "drivers-vector-opensearch", "drivers-vector-amazon-opensearch", "drivers-vector-pgvector", "drivers-vector-qdrant", "drivers-vector-astra-db", "drivers-vector-pgai", "drivers-embedding-amazon-bedrock", "drivers-embedding-amazon-sagemaker", "drivers-embedding-huggingface", "drivers-embedding-voyageai", "drivers-embedding-google", "drivers-embedding-cohere", "drivers-embedding-ollama", "drivers-embedding-twelvelabs", "drivers-web-scraper-trafilatura", "drivers-web-scraper-markdownify", "drivers-web-search-duckduckgo", "drivers-web-search-tavily", "drivers-web-search-exa", "drivers-event-listener-amazon-sqs", "drivers-event-listener-amazon-iot", "drivers-event-listener-pusher", "drivers-text-to-speech-elevenlabs", "drivers-rerank-cohere", "drivers-rerank-amazon-bedrock", "drivers-observability-opentelemetry", "drivers-observability-griptape-cloud", "drivers-observability-datadog", "drivers-image-generation-huggingface", "drivers-file-manager-amazon-s3", "loaders-pdf", "loaders-image", "loaders-email", "loaders-sql", "all"] [package.metadata.requires-dev] dev = [ @@ -5453,6 +5457,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c1/68/fa86e5a39608000f645535b2c124920126327ab731f8c4fafd5b07ff8d4b/triton-3.7.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ce061073102714b725f3660ec6939d94a1da7984b3aa99c921417cae273672f5", size = 201546766, upload-time = "2026-05-07T18:46:42.088Z" }, ] +[[package]] +name = "twelvelabs" +version = "1.2.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "pydantic" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5f/24/2604563f5a528bf5f12f9bbb7e3102e91c0e4f9bfd7aeb9963dced3a0b0f/twelvelabs-1.2.8.tar.gz", hash = "sha256:bb0846cc839845c800de8061bb7b99212a472e24bf20770d569f6f620b845e3c", size = 178449, upload-time = "2026-06-18T06:38:48.135Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/a6/6380129c65222bb2497cbef27334a496a1a31f4f38973fac46b62b9224fc/twelvelabs-1.2.8-py3-none-any.whl", hash = "sha256:c0487dd892b03728ac2afe3e4ebd4ea5e30ff404b61896667eb5a39264db282b", size = 372073, upload-time = "2026-06-18T06:38:46.747Z" }, +] + [[package]] name = "twine" version = "6.2.0"