Adding guardrail support to all strands agents. Uses existing guardrail created by user (used in several places throughout the application outside of agents already). No guardrail is used if users don't specify one at deploy time. Note this required adding a helper function, so IDPAgent's can either create their own strands.BedrockModel (not recommended) or can use the create_strands_bedrock_model helper function to create one to automatically include IDP guardrails (recommended)

Author: kaleko

lib/idp_common_pkg/idp_common/agents/analytics/agent.py

@@ -11,9 +11,9 @@
 
 import boto3
 import strands
-from strands.models import BedrockModel
 
 from ..common.config import load_result_format_description
+from ..common.strands_bedrock_model import create_strands_bedrock_model
 from .config import load_python_plot_generation_examples
 from .tools import CodeInterpreterTools, get_database_info, run_athena_query
 from .utils import register_code_interpreter_tools
@@ -138,7 +138,9 @@ def run_athena_query_with_config(
     # Get model ID from environment variable
     model_id = os.environ.get("DOCUMENT_ANALYSIS_AGENT_MODEL_ID")
 
-    bedrock_model = BedrockModel(model_id=model_id, boto_session=session)
+    bedrock_model = create_strands_bedrock_model(
+        model_id=model_id, boto_session=session
+    )
 
     # Create the Strands agent with tools and system prompt
     strands_agent = strands.Agent(

lib/idp_common_pkg/idp_common/agents/common/strands_bedrock_model.py

@@ -0,0 +1,43 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+"""
+Helper function for creating BedrockModel instances with automatic guardrail support.
+"""
+
+import os
+
+from strands.models import BedrockModel
+
+
+def create_strands_bedrock_model(
+    model_id: str, boto_session=None, **kwargs
+) -> BedrockModel:
+    """
+    Create a BedrockModel with automatic guardrail configuration from environment.
+
+    Args:
+        model_id: The Bedrock model ID to use
+        boto_session: Optional boto3 session
+        **kwargs: Additional arguments to pass to BedrockModel
+
+    Returns:
+        BedrockModel instance with guardrails applied if configured
+    """
+    # Get guardrail configuration from environment if available
+    guardrail_env = os.environ.get("GUARDRAIL_ID_AND_VERSION", "")
+    if guardrail_env:
+        try:
+            guardrail_id, guardrail_version = guardrail_env.split(":")
+            if guardrail_id and guardrail_version:
+                kwargs.update(
+                    {
+                        "guardrail_id": guardrail_id,
+                        "guardrail_version": guardrail_version,
+                        "guardrail_trace": "enabled",
+                    }
+                )
+        except ValueError:
+            pass  # Invalid format, continue without guardrails
+
+    return BedrockModel(model_id=model_id, boto_session=boto_session, **kwargs)

lib/idp_common_pkg/idp_common/agents/external_mcp/agent.py

@@ -12,10 +12,10 @@
 import boto3
 from mcp.client.streamable_http import streamablehttp_client
 from strands import Agent
-from strands.models import BedrockModel
 from strands.tools.mcp import MCPClient
 
 from ..common.oauth_auth import get_cognito_bearer_token
+from ..common.strands_bedrock_model import create_strands_bedrock_model
 
 logger = logging.getLogger(__name__)
 
@@ -130,7 +130,9 @@ def create_external_mcp_agent(
                     raise Exception(error_msg)
 
             # Create Bedrock model
-            bedrock_model = BedrockModel(model_id=model_id, boto_session=session)
+            bedrock_model = create_strands_bedrock_model(
+                model_id=model_id, boto_session=session
+            )
 
             # Create system prompt
             system_prompt = f"""

lib/idp_common_pkg/idp_common/agents/orchestrator/agent.py

@@ -15,9 +15,9 @@
 
 import strands
 from strands import tool
-from strands.models import BedrockModel
 
 from ..common.config import load_result_format_description
+from ..common.strands_bedrock_model import create_strands_bedrock_model
 
 logger = logging.getLogger(__name__)
 
@@ -161,7 +161,7 @@ def tool_func(query: str) -> str:
     )
 
     # Create the orchestrator agent
-    model = BedrockModel(
+    model = create_strands_bedrock_model(
         model_id=model_id,
         session=session,
     )

lib/idp_common_pkg/idp_common/agents/sample_calculator/agent.py

@@ -7,9 +7,10 @@
 
 import boto3
 import strands
-from strands.models import BedrockModel
 from strands_tools import calculator
 
+from ..common.strands_bedrock_model import create_strands_bedrock_model
+
 logger = logging.getLogger(__name__)
 
 
@@ -31,7 +32,7 @@ def create_sample_calculator_agent(
     model_id = "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
 
     # Create Bedrock model
-    model = BedrockModel(model_id=model_id, session=session)
+    model = create_strands_bedrock_model(model_id=model_id, session=session)
 
     # Create and return agent with calculator tool
     return strands.Agent(model=model, tools=[calculator])

template.yaml

@@ -3608,6 +3608,7 @@ Resources:
                 - !Ref ReportingBucketName
           DOCUMENT_ANALYSIS_AGENT_MODEL_ID: !Ref DocumentAnalysisAgentModelId
           AWS_STACK_NAME: !Ref AWS::StackName
+          GUARDRAIL_ID_AND_VERSION: !If [HasGuardrailConfig, !Sub "${BedrockGuardrailId}:${BedrockGuardrailVersion}", ""]
       LoggingConfig:
         LogGroup: !Ref AgentProcessorLogGroup
       Policies:
@@ -3645,6 +3646,14 @@ Resources:
               Resource: 
                 - !Sub "arn:${AWS::Partition}:bedrock:*::foundation-model/*"
                 - !Sub "arn:${AWS::Partition}:bedrock:${AWS::Region}:${AWS::AccountId}:inference-profile/*"
+            - !If
+              - HasGuardrailConfig
+              - Effect: Allow
+                Action:
+                  - "bedrock:ApplyGuardrail"
+                Resource:
+                  - !Sub "arn:${AWS::Partition}:bedrock:${AWS::Region}:${AWS::AccountId}:guardrail/${BedrockGuardrailId}"
+              - !Ref AWS::NoValue
             - Effect: Allow
               Action:
                 - appsync:GraphQL