Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: Add function schema and return of control support to aws_bedrockagent_agent_action_group #37484

Merged
merged 4 commits into from
Jul 18, 2024
Merged

feat: Add function schema and return of control support to aws_bedrockagent_agent_action_group #37484

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b2c3d44
feat: Add function schema support to aws_bedrockagent_agent_action_group
acwwat May 13, 2024
751bee7
feat: Add return of control support to aws_bedrockagent_agent_action_…
acwwat May 20, 2024
11ab4cf
Merge branch 'main' into HEAD
ewbankkit Jul 18, 2024
cb2d0ef
Run 'make fix-constants PKG=bedrockagent'.
ewbankkit Jul 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions .changelog/37484.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
```release-note:enhancement
resource/aws_bedrockagent_agent_action_group: Add `function_schema` configuration block
```

```release-note:enhancement
resource/aws_bedrockagent_agent_action_group: Add `action_group_executor.custom_control` argument
```
135 changes: 134 additions & 1 deletion internal/service/bedrockagent/agent_action_group.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,10 @@ func (r *agentActionGroupResource) Schema(ctx context.Context, request resource.
},
NestedObject: schema.NestedBlockObject{
Attributes: map[string]schema.Attribute{
"custom_control": schema.StringAttribute{
CustomType: fwtypes.StringEnumType[awstypes.CustomControlMethod](),
Optional: true,
},
"lambda": schema.StringAttribute{
CustomType: fwtypes.ARNType,
Optional: true,
Expand Down Expand Up @@ -151,6 +155,63 @@ func (r *agentActionGroupResource) Schema(ctx context.Context, request resource.
},
},
},
"function_schema": schema.ListNestedBlock{
CustomType: fwtypes.NewListNestedObjectTypeOf[functionSchemaModel](ctx),
Validators: []validator.List{
listvalidator.SizeAtMost(1),
},
NestedObject: schema.NestedBlockObject{
Blocks: map[string]schema.Block{
"functions": schema.ListNestedBlock{
CustomType: fwtypes.NewListNestedObjectTypeOf[functionModel](ctx),
NestedObject: schema.NestedBlockObject{
Attributes: map[string]schema.Attribute{
names.AttrDescription: schema.StringAttribute{
Optional: true,
Validators: []validator.String{
stringvalidator.LengthBetween(1, 1200),
},
},
names.AttrName: schema.StringAttribute{
Required: true,
Validators: []validator.String{
stringvalidator.RegexMatches(regexache.MustCompile(`^([0-9a-zA-Z][_-]?){1,100}$`), "valid characters are a-z, A-Z, 0-9, _ (underscore) and - (hyphen). The name can have up to 100 characters"),
},
},
},
Blocks: map[string]schema.Block{
names.AttrParameters: schema.SetNestedBlock{
CustomType: fwtypes.NewSetNestedObjectTypeOf[parameterDetailModel](ctx),
NestedObject: schema.NestedBlockObject{
Attributes: map[string]schema.Attribute{
names.AttrDescription: schema.StringAttribute{
Optional: true,
Validators: []validator.String{
stringvalidator.LengthBetween(1, 500),
},
},
"map_block_key": schema.StringAttribute{
Required: true,
Validators: []validator.String{
stringvalidator.RegexMatches(regexache.MustCompile(`^([0-9a-zA-Z][_-]?){1,100}$`), "valid characters are a-z, A-Z, 0-9, _ (underscore) and - (hyphen). The name can have up to 100 characters"),
},
},
"required": schema.BoolAttribute{
Optional: true,
},
names.AttrType: schema.StringAttribute{
Required: true,
CustomType: fwtypes.StringEnumType[awstypes.Type](),
},
},
},
},
},
},
},
},
},
},
},
}
}
Expand Down Expand Up @@ -224,6 +285,12 @@ func (r *agentActionGroupResource) Read(ctx context.Context, request resource.Re
// AutoFlEx doesn't yet handle union types.
data.ActionGroupExecutor = flattenActionGroupExecutor(ctx, output.ActionGroupExecutor)
data.APISchema = flattenAPISchema(ctx, output.ApiSchema)
functionSchema, diags := flattenFunctionSchema(ctx, output.FunctionSchema)
response.Diagnostics.Append(diags...)
if response.Diagnostics.HasError() {
return
}
data.FunctionSchema = functionSchema

response.Diagnostics.Append(response.State.Set(ctx, &data)...)
}
Expand All @@ -246,6 +313,7 @@ func (r *agentActionGroupResource) Update(ctx context.Context, request resource.
!new.ActionGroupState.Equal(old.ActionGroupState) ||
!new.APISchema.Equal(old.APISchema) ||
!new.Description.Equal(old.Description) ||
!new.FunctionSchema.Equal(old.FunctionSchema) ||
!new.ParentActionGroupSignature.Equal(old.ParentActionGroupSignature) {
input := &bedrockagent.UpdateAgentActionGroupInput{}
response.Diagnostics.Append(fwflex.Expand(ctx, new, input)...)
Expand Down Expand Up @@ -337,6 +405,7 @@ type agentActionGroupResourceModel struct {
AgentVersion types.String `tfsdk:"agent_version"`
APISchema fwtypes.ListNestedObjectValueOf[apiSchemaModel] `tfsdk:"api_schema"`
Description types.String `tfsdk:"description"`
FunctionSchema fwtypes.ListNestedObjectValueOf[functionSchemaModel] `tfsdk:"function_schema"`
ID types.String `tfsdk:"id"`
ParentActionGroupSignature fwtypes.StringEnum[awstypes.ActionGroupSignature] `tfsdk:"parent_action_group_signature"`
SkipResourceInUseCheck types.Bool `tfsdk:"skip_resource_in_use_check"`
Expand Down Expand Up @@ -366,7 +435,8 @@ func (m *agentActionGroupResourceModel) setID() {
}

type actionGroupExecutorModel struct {
Lambda fwtypes.ARN `tfsdk:"lambda"`
CustomControl fwtypes.StringEnum[awstypes.CustomControlMethod] `tfsdk:"custom_control"`
Lambda fwtypes.ARN `tfsdk:"lambda"`
}

var (
Expand All @@ -375,6 +445,10 @@ var (

func (m actionGroupExecutorModel) Expand(ctx context.Context) (result any, diags diag.Diagnostics) {
switch {
case !m.CustomControl.IsNull():
return &awstypes.ActionGroupExecutorMemberCustomControl{
Value: m.CustomControl.ValueEnum(),
}, diags
case !m.Lambda.IsNull():
return &awstypes.ActionGroupExecutorMemberLambda{
Value: m.Lambda.ValueString(),
Expand Down Expand Up @@ -414,6 +488,41 @@ func (m apiSchemaModel) Expand(ctx context.Context) (result any, diags diag.Diag
return nil, diags
}

type functionSchemaModel struct {
Functions fwtypes.ListNestedObjectValueOf[functionModel] `tfsdk:"functions"`
}

func (m functionSchemaModel) Expand(ctx context.Context) (result any, diags diag.Diagnostics) {
switch {
case !m.Functions.IsNull():
functions := m.Functions
memberFunctions := make([]awstypes.Function, len(functions.Elements()))
diags.Append(fwflex.Expand(ctx, functions, &memberFunctions)...)
if diags.HasError() {
return nil, diags
}

return &awstypes.FunctionSchemaMemberFunctions{
Value: memberFunctions,
}, diags
}

return nil, diags
}

type functionModel struct {
Description types.String `tfsdk:"description"`
Name types.String `tfsdk:"name"`
Parameters fwtypes.SetNestedObjectValueOf[parameterDetailModel] `tfsdk:"parameters"`
}

type parameterDetailModel struct {
Description types.String `tfsdk:"description"`
MapBlockKey types.String `tfsdk:"map_block_key"`
Required types.Bool `tfsdk:"required"`
Type fwtypes.StringEnum[awstypes.Type] `tfsdk:"type"`
}

type s3IdentifierModel struct {
S3BucketName types.String `tfsdk:"s3_bucket_name"`
S3ObjectKey types.String `tfsdk:"s3_object_key"`
Expand All @@ -427,6 +536,8 @@ func flattenActionGroupExecutor(ctx context.Context, apiObject awstypes.ActionGr
var actionGroupExecutorData actionGroupExecutorModel

switch v := apiObject.(type) {
case *awstypes.ActionGroupExecutorMemberCustomControl:
actionGroupExecutorData.CustomControl = fwtypes.StringEnumValue(v.Value)
case *awstypes.ActionGroupExecutorMemberLambda:
actionGroupExecutorData.Lambda = fwtypes.ARNValue(v.Value)
}
Expand Down Expand Up @@ -456,3 +567,25 @@ func flattenAPISchema(ctx context.Context, apiObject awstypes.APISchema) fwtypes

return fwtypes.NewListNestedObjectValueOfPtrMust(ctx, &apiSchemaData)
}

func flattenFunctionSchema(ctx context.Context, apiObject awstypes.FunctionSchema) (fwtypes.ListNestedObjectValueOf[functionSchemaModel], diag.Diagnostics) {
var diags diag.Diagnostics

if apiObject == nil {
return fwtypes.NewListNestedObjectValueOfNull[functionSchemaModel](ctx), diags
}

var functionSchemaData functionSchemaModel

switch v := apiObject.(type) {
case *awstypes.FunctionSchemaMemberFunctions:
functions := fwtypes.NewListNestedObjectValueOfSliceMust[functionModel](ctx, []*functionModel{})
diags.Append(fwflex.Flatten(ctx, v.Value, &functions)...)
if diags.HasError() {
return fwtypes.NewListNestedObjectValueOfNull[functionSchemaModel](ctx), diags
}
functionSchemaData.Functions = functions
}

return fwtypes.NewListNestedObjectValueOfPtrMust(ctx, &functionSchemaData), diags
}
134 changes: 131 additions & 3 deletions internal/service/bedrockagent/agent_action_group_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,8 @@ func TestAccBedrockAgentAgentActionGroup_s3APISchema(t *testing.T) {
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentActionGroupExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "action_group_name", rName),
resource.TestCheckResourceAttr(resourceName, "action_group_executor.#", acctest.Ct1),
resource.TestCheckResourceAttrSet(resourceName, "action_group_executor.0.lambda"),
),
},
{
Expand Down Expand Up @@ -120,6 +122,78 @@ func TestAccBedrockAgentAgentActionGroup_update(t *testing.T) {
})
}

func TestAccBedrockAgentAgentActionGroup_functionSchema(t *testing.T) {
ctx := acctest.Context(t)
rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
resourceName := "aws_bedrockagent_agent_action_group.test"
var v awstypes.AgentActionGroup

resource.ParallelTest(t, resource.TestCase{
PreCheck: func() { acctest.PreCheck(ctx, t); acctest.PreCheckPartitionHasService(t, names.BedrockEndpointID) },
ErrorCheck: acctest.ErrorCheck(t, names.BedrockAgentServiceID),
ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
CheckDestroy: testAccCheckAgentActionGroupDestroy(ctx),
Steps: []resource.TestStep{
{
Config: testAccAgentActionGroupConfig_functionSchema(rName),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentActionGroupExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "action_group_name", rName),
resource.TestCheckResourceAttr(resourceName, "function_schema.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.0.name", "sayHello"),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.0.description", "Says Hello"),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.0.parameters.#", acctest.Ct2),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.0.parameters.0.map_block_key", names.AttrMessage),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.0.parameters.0.type", "string"),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.0.parameters.0.description", "The Hello message"),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.0.parameters.0.required", acctest.CtTrue),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.0.parameters.1.map_block_key", "unused"),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.0.parameters.1.type", "integer"),
resource.TestCheckResourceAttr(resourceName, "function_schema.0.functions.0.parameters.1.required", acctest.CtFalse),
),
},
{
ResourceName: resourceName,
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"skip_resource_in_use_check"},
},
},
})
}

