Skip to content

Commit 4691f65

Browse files
ewbankkitewbankkit
committed
r/aws_bedrockagent_knowledge_base: Additions from #45468.
1 parent bedb43b commit 4691f65

File tree

4 files changed

+326
-1
lines changed

4 files changed

+326
-1
lines changed

‎internal/service/bedrockagent/bedrockagent_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ func TestAccBedrockAgent_serial(t *testing.T) {
2222
"OpenSearchServerlessSupplementalDataStorage": testAccKnowledgeBase_OpenSearchServerless_supplementalDataStorage,
2323
"KendraBasic": testAccKnowledgeBase_Kendra_basic,
2424
"OpenSearchManagedClusterBasic": testAccKnowledgeBase_OpenSearchManagedCluster_basic,
25+
"S3Vectors": testAccKnowledgeBase_S3Vectors_update,
2526
},
2627
"DataSource": {
2728
acctest.CtBasic: testAccDataSource_basic,

‎internal/service/bedrockagent/knowledge_base.go‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717
"github.com/hashicorp/terraform-plugin-framework-timeouts/resource/timeouts"
1818
"github.com/hashicorp/terraform-plugin-framework-timetypes/timetypes"
1919
"github.com/hashicorp/terraform-plugin-framework-validators/listvalidator"
20+
"github.com/hashicorp/terraform-plugin-framework-validators/objectvalidator"
2021
"github.com/hashicorp/terraform-plugin-framework-validators/stringvalidator"
2122
"github.com/hashicorp/terraform-plugin-framework/path"
2223
"github.com/hashicorp/terraform-plugin-framework/resource"
@@ -260,6 +261,7 @@ func (r *knowledgeBaseResource) Schema(ctx context.Context, request resource.Sch
260261
path.MatchRelative().AtParent().AtName("pinecone_configuration"),
261262
path.MatchRelative().AtParent().AtName("rds_configuration"),
262263
path.MatchRelative().AtParent().AtName("redis_enterprise_cloud_configuration"),
264+
path.MatchRelative().AtParent().AtName("s3_vectors_configuration"),
263265
),
264266
},
265267
PlanModifiers: []planmodifier.List{
@@ -605,6 +607,53 @@ func (r *knowledgeBaseResource) Schema(ctx context.Context, request resource.Sch
605607
},
606608
},
607609
},
610+
"s3_vectors_configuration": schema.ListNestedBlock{
611+
CustomType: fwtypes.NewListNestedObjectTypeOf[s3VectorsConfigurationModel](ctx),
612+
Validators: []validator.List{
613+
listvalidator.SizeAtMost(1),
614+
},
615+
PlanModifiers: []planmodifier.List{
616+
listplanmodifier.RequiresReplace(),
617+
},
618+
NestedObject: schema.NestedBlockObject{
619+
Attributes: map[string]schema.Attribute{
620+
"index_arn": schema.StringAttribute{
621+
CustomType: fwtypes.ARNType,
622+
Optional: true,
623+
PlanModifiers: []planmodifier.String{
624+
stringplanmodifier.RequiresReplace(),
625+
},
626+
Validators: []validator.String{
627+
stringvalidator.ConflictsWith(path.MatchRelative().AtParent().AtName("index_name"), path.MatchRelative().AtParent().AtName("vector_bucket_arn")),
628+
},
629+
},
630+
"index_name": schema.StringAttribute{
631+
Optional: true,
632+
PlanModifiers: []planmodifier.String{
633+
stringplanmodifier.RequiresReplace(),
634+
},
635+
Validators: []validator.String{
636+
stringvalidator.AlsoRequires(path.MatchRelative().AtParent().AtName("vector_bucket_arn")),
637+
stringvalidator.ConflictsWith(path.MatchRelative().AtParent().AtName("index_arn")),
638+
},
639+
},
640+
"vector_bucket_arn": schema.StringAttribute{
641+
CustomType: fwtypes.ARNType,
642+
Optional: true,
643+
PlanModifiers: []planmodifier.String{
644+
stringplanmodifier.RequiresReplace(),
645+
},
646+
Validators: []validator.String{
647+
stringvalidator.AlsoRequires(path.MatchRelative().AtParent().AtName("index_name")),
648+
stringvalidator.ConflictsWith(path.MatchRelative().AtParent().AtName("index_arn")),
649+
},
650+
},
651+
},
652+
Validators: []validator.Object{
653+
objectvalidator.AtLeastOneOf(path.MatchRelative().AtName("index_arn"), path.MatchRelative().AtName("index_name")),
654+
},
655+
},
656+
},
608657
},
609658
},
610659
},
@@ -966,6 +1015,7 @@ type storageConfigurationModel struct {
9661015
PineconeConfiguration fwtypes.ListNestedObjectValueOf[pineconeConfigurationModel] `tfsdk:"pinecone_configuration"`
9671016
RDSConfiguration fwtypes.ListNestedObjectValueOf[rdsConfigurationModel] `tfsdk:"rds_configuration"`
9681017
RedisEnterpriseCloudConfiguration fwtypes.ListNestedObjectValueOf[redisEnterpriseCloudConfigurationModel] `tfsdk:"redis_enterprise_cloud_configuration"`
1018+
S3VectorsConfiguration fwtypes.ListNestedObjectValueOf[s3VectorsConfigurationModel] `tfsdk:"s3_vectors_configuration"`
9691019
Type fwtypes.StringEnum[awstypes.KnowledgeBaseStorageType] `tfsdk:"type"`
9701020
}
9711021

