feat(bedrock): replace PrepareAgent with shouldPrepareAgent prop (#250)

Co-authored-by: Heitor Vital <[email protected]>

Author: jstrunk

apidocs/classes/bedrock.Agent.md

@@ -30,15 +30,15 @@ Deploy a Bedrock Agent.
 - [aliasId](bedrock.Agent.md#aliasid)
 - [aliasName](bedrock.Agent.md#aliasname)
 - [cdkTagManager](bedrock.Agent.md#cdktagmanager)
-- [changeIds](bedrock.Agent.md#changeids)
 - [name](bedrock.Agent.md#name)
 - [node](bedrock.Agent.md#node)
-- [prepareAgent](bedrock.Agent.md#prepareagent)
+- [resourceUpdates](bedrock.Agent.md#resourceupdates)
 - [role](bedrock.Agent.md#role)
+- [shouldPrepareAgent](bedrock.Agent.md#shouldprepareagent)
 
 ### Methods
 
-- [\_addPrepareAgentDependency](bedrock.Agent.md#_addprepareagentdependency)
+- [\_addAliasDependency](bedrock.Agent.md#_addaliasdependency)
 - [addActionGroup](bedrock.Agent.md#addactiongroup)
 - [addAlias](bedrock.Agent.md#addalias)
 - [toString](bedrock.Agent.md#tostring)
@@ -120,14 +120,6 @@ cdk.ITaggableV2.cdkTagManager
 
 ___
 
-### changeIds
-
-• `Private` **changeIds**: `string`[] = `[]`
-
-A list of values to indicate if PrepareAgent or an Alias needs to be updated.
-
-___
-
 ### name
 
 • `Readonly` **name**: `string`
@@ -148,13 +140,11 @@ Construct.node
 
 ___
 
-### prepareAgent
-
-• `Private` **prepareAgent**: `CustomResource`
+### resourceUpdates
 
-The prepareAgent custom resource.
+• `Private` **resourceUpdates**: `string`[] = `[]`
 
-Add other resources as dependencies to ensure Prepare Agent is called after they are updated.
+A list of values to indicate if PrepareAgent or an Alias needs to be updated.
 
 ___
 
@@ -164,20 +154,27 @@ ___
 
 The IAM role for the agent.
 
+___
+
+### shouldPrepareAgent
+
+• `Private` `Readonly` **shouldPrepareAgent**: `boolean`
+
+If prepare agent should be called on resource updates.
+
 ## Methods
 
-### \_addPrepareAgentDependency
+### \_addAliasDependency
 
-▸ **_addPrepareAgentDependency**(`resource`, `changeId?`): `void`
+▸ **_addAliasDependency**(`updatedAt`): `void`
 
-Register a dependency for prepareAgent.
+Register a dependency for aliases.
 
 #### Parameters
 
 | Name | Type | Description |
 | :------ | :------ | :------ |
-| `resource` | `IResource` | The resource that will be registered as a dependency. |
-| `changeId?` | `string` | The changeId of the resource that will be registered as a dependency. This is an internal core function and should not be called directly. |
+| `updatedAt` | `string` | The updatedAt of the resource that will be registered as a dependency. This is an internal core function and should not be called directly. |
 
 #### Returns
 

apidocs/interfaces/bedrock.AgentActionGroupProps.md

@@ -15,6 +15,7 @@
 - [apiSchema](bedrock.AgentActionGroupProps.md#apischema)
 - [description](bedrock.AgentActionGroupProps.md#description)
 - [parentActionGroupSignature](bedrock.AgentActionGroupProps.md#parentactiongroupsignature)
+- [shouldPrepareAgent](bedrock.AgentActionGroupProps.md#shouldprepareagent)
 
 ## Properties
 
@@ -84,3 +85,17 @@ ___
 If you specify this value as AMAZON.UserInput, the agent will prompt additional information from the user when it
 doesn't have enough information to respond to an utterance. Leave this field blank if you don't want the agent to
 prompt additional information.
+
+___
+
+### shouldPrepareAgent
+
+• `Optional` `Readonly` **shouldPrepareAgent**: `boolean`
+
+Whether to prepare the agent for use.
+
+**`Default`**
+
+```ts
+- false
+```

apidocs/interfaces/bedrock.AgentAliasProps.md

@@ -11,7 +11,7 @@
 - [agentId](bedrock.AgentAliasProps.md#agentid)
 - [agentVersion](bedrock.AgentAliasProps.md#agentversion)
 - [aliasName](bedrock.AgentAliasProps.md#aliasname)
-- [changeIds](bedrock.AgentAliasProps.md#changeids)
+- [resourceUpdates](bedrock.AgentAliasProps.md#resourceupdates)
 
 ## Properties
 
@@ -51,9 +51,8 @@ The name for the agent alias.
 
 ___
 
-### changeIds
+### resourceUpdates
 
-• `Optional` `Readonly` **changeIds**: `string`[]
+• `Optional` `Readonly` **resourceUpdates**: `string`[]
 
-The list of change ids to let CloudFormation determine when to update the alias.
-A changeId is a hash of the properties of an agent, an agent/knowledge base association, or an action group.
+The list of resource update timestamps to let CloudFormation determine when to update the alias.

apidocs/interfaces/bedrock.AgentProps.md

@@ -19,6 +19,7 @@ Properties for a Bedrock Agent.
 - [knowledgeBases](bedrock.AgentProps.md#knowledgebases)
 - [name](bedrock.AgentProps.md#name)
 - [promptOverrideConfiguration](bedrock.AgentProps.md#promptoverrideconfiguration)
+- [shouldPrepareAgent](bedrock.AgentProps.md#shouldprepareagent)
 
 ## Properties
 
@@ -133,3 +134,17 @@ Overrides for the agent.
 ```ts
 - No overrides are provided.
 ```
+
+___
+
+### shouldPrepareAgent
+
+• `Optional` `Readonly` **shouldPrepareAgent**: `boolean`
+
+Whether to prepare the agent for use.
+
+**`Default`**
+
+```ts
+- false
+```

lambda/bedrock-custom-resources/custom_resources/__init__.py

@@ -26,7 +26,6 @@
     on_event as on_event_bedrock_agent_knowledgebase,
 )
 from .bedrock_agent_action_group import on_event as on_event_bedrock_agent_action_group
-from .bedrock_prepare_agent import on_event as on_event_bedrock_prepare_agent
 
 LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
 
@@ -50,8 +49,6 @@ def on_event(event: CustomResourceRequest, context):
         return on_event_bedrock_knowledgebase(event, context)
     if resource_type == "Custom::Bedrock-DataSource":
         return on_event_bedrock_datasource(event, context)
-    if resource_type == "Custom::Bedrock-PrepareAgent":
-        return on_event_bedrock_prepare_agent(event, context)
     if resource_type == "Custom::NoOp":
         logger.info("NoOp resource type")
         # Return a response with a physical resource ID that is not empty.

lambda/bedrock-custom-resources/custom_resources/bedrock_agent.py

@@ -10,8 +10,6 @@
 # OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
 # and limitations under the License.
 #
-import pickle  # nosec B403: pickle is used to create a hashable value. The value will never be deserialized.
-
 import uuid
 
 import boto3
@@ -27,6 +25,7 @@
 
 from typing import TypedDict, NotRequired, Dict, Literal, Sequence
 
+from .bedrock_prepare_agent import prepare_agent
 from .cr_types import CustomResourceRequest, CustomResourceResponse
 from .exceptions import AWSRetryableError, can_retry
 
@@ -75,13 +74,14 @@ class AgentRequest(TypedDict):
     customerEncryptionKeyArn: NotRequired[str]
     promptOverrideConfiguration: NotRequired[PromptOverride]
     tags: NotRequired[Dict[str, str]]
+    shouldPrepareAgent: NotRequired[str]
 
 
 class AgentResponse(TypedDict):
     agentId: str
     agentArn: NotRequired[str]
     agentName: NotRequired[str]
-    changeId: NotRequired[str]
+    updatedAt: NotRequired[str]
 
 
 session = boto3.session.Session()
@@ -170,23 +170,6 @@ def validate_agent_request(request: AgentRequest) -> AgentRequest:
     return request
 
 
-def create_change_hash(request: AgentRequest) -> str:
-    return str(
-        hash(
-            # nosemgrep - Pickle is used to generate a hash. This object will never be deserialized.
-            pickle.dumps(
-                (
-                    request.get("instruction", None),
-                    request.get("foundationModel", None),
-                    request.get("description", None),
-                    request.get("idleSessionTTLInSeconds", None),
-                    request.get("promptOverrideConfiguration", None),
-                )
-            )
-        )
-    )
-
-
 def on_event(event: CustomResourceRequest[Dict], context):
     logger.debug(f"Received event: {event}")
     request_type = event["RequestType"]
@@ -213,23 +196,31 @@ def on_create(
 ) -> CustomResourceResponse:
     bedrock_agent = session.client("bedrock-agent")
 
+    if "shouldPrepareAgent" in event["ResourceProperties"]:
+        should_prepare_agent = event["ResourceProperties"]["shouldPrepareAgent"].upper() == 'TRUE'
+        del event["ResourceProperties"]["shouldPrepareAgent"]
+    else:
+        should_prepare_agent = False
+
     request = AgentRequest(**event["ResourceProperties"])
     request = validate_agent_request(request)
 
     try:
         response = bedrock_agent.create_agent(clientToken=client_token, **request)
-        return CustomResourceResponse(
+
+    except botocore.exceptions.ClientError as e:
+        can_retry(e)
+    if should_prepare_agent:
+        prepare_agent(response["agent"]["agentId"])
+    return CustomResourceResponse(
             PhysicalResourceId=response["agent"]["agentId"],
             Data=AgentResponse(
                 agentArn=response["agent"]["agentArn"],
                 agentId=response["agent"]["agentId"],
                 agentName=response["agent"]["agentName"],
-                changeId=create_change_hash(request),
+                updatedAt=str(response["agent"]["updatedAt"]),
             ),
         )
-    except botocore.exceptions.ClientError as e:
-        can_retry(e)
-
 
 @retry(
     retry=retry_if_exception_type(AWSRetryableError),
@@ -238,6 +229,11 @@ def on_create(
 )
 def on_update(event: CustomResourceRequest[Dict]) -> CustomResourceResponse:
     bedrock_agent = session.client("bedrock-agent")
+    if "shouldPrepareAgent" in event["ResourceProperties"]:
+        should_prepare_agent = event["ResourceProperties"]["shouldPrepareAgent"].upper() == 'TRUE'
+        del event["ResourceProperties"]["shouldPrepareAgent"]
+    else:
+        should_prepare_agent = False
     request = AgentRequest(**event["ResourceProperties"])
     request = validate_agent_request(request)
     if "tags" in request:
@@ -248,18 +244,19 @@ def on_update(event: CustomResourceRequest[Dict]) -> CustomResourceResponse:
             **request,
             agentId=event["PhysicalResourceId"],
         )
-        return CustomResourceResponse(
-            PhysicalResourceId=response["agent"]["agentId"],
-            Data=AgentResponse(
-                agentArn=response["agent"]["agentArn"],
-                agentId=response["agent"]["agentId"],
-                agentName=response["agent"]["agentName"],
-                changeId=create_change_hash(request),
-            ),
-        )
     except botocore.exceptions.ClientError as e:
         can_retry(e)
-
+    if should_prepare_agent:
+        prepare_agent(response["agent"]["agentId"])
+    return CustomResourceResponse(
+        PhysicalResourceId=response["agent"]["agentId"],
+        Data=AgentResponse(
+            agentArn=response["agent"]["agentArn"],
+            agentId=response["agent"]["agentId"],
+            agentName=response["agent"]["agentName"],
+            updatedAt=str(response["agent"]["updatedAt"]),
+        ),
+    )
 
 @retry(
     retry=retry_if_exception_type(AWSRetryableError),