Bedrock Knowledge Base with S3 Vector Index in AWS CDK

In this post I will explain how to create an S3 Vector Index and use it in an Amazon Bedrock Knowledge Base and ultimately in your AI agent in Bedrock, all this using AWS CDK in python.

Previous posts

I previously described how to create a Restaurant Reservation Agent in AWS CDK in these 2 posts:

• Restaurant Reservation Agent with Amazon Bedrock and AWS CDK

• Bedrock Agent with open-ended Action using SQL

and then tested it in Breaking the Restaurant Reservation Agent.

In this post, I will build on top of the previous posts, by swapping the Open Search Index with the S3 Vector Index.

The code

In this section I will explain only what is needed to create a knowledge base with S3 vectors and indexing a set of documents. The complete working example can be found in the v3 stack in the github repo of this project. There you can see how you can plug it into a Bedrock Agent.

First we define the model we will use in the knowledge base and the number of dimensions (this number depends on the model and the dimensions it supports). This model is used for indexing the documents in the knowledge base.

knowledge_base_foundation_model_vector_dimension = 1024
knowledge_base_foundation_model_id = "amazon.titan-embed-text-v2:0"

The documents to be indexed are uploaded to an S3 bucket like this

s3_bucket = s3.Bucket(
    self,
    "s3-bucket",
    bucket_name=f"{prefix}-{Aws.ACCOUNT_ID}",
    removal_policy=aws_cdk.RemovalPolicy.DESTROY,
    auto_delete_objects=True,
)

restaurant_descriptions_deployment = s3_deploy.BucketDeployment(
    self,
    "s3-deployment",
    sources=[
        s3_deploy.Source.asset(
            "./data/restaurants-v2/",
        )
    ],
    destination_bucket=s3_bucket,
    prune=True,
    retain_on_delete=False,
    destination_key_prefix="restaurants-v2/",
)

Next we create a Vector Bucket and a Vector Index

vector_index_name = "restaurant-descriptions-vector-index"

# Create the Vector Bucket and Vector Index

vector_bucket = s3vectors.CfnVectorBucket(
    self,
    "s3-vector-bucket",
    vector_bucket_name=f"{prefix}-vectors-{Aws.ACCOUNT_ID}",
)

vector_index = s3vectors.CfnIndex(
    self,
    "s3-vectors-index",
    data_type="float32",
    dimension=knowledge_base_foundation_model_vector_dimension,
    distance_metric="cosine",
    index_name=vector_index_name,
    vector_bucket_arn=vector_bucket.attr_vector_bucket_arn,
)

Before creating the knowledge base, we need to create an IAM role that

1. Allows to be assumed by the Bedrock AWS Service (bedrock.amazonaws.com). We restrict this to our own AWS account only.

2. Is allowed to invoke the model we have chosen for creating our embeddings.

3. Is allowed to read the AWS bucket where we have stored the documents to be indexed.

4. It is allowed to read from and write into the Vector Index where the embeddings will be stored.

# Define the IAM role for the Knowledge Base
embedding_model_arn = f"arn:aws:bedrock:{Aws.REGION}::foundation-model/{knowledge_base_foundation_model_id}"

knowledge_base_role = iam.Role(
    self,
    "knowledge-base-role",
    role_name=f"{prefix}-knowledge-base-role",
    assumed_by=iam.PrincipalWithConditions(
        principal=iam.ServicePrincipal("bedrock.amazonaws.com"),
        conditions={
            "StringEquals": {"aws:SourceAccount": Aws.ACCOUNT_ID},
            "ArnLike": {
                "aws:SourceArn": f"arn:aws:bedrock:{Aws.REGION}:{Aws.ACCOUNT_ID}:knowledge-base/*"
            },
        },
    ),
)

knowledge_base_role.add_to_policy(
    iam.PolicyStatement(
        effect=iam.Effect.ALLOW,
        actions=["bedrock:InvokeModel"],
        resources=[embedding_model_arn],
    )
)

knowledge_base_role.add_to_policy(
    iam.PolicyStatement(
        effect=iam.Effect.ALLOW,
        actions=["s3:ListBucket", "s3:GetObject"],
        resources=[
            s3_bucket.bucket_arn,
            s3_bucket.arn_for_objects("restaurants-v2/*"),
        ],
    )
)
knowledge_base_role.add_to_policy(
    iam.PolicyStatement(
        effect=iam.Effect.ALLOW,
        actions=[
            "s3vectors:PutVectors",
            "s3vectors:GetVectors",
            "s3vectors:DeleteVectors",
            "s3vectors:QueryVectors",
            "s3vectors:GetIndex",
        ],
        resources=[vector_index.attr_index_arn],
    )
)

