Skip to content

Commit

Permalink
Merge pull request #38944 from drewtul/bug-single-prompt-configuration
Browse files Browse the repository at this point in the history 
[Bug] Fix issues where only some prompts are overridden
  • Loading branch information
@jar-b
jar-b authored Aug 22, 2024
2 parents 67c89ea + e224fd6 commit 82a8eaf
Show file tree
Hide file tree
Showing 3 changed files with 166 additions and 0 deletions.
3 changes: 3 additions & 0 deletions .changelog/38944.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
```release-note:bug
resource/aws_bedrockagent_agent: Fixes consistency issues where only some prompts are overridden
```
22 changes: 22 additions & 0 deletions internal/service/bedrockagent/agent.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,8 @@ func (r *agentResource) Create(ctx context.Context, request resource.CreateReque
}
}

removeDefaultPrompts(agent)

// Set values for unknowns.
response.Diagnostics.Append(fwflex.Flatten(ctx, agent, &data)...)
if response.Diagnostics.HasError() {
Expand Down Expand Up @@ -252,6 +254,8 @@ func (r *agentResource) Read(ctx context.Context, request resource.ReadRequest,
return
}

removeDefaultPrompts(agent)

response.Diagnostics.Append(fwflex.Flatten(ctx, agent, &data)...)
if response.Diagnostics.HasError() {
return
Expand Down Expand Up @@ -329,6 +333,8 @@ func (r *agentResource) Update(ctx context.Context, request resource.UpdateReque
}
}

removeDefaultPrompts(agent)

// Set values for unknowns.
response.Diagnostics.Append(fwflex.Flatten(ctx, agent, &new)...)
if response.Diagnostics.HasError() {
Expand Down Expand Up @@ -539,6 +545,22 @@ func waitAgentDeleted(ctx context.Context, conn *bedrockagent.Client, id string,
return nil, err
}

func removeDefaultPrompts(agent *awstypes.Agent) {
pocm := agent.PromptOverrideConfiguration

overrides := pocm.PromptConfigurations

var filtered []awstypes.PromptConfiguration

for _, override := range overrides {
if override.PromptCreationMode == awstypes.CreationModeOverridden {
filtered = append(filtered, override)
}
}

pocm.PromptConfigurations = filtered
}

type agentResourceModel struct {
AgentARN types.String `tfsdk:"agent_arn"`
AgentID types.String `tfsdk:"agent_id"`
Expand Down
141 changes: 141 additions & 0 deletions internal/service/bedrockagent/agent_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,110 @@ func TestAccBedrockAgentAgent_full(t *testing.T) {
})
}

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

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: testAccCheckAgentDestroy(ctx),
Steps: []resource.TestStep{
{
Config: testAccAgentConfig_singlePrompt(rName, "anthropic.claude-v2", "basic claude"),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "agent_name", rName),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, names.AttrDescription, "basic claude"),
resource.TestCheckResourceAttr(resourceName, "skip_resource_in_use_check", acctest.CtTrue),
),
},
{
ResourceName: resourceName,
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"skip_resource_in_use_check"},
},
},
})
}

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

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: testAccCheckAgentDestroy(ctx),
Steps: []resource.TestStep{
{
Config: testAccAgentConfig_basic(rName, "anthropic.claude-v2", "basic claude"),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "agent_name", rName),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.0.prompt_configurations.#", acctest.Ct0),
resource.TestCheckResourceAttr(resourceName, names.AttrDescription, "basic claude"),
resource.TestCheckResourceAttr(resourceName, "skip_resource_in_use_check", acctest.CtFalse),
),
},
{
ResourceName: resourceName,
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"skip_resource_in_use_check"},
},
{
Config: testAccAgentConfig_singlePrompt(rName, "anthropic.claude-v2", "basic claude"),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "agent_name", rName),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.0.prompt_configurations.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, names.AttrDescription, "basic claude"),
resource.TestCheckResourceAttr(resourceName, "skip_resource_in_use_check", acctest.CtTrue),
),
},
{
Config: testAccAgentConfig_full(rName, "anthropic.claude-v2", "basic claude"),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "agent_name", rName),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.0.prompt_configurations.#", acctest.Ct4),
resource.TestCheckResourceAttr(resourceName, names.AttrDescription, "basic claude"),
resource.TestCheckResourceAttr(resourceName, "skip_resource_in_use_check", acctest.CtTrue),
),
},
{
Config: testAccAgentConfig_singlePrompt(rName, "anthropic.claude-v2", "basic claude"),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "agent_name", rName),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.0.prompt_configurations.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, names.AttrDescription, "basic claude"),
resource.TestCheckResourceAttr(resourceName, "skip_resource_in_use_check", acctest.CtTrue),
),
},
{
ResourceName: resourceName,
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"skip_resource_in_use_check"},
},
},
})
}

func TestAccBedrockAgentAgent_update(t *testing.T) {
ctx := acctest.Context(t)
rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
Expand Down Expand Up @@ -419,3 +523,40 @@ resource "aws_bedrockagent_agent" "test" {
}
`, rName, model, desc))
}

func testAccAgentConfig_singlePrompt(rName, model, desc string) string {
return acctest.ConfigCompose(testAccAgent_base(rName, model), fmt.Sprintf(`
resource "aws_bedrockagent_agent" "test" {
agent_name = %[1]q
agent_resource_role_arn = aws_iam_role.test_agent.arn
description = %[3]q
idle_session_ttl_in_seconds = 500
instruction = file("${path.module}/test-fixtures/instruction.txt")
foundation_model = %[2]q
skip_resource_in_use_check = true
prompt_override_configuration {
override_lambda = null
prompt_configurations = [
{
base_prompt_template = file("${path.module}/test-fixtures/post-processing.txt")
inference_configuration = [
{
max_length = 2048
stop_sequences = ["Human:"]
temperature = 0
top_k = 250
top_p = 1
},
]
parser_mode = "DEFAULT"
prompt_creation_mode = "OVERRIDDEN"
prompt_state = "DISABLED"
prompt_type = "POST_PROCESSING"
},
]
}
}
`, rName, model, desc))
}

0 comments on commit 82a8eaf

Please sign in to comment.