Skip to content

Commit

Permalink
F aws bedrock foundation model datasource (#34148)
Browse files Browse the repository at this point in the history
* move bedrock to v2

* make gen

* add dependencies

* add aws_bedrock_foundation_models datasource

* add new datasource aws_bedrock_foundation_model

* make gen

* add docs

* add docs

* linting fixes

* chore: changelog

* d/aws_bedrock_foundation_models: awstypes alias

* d/aws_bedrock: add name to resource annotations

* d/aws_bedrock: tidy skaff comments

* d/aws_bedrock_foundation_models: alphabetize attributes

* d/aws_bedrock_foundation_model: alphabetize attributes

* d/aws_bedrock_foundation_models(doc): tidy

* d/aws_bedrock_foundation_model(doc): tidy

* d/aws_bedrock_foundation_model(doc): markdownlint

* d/aws_bedrock_foundation_model: custom string type conversion helper

* d/aws_bedrock_foundation_models: add filters on ListFoundationModels

* d/aws_bedrock_foundation_models: simplify flatten func

* d/aws_bedrock_foundation_models(test): add tests with filters

* chore: fix moved function after merge

* d/aws_bedrock_foundation_models(doc): document id attribute

* d/aws_bedrock: use diags constructor funcs

---------

Co-authored-by: Jared Baker <jared.baker@hashicorp.com>
  • Loading branch information
skyscrapr and jar-b committed Nov 6, 2023
1 parent da1568d commit 6bf4934
Show file tree
Hide file tree
Showing 12 changed files with 619 additions and 12 deletions.
6 changes: 6 additions & 0 deletions .changelog/34148.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
```release-note:new-data-source
aws_bedrock_foundation_model
```
```release-note:new-data-source
aws_bedrock_foundation_models
```
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ require (
github.com/aws/aws-sdk-go-v2/service/appconfig v1.25.0
github.com/aws/aws-sdk-go-v2/service/athena v1.34.0
github.com/aws/aws-sdk-go-v2/service/auditmanager v1.29.0
github.com/aws/aws-sdk-go-v2/service/bedrock v1.3.0
github.com/aws/aws-sdk-go-v2/service/cleanrooms v1.6.0
github.com/aws/aws-sdk-go-v2/service/cloudcontrol v1.14.0
github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs v1.26.0
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ github.com/aws/aws-sdk-go-v2/service/athena v1.34.0 h1:wZaywaONfHerySRsAFWhA1Myo
github.com/aws/aws-sdk-go-v2/service/athena v1.34.0/go.mod h1:WSRHaccpHc65b4OP63dElo5nLuMUoGhoDDzYr5tP5/k=
github.com/aws/aws-sdk-go-v2/service/auditmanager v1.29.0 h1:CIG6h8XkwGaS9/k4z+B72Rzh9KdlyrD/HzbidUD/C9U=
github.com/aws/aws-sdk-go-v2/service/auditmanager v1.29.0/go.mod h1:xXtiggqGVD59WdBNclsih1qfeeXPOasMzpBKfj7k+wE=
github.com/aws/aws-sdk-go-v2/service/bedrock v1.3.0 h1:RTPbj0iKv6FSecxR+vH3V944Wt4mufsGuUd+tZy1M8Y=
github.com/aws/aws-sdk-go-v2/service/bedrock v1.3.0/go.mod h1:TSSHBioEPEx4wRuWyT5nFdRPjHXiNdixNr1wyNIC61s=
github.com/aws/aws-sdk-go-v2/service/cleanrooms v1.6.0 h1:CVezvdO+/Aa1S+lQpNqe1nHnn4a9QPxr5LVKMGYUDmo=
github.com/aws/aws-sdk-go-v2/service/cleanrooms v1.6.0/go.mod h1:Zi1s46alfbcte2MyRLdlxuv5+Z6YKJ+ODPggBnSZWhc=
github.com/aws/aws-sdk-go-v2/service/cloudcontrol v1.14.0 h1:OkgbfNtZKs8cqCOA8cZrw1GLYR9/ghlRDnKkdzADo+E=
Expand Down
6 changes: 3 additions & 3 deletions internal/conns/awsclient_gen.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

137 changes: 137 additions & 0 deletions internal/service/bedrock/foundation_model_data_source.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,137 @@
// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package bedrock

import (
"context"

"github.com/aws/aws-sdk-go-v2/service/bedrock"
"github.com/hashicorp/terraform-plugin-framework/datasource"
"github.com/hashicorp/terraform-plugin-framework/datasource/schema"
"github.com/hashicorp/terraform-plugin-framework/types"
"github.com/hashicorp/terraform-provider-aws/internal/create"
"github.com/hashicorp/terraform-provider-aws/internal/framework"
"github.com/hashicorp/terraform-provider-aws/internal/framework/flex"
"github.com/hashicorp/terraform-provider-aws/names"
)

const DSNameFoundationModel = "Foundation Model Data Source"

// @FrameworkDataSource(name="Foundation Model")
func newDataSourceFoundationModel(context.Context) (datasource.DataSourceWithConfigure, error) {
return &dataSourceFoundationModel{}, nil
}

type dataSourceFoundationModel struct {
framework.DataSourceWithConfigure
}

func (d *dataSourceFoundationModel) Metadata(_ context.Context, request datasource.MetadataRequest, response *datasource.MetadataResponse) {
response.TypeName = "aws_bedrock_foundation_model"
}

func (d *dataSourceFoundationModel) Schema(ctx context.Context, req datasource.SchemaRequest, resp *datasource.SchemaResponse) {
resp.Schema = schema.Schema{
Attributes: map[string]schema.Attribute{
"customizations_supported": schema.SetAttribute{
ElementType: types.StringType,
Computed: true,
},
"id": schema.StringAttribute{
Computed: true,
},
"inference_types_supported": schema.SetAttribute{
ElementType: types.StringType,
Computed: true,
},
"input_modalities": schema.SetAttribute{
ElementType: types.StringType,
Computed: true,
},
"model_arn": schema.StringAttribute{
Computed: true,
},
"model_id": schema.StringAttribute{
Required: true,
},
"model_name": schema.StringAttribute{
Computed: true,
},
"output_modalities": schema.SetAttribute{
ElementType: types.StringType,
Computed: true,
},
"provider_name": schema.StringAttribute{
Computed: true,
},
"response_streaming_supported": schema.BoolAttribute{
Computed: true,
},
},
}
}

func (d *dataSourceFoundationModel) Read(ctx context.Context, request datasource.ReadRequest, response *datasource.ReadResponse) {
conn := d.Meta().BedrockClient(ctx)

var data foundationModel
response.Diagnostics.Append(request.Config.Get(ctx, &data)...)
if response.Diagnostics.HasError() {
return
}

input := &bedrock.GetFoundationModelInput{
ModelIdentifier: data.ModelID.ValueStringPointer(),
}
model, err := conn.GetFoundationModel(ctx, input)
if err != nil {
response.Diagnostics.AddError(
create.ProblemStandardMessage(names.Bedrock, create.ErrActionReading, DSNameFoundationModel, data.ModelID.String(), err),
err.Error(),
)
return
}

data.refreshFromOutput(ctx, model)
response.Diagnostics.Append(response.State.Set(ctx, &data)...)
}

type foundationModel struct {
CustomizationsSupported types.Set `tfsdk:"customizations_supported"`
ID types.String `tfsdk:"id"`
InferenceTypesSupported types.Set `tfsdk:"inference_types_supported"`
InputModalities types.Set `tfsdk:"input_modalities"`
ModelArn types.String `tfsdk:"model_arn"`
ModelID types.String `tfsdk:"model_id"`
ModelName types.String `tfsdk:"model_name"`
OutputModalities types.Set `tfsdk:"output_modalities"`
ProviderName types.String `tfsdk:"provider_name"`
ResponseStreamingSupported types.Bool `tfsdk:"response_streaming_supported"`
}

func (data *foundationModel) refreshFromOutput(ctx context.Context, model *bedrock.GetFoundationModelOutput) {
if model == nil {
return
}

data.ID = flex.StringToFramework(ctx, model.ModelDetails.ModelId)
data.ModelArn = flex.StringToFramework(ctx, model.ModelDetails.ModelArn)
data.ModelID = flex.StringToFramework(ctx, model.ModelDetails.ModelId)
data.ModelName = flex.StringToFramework(ctx, model.ModelDetails.ModelName)
data.ProviderName = flex.StringToFramework(ctx, model.ModelDetails.ProviderName)
data.CustomizationsSupported = flex.FlattenFrameworkStringValueSet(ctx, toStringSlice(model.ModelDetails.CustomizationsSupported))
data.InferenceTypesSupported = flex.FlattenFrameworkStringValueSet(ctx, toStringSlice(model.ModelDetails.InferenceTypesSupported))
data.InputModalities = flex.FlattenFrameworkStringValueSet(ctx, toStringSlice(model.ModelDetails.InputModalities))
data.OutputModalities = flex.FlattenFrameworkStringValueSet(ctx, toStringSlice(model.ModelDetails.OutputModalities))
data.ResponseStreamingSupported = flex.BoolToFramework(ctx, model.ModelDetails.ResponseStreamingSupported)
}

// toStringSlice converts a slice of custom string types to a slice of strings
func toStringSlice[T ~string](values []T) []string {
var out []string
for _, v := range values {
out = append(out, string(v))
}
return out
}
39 changes: 39 additions & 0 deletions internal/service/bedrock/foundation_model_data_source_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package bedrock_test

import (
"testing"

"github.com/hashicorp/terraform-plugin-testing/helper/resource"
"github.com/hashicorp/terraform-provider-aws/internal/acctest"
)

func TestAccBedrockFoundationModelDataSource_basic(t *testing.T) {
ctx := acctest.Context(t)

resource.ParallelTest(t, resource.TestCase{
PreCheck: func() { acctest.PreCheck(ctx, t) },
ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
Steps: []resource.TestStep{
{
Config: testAccFoundationModelDataSourceConfig_basic(),
Check: resource.ComposeTestCheckFunc(
resource.TestCheckResourceAttrSet("data.aws_bedrock_foundation_model.test", "id"),
resource.TestCheckResourceAttrSet("data.aws_bedrock_foundation_model.test", "model_id"),
),
},
},
})
}

func testAccFoundationModelDataSourceConfig_basic() string {
return `
data "aws_bedrock_foundation_models" "test" {}
data "aws_bedrock_foundation_model" "test" {
model_id = data.aws_bedrock_foundation_models.test.model_summaries[0].model_id
}
`
}

0 comments on commit 6bf4934

Please sign in to comment.