Skip to content

Embedding Drivers

Overview

Embeddings in Griptape are multidimensional representations of text data. Embeddings carry semantic information, which makes them useful for extracting relevant chunks from large bodies of text for search and querying.

Griptape provides a way to build Embedding Drivers that are reused in downstream framework components. Every Embedding Driver has two basic methods that can be used to generate embeddings:

You can optionally provide a Tokenizer via the tokenizer field to have the Driver automatically chunk the input text to fit into the token limit.

Embedding Drivers

OpenAI

The OpenAiEmbeddingDriver uses the OpenAI Embeddings API.

from griptape.drivers import OpenAiEmbeddingDriver

embeddings = OpenAiEmbeddingDriver().embed_string("Hello Griptape!")

# display the first 3 embeddings
print(embeddings[:3])

[0.0017853748286142945, 0.006118456833064556, -0.005811543669551611]

Azure OpenAI

The AzureOpenAiEmbeddingDriver uses the same parameters as OpenAiEmbeddingDriver with updated defaults.

Bedrock Titan

Info

This driver requires the drivers-embedding-amazon-bedrock extra.

The AmazonBedrockTitanEmbeddingDriver uses the Amazon Bedrock Embeddings API.

from griptape.drivers import AmazonBedrockTitanEmbeddingDriver

embeddings = AmazonBedrockTitanEmbeddingDriver().embed_string("Hello world!")

# display the first 3 embeddings
print(embeddings[:3])

[-0.234375, -0.024902344, -0.14941406]

Google

Info

This driver requires the drivers-embedding-google extra.

The GoogleEmbeddingDriver uses the Google Embeddings API.

from griptape.drivers import GoogleEmbeddingDriver

embeddings = GoogleEmbeddingDriver().embed_string("Hello world!")

# display the first 3 embeddings
print(embeddings[:3])

[0.0588633, 0.0033929371, -0.072810836]

Hugging Face Hub

Info

This driver requires the drivers-embedding-huggingface extra.

The HuggingFaceHubEmbeddingDriver connects to the Hugging Face Hub API. It supports models with the following tasks:

  • feature-extraction
import os
from griptape.drivers import HuggingFaceHubEmbeddingDriver
from griptape.tokenizers import HuggingFaceTokenizer
from transformers import AutoTokenizer

driver = HuggingFaceHubEmbeddingDriver(
    api_token=os.environ["HUGGINGFACE_HUB_ACCESS_TOKEN"],
    model="sentence-transformers/all-MiniLM-L6-v2",
    tokenizer=HuggingFaceTokenizer(
        model="sentence-transformers/all-MiniLM-L6-v2",
        max_output_tokens=512,
    ),
)

embeddings = driver.embed_string("Hello world!")

# display the first 3 embeddings
print(embeddings[:3])

Amazon SageMaker Jumpstart

The AmazonSageMakerJumpstartEmbeddingDriver uses the Amazon SageMaker Endpoints to generate embeddings on AWS.

Info

This driver requires the drivers-embedding-amazon-sagemaker extra.

PYTEST_IGNORE
import os
from griptape.drivers import AmazonSageMakerJumpstartEmbeddingDriver, SageMakerTensorFlowHubEmbeddingModelDriver

driver = AmazonSageMakerJumpstartEmbeddingDriver(
    model=os.environ["SAGEMAKER_TENSORFLOW_HUB_MODEL"],
)

embeddings = driver.embed_string("Hello world!")

# display the first 3 embeddings
print(embeddings[:3])

VoyageAI

The VoyageAiEmbeddingDriver uses the VoyageAI Embeddings API.

Info

This driver requires the drivers-embedding-voyageai extra.

import os
from griptape.drivers import VoyageAiEmbeddingDriver

driver = VoyageAiEmbeddingDriver(
    api_key=os.environ["VOYAGE_API_KEY"]
)

embeddings = driver.embed_string("Hello world!")

# display the first 3 embeddings
print(embeddings[:3])

Cohere

The CohereEmbeddingDriver uses the Cohere Embeddings API.

Info

This driver requires the drivers-embedding-cohere extra.

import os
from griptape.drivers import CohereEmbeddingDriver

embedding_driver=CohereEmbeddingDriver(
    model="embed-english-v3.0",
    api_key=os.environ["COHERE_API_KEY"],
    input_type="search_document",
)

embeddings = embedding_driver.embed_string("Hello world!")

# display the first 3 embeddings
print(embeddings[:3])

Override Default Structure Embedding Driver

Here is how you can override the Embedding Driver that is used by default in Structures. 

from griptape.structures import Agent
from griptape.tools import WebScraper, TaskMemoryClient
from griptape.drivers import (
    OpenAiChatPromptDriver,
    VoyageAiEmbeddingDriver,
)
from griptape.config import StructureConfig

agent = Agent(
    tools=[WebScraper(off_prompt=True), TaskMemoryClient(off_prompt=False)],
    config=StructureConfig(
        prompt_driver=OpenAiChatPromptDriver(model="gpt-4o"),
        embedding_driver=VoyageAiEmbeddingDriver(),
    ),
)

agent.run("based on https://www.griptape.ai/, tell me what Griptape is")