Now we are ready to define the Knowledge Base.

# Define the knowledge base
restaurant_descriptions_knowledge_base = bedrock.CfnKnowledgeBase(
    self,
    "knowledge-base-restaurant-descriptions",
    name=f"{prefix}-descriptions-knowledge-base",
    role_arn=knowledge_base_role.role_arn,
    knowledge_base_configuration=bedrock.CfnKnowledgeBase.KnowledgeBaseConfigurationProperty(
        type="VECTOR",
        vector_knowledge_base_configuration=bedrock.CfnKnowledgeBase.VectorKnowledgeBaseConfigurationProperty(
            embedding_model_arn=embedding_model_arn,
        ),
    ),
    storage_configuration=bedrock.CfnKnowledgeBase.StorageConfigurationProperty(
        type="S3_VECTORS",
        s3_vectors_configuration=bedrock.CfnKnowledgeBase.S3VectorsConfigurationProperty(
            index_arn=vector_index.attr_index_arn,
            vector_bucket_arn=vector_bucket.attr_vector_bucket_arn,
        ),
    ),
)
restaurant_descriptions_knowledge_base.node.add_dependency(knowledge_base_role)

Note that I had to add a dependency to the IAM role to get rid of some permission errors I got when deploying the stack. This dependency solved the issue.

Next, we add the restaurant descriptions stored in the S3 bucket (the one with the documents not the Vector Bucket) as a data source to the knowledge base.

restaurant_descriptions_data_source = bedrock.CfnDataSource(
    self,
    "knowledge-base-data-source-restaurant-descriptions",
    name=f"{prefix}-data-source",
    knowledge_base_id=restaurant_descriptions_knowledge_base.attr_knowledge_base_id,
    # We will delete the collection anyway.
    # If we do not RETAIN the cloudformation cannot be deleted smoothly.
    data_deletion_policy="RETAIN",
    data_source_configuration=bedrock.CfnDataSource.DataSourceConfigurationProperty(
        s3_configuration=bedrock.CfnDataSource.S3DataSourceConfigurationProperty(
            bucket_arn=s3_bucket.bucket_arn,
            inclusion_prefixes=["restaurants-v2/descriptions/"],
        ),
        type="S3",
    ),
    vector_ingestion_configuration=bedrock.CfnDataSource.VectorIngestionConfigurationProperty(
        chunking_configuration=bedrock.CfnDataSource.ChunkingConfigurationProperty(
            chunking_strategy="FIXED_SIZE",
            fixed_size_chunking_configuration=bedrock.CfnDataSource.FixedSizeChunkingConfigurationProperty(
                max_tokens=300, overlap_percentage=20
            ),
        )
    ),
)
restaurant_descriptions_data_source.add_dependency(
    restaurant_descriptions_knowledge_base
)
restaurant_descriptions_data_source.node.add_dependency(
    restaurant_descriptions_deployment
)

Finally, we add a custom resource for triggering a sync on the Data Source during deployment. This will make sure that the index is populated when the stack is deployed and we will not need a manual step before using the agent.

# Sync the Data Source
sync_data_source = cr.AwsCustomResource(
    self,
    "sync-data-source",
    on_create=cr.AwsSdkCall(
        service="bedrock-agent",
        action="startIngestionJob",
        parameters={
            "dataSourceId": restaurant_descriptions_data_source.attr_data_source_id,
            "knowledgeBaseId": restaurant_descriptions_knowledge_base.attr_knowledge_base_id,
        },
        physical_resource_id=cr.PhysicalResourceId.of("Parameter.ARN"),
    ),
    policy=cr.AwsCustomResourcePolicy.from_sdk_calls(
        resources=cr.AwsCustomResourcePolicy.ANY_RESOURCE
    ),
)

Comparison with Open Search implementation

Comparing the two implementations (v2 with Open Search Index and v3 with S3 Vector Index), the S3 Vector Index requires less code and has lower fixed costs. Specifically:

• With Open Search we had to write a custom lambda function for creating the index (see here).

• With the S3 Vector Index the access policy is much simpler, we only need to add a simple IAM policy statement to the role passed to the knowledge base.

• With S3 Vector Index we only pay-as-you-go so it is much cheaper for building prototypes like this one. Previously, Open Search was the main cost for this stack and I had to make sure I deleted my stack when I was done testing.

Conclusion

Creating an S3 Vector Index and using it in Bedrock Knowledge Bases and AI Agents is very straightforward with AWS CDK. It is much simpler than setting up an index in Open Search. In addition, it offers lower cost for prototypes or non-production workflows as we pay only for what we use with lower per hour costs.