@@ -1034,3 +1084,9 @@ type redisEnterpriseCloudFieldMappingModel struct {
10341084
TextField types.String `tfsdk:"text_field"`
10351085
VectorField types.String `tfsdk:"vector_field"`
10361086
}
1087+
1088+
type s3VectorsConfigurationModel struct {
1089+
IndexARN fwtypes.ARN `tfsdk:"index_arn"`
1090+
IndexName types.String `tfsdk:"index_name"`
1091+
VectorBucketARN fwtypes.ARN `tfsdk:"vector_bucket_arn"`
1092+
}

‎internal/service/bedrockagent/knowledge_base_test.go‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -488,6 +488,79 @@ func testAccKnowledgeBase_OpenSearchManagedCluster_basic(t *testing.T) {
488488
"pinecone_configuration": knownvalue.ListSizeExact(0),
489489
"rds_configuration": knownvalue.ListSizeExact(0),
490490
"redis_enterprise_cloud_configuration": knownvalue.ListSizeExact(0),
491+
"s3_vectors": knownvalue.ListSizeExact(0),
492+
}),
493+
})),
494+
},
495+
},
496+
},
497+
})
498+
}
499+
500+
func testAccKnowledgeBase_S3Vectors_update(t *testing.T) {
501+
ctx := acctest.Context(t)
502+
var knowledgebase awstypes.KnowledgeBase
503+
rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
504+
resourceName := "aws_bedrockagent_knowledge_base.test"
505+
foundationModel := "amazon.titan-embed-text-v2:0"
506+
507+
resource.Test(t, resource.TestCase{
508+
PreCheck: func() {
509+
acctest.PreCheck(ctx, t)
510+
},
511+
ErrorCheck: acctest.ErrorCheck(t, names.BedrockAgentServiceID),
512+
ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
513+
CheckDestroy: testAccCheckKnowledgeBaseDestroy(ctx),
514+
Steps: []resource.TestStep{
515+
{
516+
Config: testAccKnowledgeBaseConfig_S3VectorsByIndexARN(rName, foundationModel),
517+
ConfigPlanChecks: resource.ConfigPlanChecks{
518+
PreApply: []plancheck.PlanCheck{
519+
plancheck.ExpectResourceAction(resourceName, plancheck.ResourceActionCreate),
520+
},
521+
},
522+
Check: resource.ComposeTestCheckFunc(
523+
testAccCheckKnowledgeBaseExists(ctx, resourceName, &knowledgebase),
524+
),
525+
ConfigStateChecks: []statecheck.StateCheck{
526+
statecheck.ExpectKnownValue(resourceName, tfjsonpath.New("storage_configuration"), knownvalue.ListExact([]knownvalue.Check{
527+
knownvalue.MapExact(map[string]knownvalue.Check{
528+
"opensearch_managed_cluster_configuration": knownvalue.ListSizeExact(0),
529+
"opensearch_serverless_configuration": knownvalue.ListSizeExact(0),
530+
names.AttrType: tfknownvalue.StringExact(awstypes.KnowledgeBaseStorageTypeS3Vectors),
531+
"pinecone_configuration": knownvalue.ListSizeExact(0),
532+
"rds_configuration": knownvalue.ListSizeExact(0),
533+
"redis_enterprise_cloud_configuration": knownvalue.ListSizeExact(0),
534+
"s3_vectors_configuration": knownvalue.ListSizeExact(1),
535+
}),
536+
})),
537+
},
538+
},
539+
{
540+
ResourceName: resourceName,
541+
ImportState: true,
542+
ImportStateVerify: true,
543+
},
544+
{
545+
Config: testAccKnowledgeBaseConfig_S3VectorsByIndexName(rName, foundationModel),
546+
ConfigPlanChecks: resource.ConfigPlanChecks{
547+
PreApply: []plancheck.PlanCheck{
548+
plancheck.ExpectResourceAction(resourceName, plancheck.ResourceActionDestroyBeforeCreate),
549+
},
550+
},
551+
Check: resource.ComposeTestCheckFunc(
552+
testAccCheckKnowledgeBaseExists(ctx, resourceName, &knowledgebase),
553+
),
554+
ConfigStateChecks: []statecheck.StateCheck{
555+
statecheck.ExpectKnownValue(resourceName, tfjsonpath.New("storage_configuration"), knownvalue.ListExact([]knownvalue.Check{
556+
knownvalue.MapExact(map[string]knownvalue.Check{
557+
"opensearch_managed_cluster_configuration": knownvalue.ListSizeExact(0),
558+
"opensearch_serverless_configuration": knownvalue.ListSizeExact(0),
559+
names.AttrType: tfknownvalue.StringExact(awstypes.KnowledgeBaseStorageTypeS3Vectors),
560+
"pinecone_configuration": knownvalue.ListSizeExact(0),
561+
"rds_configuration": knownvalue.ListSizeExact(0),
562+
"redis_enterprise_cloud_configuration": knownvalue.ListSizeExact(0),
563+
"s3_vectors_configuration": knownvalue.ListSizeExact(1),
491564
}),
492565
})),
493566
},
@@ -1416,3 +1489,140 @@ resource "aws_bedrockagent_knowledge_base" "test" {
14161489
}
14171490
`, rName, model))
14181491
}
1492+
1493+
func testAccKnowledgeBaseConfig_baseS3Vectors(rName string) string {
1494+
return fmt.Sprintf(`
1495+
data "aws_region" "current" {}
1496+
data "aws_partition" "current" {}
1497+
1498+
data "aws_iam_policy_document" "assume_role_bedrock" {
1499+
statement {
1500+
effect = "Allow"
1501+
principals {
1502+
type = "Service"
1503+
identifiers = ["bedrock.amazonaws.com"]
1504+
}
1505+
actions = ["sts:AssumeRole"]
1506+
}
1507+
}
1508+
1509+
data "aws_iam_policy_document" "bedrock" {
1510+
statement {
1511+
effect = "Allow"
1512+
actions = ["bedrock:InvokeModel"]
1513+
resources = ["*"]
1514+
}
1515+
statement {
1516+
effect = "Allow"
1517+
actions = ["s3:ListBucket", "s3:GetObject"]
1518+
resources = ["*"]
1519+
}
1520+
statement {
1521+
effect = "Allow"
1522+
actions = [
1523+
"s3vectors:GetIndex",
1524+
"s3vectors:QueryVectors",
1525+
"s3vectors:PutVectors",
1526+
"s3vectors:GetVectors",
1527+
"s3vectors:DeleteVectors"
1528+
]
1529+
resources = ["*"]
1530+
}
1531+
}
1532+
1533+
resource "aws_iam_role" "test" {
1534+
assume_role_policy = data.aws_iam_policy_document.assume_role_bedrock.json
1535+
name = %[1]q
1536+
}
1537+
1538+
resource "aws_iam_role_policy" "test" {
1539+
role = aws_iam_role.test.name
1540+
policy = data.aws_iam_policy_document.bedrock.json
1541+
}
1542+
1543+
resource "aws_s3vectors_vector_bucket" "test" {
1544+
vector_bucket_name = %[1]q
1545+
force_destroy = true
1546+
}
1547+
1548+
resource "aws_s3vectors_index" "test" {
1549+
index_name = %[1]q
1550+
vector_bucket_name = aws_s3vectors_vector_bucket.test.vector_bucket_name
1551+
1552+
data_type = "float32"
1553+
dimension = 256
1554+
distance_metric = "euclidean"
1555+
}
1556+
`, rName)
1557+
}
1558+
1559+
func testAccKnowledgeBaseConfig_S3VectorsByIndexARN(rName, model string) string {
1560+
return acctest.ConfigCompose(testAccKnowledgeBaseConfig_baseS3Vectors(rName), fmt.Sprintf(`
1561+
resource "aws_bedrockagent_knowledge_base" "test" {
1562+
depends_on = [
1563+
aws_iam_role_policy.test,
1564+
]
1565+
1566+
name = %[1]q
1567+
role_arn = aws_iam_role.test.arn
1568+
1569+
knowledge_base_configuration {
1570+
type = "VECTOR"
1571+
1572+
vector_knowledge_base_configuration {
1573+
embedding_model_arn = "arn:${data.aws_partition.current.partition}:bedrock:${data.aws_region.current.region}::foundation-model/%[2]s"
1574+
embedding_model_configuration {
1575+
bedrock_embedding_model_configuration {
1576+
dimensions = 256
1577+
embedding_data_type = "FLOAT32"
1578+
}
1579+
}
1580+
}
1581+
}
1582+
1583+
storage_configuration {
1584+
type = "S3_VECTORS"
1585+
1586+
s3_vectors_configuration {
1587+
index_arn = aws_s3vectors_index.test.index_arn
1588+
}
1589+
}
1590+
}
1591+
`, rName, model))
1592+
}
1593+
1594+
func testAccKnowledgeBaseConfig_S3VectorsByIndexName(rName, model string) string {
1595+
return acctest.ConfigCompose(testAccKnowledgeBaseConfig_baseS3Vectors(rName), fmt.Sprintf(`
1596+
resource "aws_bedrockagent_knowledge_base" "test" {
1597+
depends_on = [
1598+
aws_iam_role_policy.test,
1599+
]
1600+
1601+
name = %[1]q
1602+
role_arn = aws_iam_role.test.arn
1603+
1604+
knowledge_base_configuration {
1605+
type = "VECTOR"
1606+
1607+
vector_knowledge_base_configuration {
1608+
embedding_model_arn = "arn:${data.aws_partition.current.partition}:bedrock:${data.aws_region.current.region}::foundation-model/%[2]s"
1609+
embedding_model_configuration {
1610+
bedrock_embedding_model_configuration {
1611+
dimensions = 256
1612+
embedding_data_type = "FLOAT32"
1613+
}
1614+
}
1615+
}
1616+
}
1617+
1618+
storage_configuration {
1619+
type = "S3_VECTORS"
1620+
1621+
s3_vectors_configuration {
1622+
index_name = aws_s3vectors_index.test.index_name
1623+
vector_bucket_arn = aws_s3vectors_vector_bucket.test.vector_bucket_arn
1624+
}
1625+
}
1626+
}
1627+
`, rName, model))
1628+
}

0 commit comments

Comments
 (0)