func TestAccBedrockAgentAgentActionGroup_returnControl(t *testing.T) {
ctx := acctest.Context(t)
rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
resourceName := "aws_bedrockagent_agent_action_group.test"
var v awstypes.AgentActionGroup

resource.ParallelTest(t, resource.TestCase{
PreCheck: func() { acctest.PreCheck(ctx, t); acctest.PreCheckPartitionHasService(t, names.BedrockEndpointID) },
ErrorCheck: acctest.ErrorCheck(t, names.BedrockAgentServiceID),
ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
CheckDestroy: testAccCheckAgentActionGroupDestroy(ctx),
Steps: []resource.TestStep{
{
Config: testAccAgentActionGroupConfig_returnControl(rName),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentActionGroupExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "action_group_name", rName),
resource.TestCheckResourceAttr(resourceName, "action_group_executor.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, "action_group_executor.0.custom_control", "RETURN_CONTROL"),
),
},
{
ResourceName: resourceName,
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"skip_resource_in_use_check"},
},
},
})
}

func testAccCheckAgentActionGroupDestroy(ctx context.Context) resource.TestCheckFunc {
return func(s *terraform.State) error {
conn := acctest.Provider.Meta().(*conns.AWSClient).BedrockAgentClient(ctx)
Expand Down Expand Up @@ -168,7 +242,7 @@ func testAccCheckAgentActionGroupExists(ctx context.Context, n string, v *awstyp

func testAccAgentActionGroupConfig_basic(rName string) string {
return acctest.ConfigCompose(testAccAgentConfig_basic(rName, "anthropic.claude-v2", "basic claude"),
testAccAgentActionGroupConfig_lamba(rName),
testAccAgentActionGroupConfig_lambda(rName),
fmt.Sprintf(`
resource "aws_bedrockagent_agent_action_group" "test" {
action_group_name = %[1]q
Expand All @@ -188,7 +262,7 @@ resource "aws_bedrockagent_agent_action_group" "test" {

func testAccAgentActionGroupConfig_s3APISchema(rName string) string {
return acctest.ConfigCompose(testAccAgentConfig_basic(rName, "anthropic.claude-v2", "basic claude"),
testAccAgentActionGroupConfig_lamba(rName),
testAccAgentActionGroupConfig_lambda(rName),
fmt.Sprintf(`
resource "aws_s3_bucket" "test" {
bucket = %[1]q
Expand Down Expand Up @@ -219,7 +293,7 @@ resource "aws_bedrockagent_agent_action_group" "test" {
`, rName))
}

func testAccAgentActionGroupConfig_lamba(rName string) string {
func testAccAgentActionGroupConfig_lambda(rName string) string {
return fmt.Sprintf(`
data "aws_iam_policy_document" "lambda_assume" {
statement {
Expand Down Expand Up @@ -251,3 +325,57 @@ resource "aws_lambda_function" "test_lambda" {
}
`, rName)
}

func testAccAgentActionGroupConfig_functionSchema(rName string) string {
return acctest.ConfigCompose(testAccAgentConfig_basic(rName, "anthropic.claude-v2", "basic claude"),
testAccAgentActionGroupConfig_lambda(rName),
fmt.Sprintf(`
resource "aws_bedrockagent_agent_action_group" "test" {
action_group_name = %[1]q
agent_id = aws_bedrockagent_agent.test.agent_id
agent_version = "DRAFT"
description = "Basic Agent Action"
skip_resource_in_use_check = true
action_group_executor {
lambda = aws_lambda_function.test_lambda.arn
}
function_schema {
functions {
name = "sayHello"
description = "Says Hello"
parameters {
map_block_key = "message"
type = "string"
description = "The Hello message"
required = true
}
parameters {
map_block_key = "unused"
type = "integer"
required = false
}
}
}
}
`, rName))
}

func testAccAgentActionGroupConfig_returnControl(rName string) string {
return acctest.ConfigCompose(testAccAgentConfig_basic(rName, "anthropic.claude-v2", "basic claude"),
testAccAgentActionGroupConfig_lambda(rName),
fmt.Sprintf(`
resource "aws_bedrockagent_agent_action_group" "test" {
action_group_name = %[1]q
agent_id = aws_bedrockagent_agent.test.agent_id
agent_version = "DRAFT"
description = "Basic Agent Action"
skip_resource_in_use_check = true
action_group_executor {
custom_control = "RETURN_CONTROL"
}
api_schema {
payload = file("${path.module}/test-fixtures/api_schema.yaml")
}
}
`, rName))
}
Loading
Loading