AWSTemplateFormatVersion: '2010-09-09'
Description: CloudFormation stack for creating an Amazon Bedrock agent with custom action groups, create ses identity for sending mail minimal deployment and creating a SageMaker Studio Domain which will default to JupyterLab spaces
Parameters:
  AgentName:
    Type: String
    Description: The name of the Bedrock agent
    Default: PortfolioRecommenderAgent
  EmailIdentityParameter:
    Type: String
    Description: mail id for ses identity
    AllowedPattern: "^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$"
  UserProfileName:
    Type: String
    Description: The user profile name for the SageMaker workshop
    Default: 'sagemakeruser'
  DomainName:
    Type: String
    Description: The domain name of the Sagemaker studio instance
    Default: 'bedrock-return-of-control-demo'
  SpaceName:
    Type: String
    Description: The space name of the Sagemaker studio JupyterLab instance
    Default: 'bedrock-agent-space'
  ApplicationName:
    Type: String
    Default: 'amazonbedrockagentroc'
    Description: Name of the application
  InstanceType:
    Type: String
    Default: ml.t3.medium
    Description: Instance type for JupyterLab  
  
Metadata: 
  AWS::CloudFormation::Interface: 
    ParameterGroups: 
      - 
        Label: 
          default: "Configuration"
        Parameters: 
          - EmailIdentityParameter    
Resources:
  BedrockAgentRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: 
                - bedrock.amazonaws.com
            Action: 
              - sts:AssumeRole
      Path: "/"
      Policies:
        - PolicyName: InvokeModelPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Sid: "AmazonBedrockAgentBedrockFoundationModelPolicyProd"
                Effect: Allow
                Action: "bedrock:InvokeModel"
                Resource:
                  - "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0"
  MyBedrockAgent:
    Type: AWS::Bedrock::Agent
    Properties:
      AgentName: !Ref AgentName
      AgentResourceRoleArn: !GetAtt BedrockAgentRole.Arn

      FoundationModel: "anthropic.claude-3-sonnet-20240229-v1:0"
      Instruction: "You are a personal assistant. You are friendly, polite and casual. You help with deciding detailed asset investment portfolio recommendation. User will provide their <email id>, <financial goals (retirement, wealth accumulation or passive income generation)>, <risk tolerance (low, medium or high)>, <investment horizon (short-term which is till 0 to 3 years or long term which is more than 3 years)> , <ESG preference> and <mail to be sent preference>. Based on this data, we need to create a customized portfolio for the user. 
