Skip to content

Commit 01f5ad0

Browse files
ewbankkitewbankkit
authored
Merge pull request hashicorp#45075 from Itsindigo/f-aws_bedrock_rds_storage_config_add_field
Add `custom_metadata_field` to `aws_bedrockagent_knowledge_base` knowledge base resource
2 parents ea3a2b2 + 6f7465c commit 01f5ad0

File tree

5 files changed

+49
-26
lines changed

5 files changed

+49
-26
lines changed

‎.changelog/45075.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
```release-note:enhancement
2+
resource/aws_bedrockagent_knowledge_base: Add `storage_configuration.rds_configuration.field_mapping.custom_metadata_field` argument
3+
```

‎internal/acctest/configs.go‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -752,17 +752,22 @@ resource "aws_iam_role_policy_attachment" "secrets_manager_read_write" {
752752
policy_arn = "arn:${data.aws_partition.current.partition}:iam::${data.aws_partition.current.partition}:policy/SecretsManagerReadWrite"
753753
}
754754
755+
data "aws_rds_orderable_db_instance" "test" {
756+
engine = "aurora-postgresql"
757+
engine_latest_version = true
758+
preferred_instance_classes = ["db.serverless"]
759+
}
760+
755761
resource "aws_rds_cluster" "test" {
756762
cluster_identifier = %[1]q
757-
engine = "aurora-postgresql"
758-
engine_mode = "provisioned"
759-
engine_version = "15.4"
760-
database_name = "test"
761763
master_username = "test"
762764
manage_master_user_password = true
765+
database_name = "test"
766+
skip_final_snapshot = true
767+
engine = data.aws_rds_orderable_db_instance.test.engine
768+
engine_version = data.aws_rds_orderable_db_instance.test.engine_version
763769
enable_http_endpoint = true
764770
vpc_security_group_ids = [aws_security_group.test.id]
765-
skip_final_snapshot = true
766771
db_subnet_group_name = aws_db_subnet_group.test.name
767772
768773
serverlessv2_scaling_configuration {
@@ -846,9 +851,10 @@ resource "null_resource" "db_setup" {
846851
psql -h ${aws_rds_cluster.test.endpoint} -U ${aws_rds_cluster.test.master_username} -d ${aws_rds_cluster.test.database_name} -c "CREATE SCHEMA IF NOT EXISTS bedrock_new;"
847852
psql -h ${aws_rds_cluster.test.endpoint} -U ${aws_rds_cluster.test.master_username} -d ${aws_rds_cluster.test.database_name} -c "CREATE ROLE bedrock_user WITH PASSWORD '$PGPASSWORD' LOGIN;"
848853
psql -h ${aws_rds_cluster.test.endpoint} -U ${aws_rds_cluster.test.master_username} -d ${aws_rds_cluster.test.database_name} -c "GRANT ALL ON SCHEMA bedrock_integration TO bedrock_user;"
849-
psql -h ${aws_rds_cluster.test.endpoint} -U ${aws_rds_cluster.test.master_username} -d ${aws_rds_cluster.test.database_name} -c "CREATE TABLE bedrock_integration.bedrock_kb (id uuid PRIMARY KEY, embedding vector(1536), chunks text, metadata json);"
854+
psql -h ${aws_rds_cluster.test.endpoint} -U ${aws_rds_cluster.test.master_username} -d ${aws_rds_cluster.test.database_name} -c "CREATE TABLE bedrock_integration.bedrock_kb (id uuid PRIMARY KEY, embedding vector(1536), chunks text, metadata json, custom_metadata jsonb);"
850855
psql -h ${aws_rds_cluster.test.endpoint} -U ${aws_rds_cluster.test.master_username} -d ${aws_rds_cluster.test.database_name} -c "CREATE INDEX ON bedrock_integration.bedrock_kb USING hnsw (embedding vector_cosine_ops);"
851856
psql -h ${aws_rds_cluster.test.endpoint} -U ${aws_rds_cluster.test.master_username} -d ${aws_rds_cluster.test.database_name} -c "CREATE INDEX ON bedrock_integration.bedrock_kb USING gin (to_tsvector('simple', chunks));"
857+
psql -h ${aws_rds_cluster.test.endpoint} -U ${aws_rds_cluster.test.master_username} -d ${aws_rds_cluster.test.database_name} -c "CREATE INDEX ON bedrock_integration.bedrock_kb USING gin (custom_metadata);"
852858
EOT
853859
}
854860
}

‎internal/service/bedrockagent/knowledge_base.go‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -336,6 +336,12 @@ func (r *knowledgeBaseResource) Schema(ctx context.Context, request resource.Sch
336336
},
337337
NestedObject: schema.NestedBlockObject{
338338
Attributes: map[string]schema.Attribute{
339+
"custom_metadata_field": schema.StringAttribute{
340+
Optional: true,
341+
PlanModifiers: []planmodifier.String{
342+
stringplanmodifier.RequiresReplace(),
343+
},
344+
},
339345
"metadata_field": schema.StringAttribute{
340346
Required: true,
341347
PlanModifiers: []planmodifier.String{
@@ -870,10 +876,11 @@ type rdsConfigurationModel struct {
870876
}
871877

872878
type rdsFieldMappingModel struct {
873-
MetadataField types.String `tfsdk:"metadata_field"`
874-
PrimaryKeyField types.String `tfsdk:"primary_key_field"`
875-
TextField types.String `tfsdk:"text_field"`
876-
VectorField types.String `tfsdk:"vector_field"`
879+
CustomMetadataField types.String `tfsdk:"custom_metadata_field"`
880+
MetadataField types.String `tfsdk:"metadata_field"`
881+
PrimaryKeyField types.String `tfsdk:"primary_key_field"`
882+
TextField types.String `tfsdk:"text_field"`
883+
VectorField types.String `tfsdk:"vector_field"`
877884
}
878885

879886
type redisEnterpriseCloudConfigurationModel struct {

‎internal/service/bedrockagent/knowledge_base_test.go‎

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ func testAccKnowledgeBase_basic(t *testing.T) {
7777
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.rds_configuration.0.field_mapping.0.text_field", "chunks"),
7878
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.rds_configuration.0.field_mapping.0.metadata_field", "metadata"),
7979
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.rds_configuration.0.field_mapping.0.primary_key_field", names.AttrID),
80+
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.rds_configuration.0.field_mapping.0.custom_metadata_field", "custom_metadata"),
8081
),
8182
},
8283
{
@@ -108,6 +109,7 @@ func testAccKnowledgeBase_basic(t *testing.T) {
108109
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.rds_configuration.0.field_mapping.0.text_field", "chunks"),
109110
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.rds_configuration.0.field_mapping.0.metadata_field", "metadata"),
110111
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.rds_configuration.0.field_mapping.0.primary_key_field", names.AttrID),
112+
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.rds_configuration.0.field_mapping.0.custom_metadata_field", "custom_metadata"),
111113
),
112114
},
113115
},
@@ -515,10 +517,11 @@ resource "aws_bedrockagent_knowledge_base" "test" {
515517
database_name = aws_rds_cluster.test.database_name
516518
table_name = "bedrock_integration.bedrock_kb"
517519
field_mapping {
518-
vector_field = "embedding"
519-
text_field = "chunks"
520-
metadata_field = "metadata"
521-
primary_key_field = "id"
520+
vector_field = "embedding"
521+
text_field = "chunks"
522+
metadata_field = "metadata"
523+
primary_key_field = "id"
524+
custom_metadata_field = "custom_metadata"
522525
}
523526
}
524527
}
@@ -551,10 +554,11 @@ resource "aws_bedrockagent_knowledge_base" "test" {
551554
database_name = aws_rds_cluster.test.database_name
552555
table_name = "bedrock_integration.bedrock_kb"
553556
field_mapping {
554-
vector_field = "embedding"
555-
text_field = "chunks"
556-
metadata_field = "metadata"
557-
primary_key_field = "id"
557+
vector_field = "embedding"
558+
text_field = "chunks"
559+
metadata_field = "metadata"
560+
primary_key_field = "id"
561+
custom_metadata_field = "custom_metadata"
558562
}
559563
}
560564
}
@@ -585,10 +589,11 @@ resource "aws_bedrockagent_knowledge_base" "test" {
585589
database_name = aws_rds_cluster.test.database_name
586590
table_name = "bedrock_integration.bedrock_kb"
587591
field_mapping {
588-
vector_field = "embedding"
589-
text_field = "chunks"
590-
metadata_field = "metadata"
591-
primary_key_field = "id"
592+
vector_field = "embedding"
593+
text_field = "chunks"
594+
metadata_field = "metadata"
595+
primary_key_field = "id"
596+
custom_metadata_field = "custom_metadata"
592597
}
593598
}
594599
}
@@ -623,10 +628,11 @@ resource "aws_bedrockagent_knowledge_base" "test" {
623628
database_name = aws_rds_cluster.test.database_name
624629
table_name = "bedrock_integration.bedrock_kb"
625630
field_mapping {
626-
vector_field = "embedding"
627-
text_field = "chunks"
628-
metadata_field = "metadata"
629-
primary_key_field = "id"
631+
vector_field = "embedding"
632+
text_field = "chunks"
633+
metadata_field = "metadata"
634+
primary_key_field = "id"
635+
custom_metadata_field = "custom_metadata"
630636
}
631637
}
632638
}

‎website/docs/r/bedrockagent_knowledge_base.html.markdown‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,7 @@ The `pinecone_configuration` configuration block supports the following argument
184184
* `credentials_secret_arn` - (Required) ARN of the secret that you created in AWS Secrets Manager that is linked to your Amazon RDS database.
185185
* `database_name` - (Required) Name of your Amazon RDS database.
186186
* `field_mapping` - (Required) Names of the fields to which to map information about the vector store. This block supports the following arguments:
187+
* `custom_metadata_field` - (Optional) Name for the universal metadata field where Amazon Bedrock will store any custom metadata from your data source.
187188
* `metadata_field` - (Required) Name of the field in which Amazon Bedrock stores metadata about the vector store.
188189
* `primary_key_field` - (Required) Name of the field in which Amazon Bedrock stores the ID for each entry.
189190
* `text_field` - (Required) Name of the field in which Amazon Bedrock stores the raw text from your data. The text is split according to the chunking strategy you choose.

0 commit comments

Comments
 (0)