AWS Cloudformation Bedrock

Parameters:
  FoundationModel:
    Type: String
    Default: anthropic.claude-3-sonnet-20240229-v1:0
    AllowedValues:
      - anthropic.claude-3-sonnet-20240229-v1:0
      - anthropic.claude-3-haiku-20240307-v1:0
      - amazon.titan-text-express-v1
    Description: Foundation model for agent

Resources:
  AgentRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: bedrock.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: BedrockPermissions
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - bedrock:InvokeModel
                Resource: !Sub "arn:aws:bedrock:${AWS::Region}:${AWS::AccountId}:foundation-model/${FoundationModel}"

  BedrockAgent:
    Type: AWS::Bedrock::Agent
    Properties:
      AgentName: !Sub "${AWS::StackName}-agent"
      AgentResourceRoleArn: !GetAtt AgentRole.Arn
      FoundationModelArn: !Sub "arn:aws:bedrock:${AWS::Region}::foundation-model/${FoundationModel}"
      AutoPrepare: true
      Instruction: |
        You are a helpful assistant. Use the knowledge base to answer questions.

  KnowledgeBaseRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: bedrock.amazonaws.com
            Action: sts:AssumeRole

  KnowledgeBase:
    Type: AWS::Bedrock::KnowledgeBase
    Properties:
      Name: !Sub "${AWS::StackName}-kb"
      RoleArn: !GetAtt KnowledgeBaseRole.Arn
      KnowledgeBaseConfiguration:
        Type: VECTOR
        VectorKnowledgeBaseConfiguration:
          EmbeddingModelArn: !Sub "arn:aws:bedrock:${AWS::Region}::embedding-model/amazon.titan-embed-text-v1"

  DataBucket:
    Type: AWS::S3::Bucket

  S3DataSource:
    Type: AWS::Bedrock::DataSource
    Properties:
      KnowledgeBaseId: !Ref KnowledgeBase
      Name: s3-data-source
      Type: S3
      DataSourceConfiguration:
        S3Configuration:
          BucketArn: !GetAtt DataBucket.Arn
          InclusionPrefixes:
            - documents/

  Guardrail:
    Type: AWS::Bedrock::Guardrail
    Properties:
      Name: !Sub "${AWS::StackName}-guardrail"
      BlockedInputMessaging: "I cannot help with that request."
      ContentPolicyConfig:
        filtersConfig:
          - type: PROFANITY
          - type: MISCONDUCT

  ActionLambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Runtime: python3.12
      Handler: index.handler
      Role: !GetAtt ActionLambdaRole.Arn
      Code:
        ZipFile: |
          def handler(event, context):
              return {"statusCode": 200, "body": "{\"result\": \"success\"}"}

  ActionGroup:
    Type: AWS::Bedrock::AgentActionGroup
    Properties:
      ActionGroupName: api-operations
      ActionGroupState: ENABLED
      AgentId: !GetAtt BedrockAgent.AgentId
      ActionGroupExecutor:
        Lambda: !Ref ActionLambdaFunction
      FunctionSchema:
        functionConfigurations:
          - function: |
              { "name": "get_inventory", "description": "Get current inventory status", "parameters": { "type": "object", "properties": { "sku": { "type": "string" } }, "required": [] } }

aws cloudformation validate-template --template-body file://bedrock-template.yaml

# Check agent status
aws bedrock-agent get-agent --agent-id $(aws cloudformation describe-stacks --stack-name STACK_NAME --query 'Stacks[0].Outputs[?OutputKey==`AgentId`].OutputValue' --output text)

# Check knowledge base sync status
aws bedrock-agent list-knowledge-bases --agent-id AGENT_ID

# Test guardrail
aws bedrock-runtime apply_guardrail --guardrail-identifier GUARDRAIL_ID --source SOURCE

AWSTemplateFormatVersion: "2010-09-09"
Description: "Bedrock RAG Agent with Knowledge Base"

Parameters:
  FoundationModel:
    Type: String
    Default: anthropic.claude-3-sonnet-20240229-v1:0

Resources:
  # IAM Role for Agent
  AgentRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub "${AWS::StackName}-agent-role"
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: bedrock.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: InvokeModel
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: bedrock:InvokeModel
                Resource: "*"

  # IAM Role for Knowledge Base
  KnowledgeBaseRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub "${AWS::StackName}-kb-role"
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: bedrock.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: S3Access
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: s3:GetObject
                Resource: !Sub "${DataBucket.Arn}/*"

  # S3 Bucket for Documents
  DataBucket:
    Type: AWS::S3::Bucket

  # Knowledge Base
  KnowledgeBase:
    Type: AWS::Bedrock::KnowledgeBase
    Properties:
      Name: !Sub "${AWS::StackName}-kb"
      RoleArn: !GetAtt KnowledgeBaseRole.Arn
      KnowledgeBaseConfiguration:
        Type: VECTOR
        VectorKnowledgeBaseConfiguration:
          EmbeddingModelArn: !Sub "arn:aws:bedrock:${AWS::Region}::embedding-model/amazon.titan-embed-text-v1"

  # Data Source
  DataSource:
    Type: AWS::Bedrock::DataSource
    Properties:
      KnowledgeBaseId: !Ref KnowledgeBase
      Name: !Sub "${AWS::StackName}-ds"
      Type: S3
      DataSourceConfiguration:
        S3Configuration:
          BucketArn: !GetAtt DataBucket.Arn

  # Bedrock Agent
  BedrockAgent:
    Type: AWS::Bedrock::Agent
    Properties:
      AgentName: !Sub "${AWS::StackName}-agent"
      AgentResourceRoleArn: !GetAtt AgentRole.Arn
      FoundationModelArn: !Sub "arn:aws:bedrock:${AWS::Region}::foundation-model/${FoundationModel}"
      AutoPrepare: true
      Instruction: |
        You are a helpful assistant. Use the knowledge base to answer user questions accurately.

Outputs:
  AgentId:
    Description: Bedrock Agent ID
    Value: !GetAtt BedrockAgent.AgentId
  KnowledgeBaseId:
    Description: Knowledge Base ID
    Value: !Ref KnowledgeBase

Resources:
  Guardrail:
    Type: AWS::Bedrock::Guardrail
    Properties:
      Name: !Sub "${AWS::StackName}-guardrail"
      blockedInputMessaging: "Content blocked by safety filters."
      blockedOutputMessaging: "Response filtered for safety."
      contentPolicyConfig:
        filtersConfig:
          - type: PROFANITY
            inputStrength: HIGH
            outputStrength: HIGH
          - type: MISCONDUCT
            inputStrength: HIGH
            outputStrength: HIGH
      sensitiveInformationPolicyConfig:
        piiEntitiesConfig:
          - type: EMAIL
            action: ANONYMIZE
          - type: SSN
            action: BLOCK