Do not ask for any confirmation from the user to send the customized portfolio, thats what they are here for, so proceed without any further inputs other than the financial goals, risk tolerance, investment horizon, ESG preference, mail preference. 
For fetching the customized portfolio strictly check for the set of internal function you have been provided with. Once you get the customized portfolio, your job is to format that customized portfolio into an email format. 
If user has requested for mail to be sent as prefered choice, make sure to call the available function available to you to send the mail.
Think logically based on user mail to be sent preference and call the function with out fail to send the mail."
      Description: "Portfolio Recommender Agent"
      IdleSessionTTLInSeconds: 3600


      AutoPrepare: true 
      ActionGroups:
        - ActionGroupName: "Portfolio-Recommender-Action-Group"
          Description: "Portfolio Recommender Action Group"
          ActionGroupExecutor:
            CustomControl: RETURN_CONTROL
          FunctionSchema:
            Functions:
              - Name: "fetch_custom_portfolio"
                Description: "This function is used to trigger internal system to fetch custom portfolio"
                Parameters:
                  ESG_Preference:
                    Description: "ESG Preference of the customer as yes or no"
                    Type: "string"
                    Required: true
                  email_id:
                    Description: "Email id of the customer"
                    Type: "string"
                    Required: true
                  financial_goals:
                    Description: "Financial goals of customer as retirement, wealth accumulation or passive income generation"
                    Type: "string"
                    Required: true
                  investment_horizon:
                    Description: "investment horizon of the customer as short term or long term"
                    Type: "string"
                    Required: true
                  risk_tolerance:
                    Description: "Risk tolerance of the customer as low, medium or high"
                    Type: "string"
                    Required: true                    

              - Name: "send_custom_portfolio_as_email"
                Description: "This function is used to send the email content of the custom portfolio to the user mail id. If user preference is to get the mail with custom portfolio recommendation, then make sure to run this function."
                Parameters:
                  custom_portfolio_email_content_to_be_sent_to_customer:
                    Description: "custom portfolio email content to be sent to the customer"
                    Type: "string"
                    Required: true
                  email_id:
                    Description: "Email id of the customer"
                    Type: "string"
                    Required: true


  MyBedrockAgentAlias:
    Type: AWS::Bedrock::AgentAlias
    Properties:
      AgentId: !Ref MyBedrockAgent
      AgentAliasName: "SampleAlias"
      Description: "Alias for SampleAgent"

  EmailIdentity:
    Type: AWS::SES::EmailIdentity
    Properties:
      EmailIdentity: !Ref EmailIdentityParameter
  BedrockAgentParameter:
    Type: AWS::SSM::Parameter
    Properties:
      Name: !Sub '/${ApplicationName}/agentID'
      Type: String
      Value: !Ref MyBedrockAgent
      Description: 'BedrockAgent AgentID'
  BedrockAgentAliasParameter:
    Type: AWS::SSM::Parameter
    Properties:
      Name: !Sub '/${ApplicationName}/agentaliasID'
      Type: String
      Value: !Ref MyBedrockAgentAlias
      Description: 'BedrockAgent agentaliasID'
  BedrockAgentSenderSNSMailParameter:
    Type: AWS::SSM::Parameter
    Properties:
      Name: !Sub '/${ApplicationName}/senderSNSMail'
      Type: String
      Value: !Ref EmailIdentityParameter
      Description: 'EmailIdentitysender'    
# Lambda execution role has limited actions DescribeVpcs and DescribeSubnets and does not support resource level ARN's or resource level permissions. We have limited the scope by adding conditions.
  LambdaExecutionRole:
    Type: "AWS::IAM::Role"
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - "sts:AssumeRole"
      Path: /
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
      Policies: 
        - PolicyName: iam-access
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - iam:GetRole
                  - iam:GetRolePolicy
                Resource: !Sub 'arn:aws:iam::${AWS::AccountId}:role/*'
        - PolicyName: ec2-access
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - ec2:DescribeVpcs
                  - ec2:DescribeSubnets
                Resource: '*'
                Condition:
                  StringEquals:
                    'aws:RequestedRegion': !Ref 'AWS::Region'
                    'aws:PrincipalAccount': !Ref 'AWS::AccountId'

  SageMakerExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument: 
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal: 
              Service: 
                - sagemaker.amazonaws.com
            Action: 
              - sts:AssumeRole
          - Effect: Allow
            Principal: 
              Service: 
                - bedrock.amazonaws.com
            Action: 
              - sts:AssumeRole
      Path: /
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/AmazonSageMakerFullAccess"
        - "arn:aws:iam::aws:policy/AWSCloudFormationReadOnlyAccess"
        - "arn:aws:iam::aws:policy/AmazonBedrockReadOnly"
      Policies: 
        - PolicyName: iam-access
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - iam:GetRole
                  - iam:GetRolePolicy
                  - iam:PutRolePolicy
                  - iam:CreateRole
                  - iam:DeleteRole
                  - iam:CreatePolicy
                  - iam:AttachRolePolicy
                  - iam:CreateServiceLinkedRole
                  - iam:DeleteRole
                  - iam:DeletePolicy
                  - iam:PassRole
                  - iam:GetPolicy
                  - iam:DetachRolePolicy
                  - iam:DeleteRolePolicy
                  - iam:ListAttachedRolePolicies
                Resource: !Sub 'arn:aws:iam::${AWS::AccountId}:role/*'
        - PolicyName: bedrock-agents-access
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - bedrock:CreateAgent
                  - bedrock:AssociateAgentKnowledgeBase
                  - bedrock:CreateAgentActionGroup
                  - bedrock:CreateAgentAlias
                  - bedrock:DeleteAgent
                  - bedrock:DeleteAgentMemory
                  - bedrock:DeleteAgentActionGroup
                  - bedrock:DeleteAgentVersion
                  - bedrock:DeleteAgentAlias
                  - bedrock:PrepareAgent
                  - bedrock:DisassociateAgentKnowledgeBase
                  - bedrock:UpdateAgentAlias
                  - bedrock:UpdateAgentActionGroup
                  - bedrock:UpdateAgent
                  - bedrock:UpdateAgentKnowledgeBase
                  - bedrock:TagResource
                  - bedrock:UntagResource
                  - bedrock:GetAgent
                  - bedrock:GetAgentKnowledgeBase
                  - bedrock:GetAgentAlias
                  - bedrock:GetAgentVersion
                  - bedrock:GetAgentActionGroup
                  - bedrock:GetAgentMemory
                  - bedrock:InvokeInlineAgent
                  - bedrock:InvokeAgent
                  - bedrock:ListAgentActionGroups
                  - bedrock:ListAgentAliases
                  - bedrock:ListAgentKnowledgeBases
                  - bedrock:ListAgentVersions
                  - bedrock:ListAgents
                  - bedrock:AssociateAgentCollaborator
                  - bedrock:DisassociateAgentCollaborator
                  - bedrock:UpdateAgentCollaborator
                  - bedrock:GetAgentCollaborator
                  - bedrock:ListAgentCollaborators
                Resource: [!Sub 'arn:aws:bedrock:${AWS::Region}:${AWS::AccountId}:agent/*',!Sub 'arn:aws:bedrock:${AWS::Region}:${AWS::AccountId}:agent-alias/*/*']
        - PolicyName: bedrock-deny
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Deny
                Action:
                  - bedrock:CreateModelCustomizationJob
                  - bedrock:StopModelCustomizationJob
                  - bedrock:GetModelCustomizationJob
                  - bedrock:ListModelCustomizationJobs
                Resource: "*"
        - PolicyName: send-mail
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - ses:SendEmail
                Resource: !Sub 'arn:aws:ses:${AWS::Region}:${AWS::AccountId}:identity/*'
        - PolicyName: ssn_get_parameter
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - ssm:GetParameter
                Resource: !Sub 'arn:aws:ssm:${AWS::Region}:${AWS::AccountId}:parameter/${ApplicationName}/*'
# Lambda function to get the default vpc. Not connecting to any resources with in any vpc.              
  DefaultVpcLambda:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: CFGetDefaultVpcId
      Code:
        ZipFile: |
          import json
          import boto3
          import cfnresponse

          ec2 = boto3.client('ec2')

          def lambda_handler(event, context):              
              if 'RequestType' in event and event['RequestType'] == 'Create':
                  vpc_id = get_default_vpc_id()
                  subnets = get_subnets_for_vpc(vpc_id)
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, {'VpcId': vpc_id, 'Subnets': subnets}, '')
              else:
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, {}, '')

          def get_default_vpc_id():
              vpcs = ec2.describe_vpcs(Filters=[{'Name': 'is-default', 'Values': ['true']}])
              vpc_id = vpcs['Vpcs'][0]['VpcId']
              return vpc_id

          def get_subnets_for_vpc(vpcId):
              response = ec2.describe_subnets(Filters=[{'Name': 'vpc-id', 'Values': [vpcId]}])
              subnet_ids = [subnet['SubnetId'] for subnet in response['Subnets']]
              return subnet_ids 
      Description: Return default VPC ID and Subnets
      Handler: index.lambda_handler
      MemorySize: 512
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: python3.12
      Timeout: 30
      ReservedConcurrentExecutions: 5

  DefaultVpcFinder:
    Type: Custom::ResourceForFindingDefaultVpc
    Properties:
      ServiceToken: !GetAtt DefaultVpcLambda.Arn
      ServiceTimeout: 30

  StudioDomain:
    Type: AWS::SageMaker::Domain
    Properties: 
      AppNetworkAccessType: PublicInternetOnly
      AuthMode: IAM
      DefaultUserSettings: 
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn
        JupyterLabAppSettings:
          DefaultResourceSpec:
            InstanceType: !Ref InstanceType
            LifecycleConfigArn: !GetAtt JupyterLabLifecycleConfig.StudioLifecycleConfigArn
        StudioWebPortal: ENABLED
        StudioWebPortalSettings:
          HiddenAppTypes:
            - Canvas
            - JupyterServer # Studio Classic
            - RStudioServerPro
            - TensorBoard
            - DetailedProfiler
          HiddenMlTools:
            - FeatureStore
            - EmrClusters
      DomainName: !Ref DomainName
      SubnetIds: !GetAtt DefaultVpcFinder.Subnets
      VpcId: !GetAtt DefaultVpcFinder.VpcId
      DefaultSpaceSettings: 
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn
      Tags:
        - Key: 'env'
          Value: 'aws-bedrock-agents-imd'
        - Key: 'version'
          Value: '1.0'
  JupyterLabLifecycleConfig:
    Type: AWS::SageMaker::StudioLifecycleConfig
    Properties:
      StudioLifecycleConfigName: !Sub ${AWS::StackName}-jupyter-lifecycle
      StudioLifecycleConfigAppType: JupyterLab
      StudioLifecycleConfigContent: !Base64 |
        #!/bin/bash
        set -e
        mkdir -p /home/sagemaker-user/scripts

        curl https://aws-blogs-artifacts-public.s3.us-east-1.amazonaws.com/ML-16807/external-bedrock-agent-api.py >/home/sagemaker-user/scripts/external-bedrock-agent-api.py
        curl https://aws-blogs-artifacts-public.s3.us-east-1.amazonaws.com/ML-16807/streamlit-app-portfolio-recommender.py >/home/sagemaker-user/scripts/streamlit-app-portfolio-recommender.py
        curl https://aws-blogs-artifacts-public.s3.us-east-1.amazonaws.com/ML-16807/requirements.txt >/home/sagemaker-user/scripts/requirements.txt
        # Set appropriate permissions
        chmod -R 755 /home/sagemaker-user/scripts/
        
        # Create a success marker
        touch /home/sagemaker-user/scripts/.init_complete
  UserProfile:
    Type: AWS::SageMaker::UserProfile
    Properties: 
      DomainId: !GetAtt StudioDomain.DomainId
      UserProfileName: !Ref UserProfileName
      UserSettings:
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn
        JupyterLabAppSettings:
          DefaultResourceSpec:
            InstanceType: !Ref InstanceType
            LifecycleConfigArn: !GetAtt JupyterLabLifecycleConfig.StudioLifecycleConfigArn
          LifecycleConfigArns:
          - !GetAtt JupyterLabLifecycleConfig.StudioLifecycleConfigArn
  JupyterLabSpace:
    Type: AWS::SageMaker::Space
    DependsOn:
      - UserProfile
    Properties:
      DomainId: !GetAtt StudioDomain.DomainId
      SpaceName: !Ref SpaceName
      SpaceSettings:
        AppType: JupyterLab
        SpaceStorageSettings:
          EbsStorageSettings: 
             EbsVolumeSizeInGb: 5
        JupyterLabAppSettings:
          DefaultResourceSpec: 
            InstanceType: !Ref InstanceType
            LifecycleConfigArn: !GetAtt JupyterLabLifecycleConfig.StudioLifecycleConfigArn
      SpaceSharingSettings:
        SharingType: Private
      OwnershipSettings:
        OwnerUserProfileName: !Ref UserProfileName

Outputs:
  AgentId:
    Description: The ID of the created Bedrock agent
    Value: !Ref MyBedrockAgent

  AgentAliasId:
    Description: The ID of the created Bedrock agent alias
    Value: !Ref MyBedrockAgentAlias

