Some security improvements and fixes

Author: arya23065

README.md

@@ -2,7 +2,7 @@
 
 In this repository, we present a solution that harnesses the power of Generative AI to streamline the user onboarding process for financial services through a digital assistant. Onboarding new customers in the banking is a crucial step in the customer journey, involving a series of activities designed to fulfill Know-Your-Customer (KYC) requirements, conduct necessary verifications, and introduce them to the bank's products or services. Traditionally, user onboarding has been a tedious and heavily manual process. Our solution provides practical guidance on addressing this challenge by leveraging a Generative AI assistant on AWS. 
 
-Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs) from leading AI companies like AI21 Labs, Anthropic, Cohere, Meta, Mistral AI, Stability AI, and Amazon via a single API, along with a broad set of capabilities you need to build generative AI applications with security, privacy, and responsible AI. Using Anthropic Claude 3 Haiku on Amazon Bedrock, we build a digital assistant that automates paperwork, identity verifications, and engages customers through conversational interactions, called Amazon Penny. As a result, customers can be onboarded in a matter of minutes through secure, automated workflows. 
+Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs) from leading AI companies like AI21 Labs, Anthropic, Cohere, Meta, Mistral AI, Stability AI, and Amazon via a single API, along with a broad set of capabilities you need to build generative AI applications with security, privacy, and responsible AI. Using Anthropic Claude 3.5 Sonnet on Amazon Bedrock, we build a digital assistant that automates paperwork, identity verifications, and engages customers through conversational interactions, called Amazon Penny. As a result, customers can be onboarded in a matter of minutes through secure, automated workflows. 
 
 ## Architecture
 
@@ -13,9 +13,9 @@ The flow of the application is as follows:
 
 1. Users access the front-end website hosted within the AWS Amplify. 
 
-2. The website invokes an endpoint to interact with the digital assistant, Penny, which is containerized and deployed in AWS Fargate. 
+2. The website invokes an Amazon CloudFront endpoint to interact with the digital assistant, Penny, which is containerized and deployed in AWS Fargate. 
 
-3. The digital assistant uses a custom Langchain Agent to answer questions on the bank's products and services and orchestrate the onboarding flow. The Large Language Model (LLM) used by the agent is Anthropic Claude 3 Haiku, provided by Amazon Bedrock. 
+3. The digital assistant uses a custom Langchain Agent to answer questions on the bank's products and services and orchestrate the onboarding flow. The Large Language Model (LLM) used by the agent is Anthropic Claude 3.5 Sonnet, provided by Amazon Bedrock. 
 
 4. If the user asks a general question related to the bank's products or services, the agent will utilize a custom Langchain Tool called the Product Search Tool. This tool uses Amazon Kendra linked with an S3 data source which contains the bank's data. 
 
@@ -41,7 +41,7 @@ Identify the AWS Account where you would like to deploy this solution and ensure
 * Install [NodeJS](https://nodejs.org/en) and [Docker](https://www.docker.com/). 
 
 In your chosen AWS Account, complete the following steps:
-* Ensure that the [Anthropic Claude 3 Haiku Model is enabled in Amazon Bedrock by adding model access](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html#model-access-add).
+* Ensure that the [Anthropic Claude 3.5 Sonnet Model is enabled in Amazon Bedrock by adding model access](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html#model-access-add).
 * Ensure that docker is running. This can be done by running the command `sudo docker info`. If Docker is running, information about Docker is displayed.
 
 > Note: Under a set of assumptions made on a monthly basis, running this workload would have an estimated hourly cost of around $1.34. Make sure to check the pricing details for each individual service to understand the costs you may be charged for different usage tiers and resource configurations

api/.gitignore

@@ -0,0 +1,3 @@
+llm/app/__pycache__/
+llm/app/penny/__pycache__/
+llm/venv/

api/lambdas/verify-face.py

@@ -26,13 +26,6 @@ def compare_faces(id_file_name, selfie_file_name):
     for faceMatch in response['FaceMatches']:
         position = faceMatch['Face']['BoundingBox']
         similarity = str(faceMatch['Similarity'])
-        print('The face at ' +
-              str(position['Left']) + ' ' +
-              str(position['Top']) +
-              ' matches with ' + similarity + '% confidence')
-
-
-    print(response['FaceMatches'])
 
     return len(response['FaceMatches'])
 
@@ -47,13 +40,12 @@ def main(event, context):
         "Access-Control-Allow-Methods": "POST"
     }
 
-    print(event)
     id_file_name = event["id_file_name"]
     selfie_file_name = event["selfie_file_name"]
+    body = ''
 
     try:
         face_matches = compare_faces(id_file_name, selfie_file_name)
-        print("Compared faces:" + str(face_matches))
 
         if face_matches > 0: 
             body = 'Face match verified'

api/lambdas/verify-id.py

@@ -29,13 +29,10 @@ def extract_text(id_file_name):
                 if "ValueDetection" in str(key):
                     curr_val = str(val['Text'])
             id_field_values[curr_type] = curr_val
-            print(id_field_values)
 
-    print(id_field_values)
     return id_field_values
 
 def compare_fields(id_field_values, required_field_values):
-    
     for key, val in required_field_values.items():
         if not ((key in id_field_values) and (val.lower() == id_field_values[key].lower())):
             return "The details you provided for " + key + " do not match your ID"
@@ -52,16 +49,13 @@ def main(event, context):
         "Access-Control-Allow-Origin": "*",
         "Access-Control-Allow-Methods": "POST"
     }
+    body= ''
 
-    print(event)
-
     id_file_name = event["file_name"]
     required_field_values = event["required_field_values"]
 
     try:
         id_file_values = extract_text(id_file_name)
-        print("Extracted info:" + str(id_file_values))
-        
         body = compare_fields(id_file_values, required_field_values)
 
     except Exception as e:

api/llm/Dockerfile

@@ -1,4 +1,4 @@
-FROM public.ecr.aws/docker/library/python:3.9
+FROM public.ecr.aws/docker/library/python:3.12
 
 WORKDIR /code
 

api/llm/app/main.py

@@ -35,10 +35,10 @@
     use_tools=True
 )
 llm = BedrockChat(
-        model_id='anthropic.claude-3-haiku-20240307-v1:0', 
+        model_id='anthropic.claude-3-5-sonnet-20240620-v1:0', 
         client=bedrock,
         model_kwargs={
-            "temperature": 0,
+            "temperature": 0.5,
             "top_k": 250,
             "top_p": 0.999,
             "stop_sequences": ["\\n\\nHuman:"]
@@ -61,7 +61,7 @@ async def question(request: Request) -> Response:
     response = agent.step()
     return JSONResponse(content={"message": response})
 
-@app.post("/uploadId")
+@app.post("/uploadDoc")
 async def id(file: UploadFile = File(...)) -> Response:
     try:
         contents = file.file.read()

api/llm/app/penny/tools.py

@@ -3,27 +3,22 @@
 import boto3
 import json
 from langchain.agents import Tool
-from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
 from langchain_community.chat_models import BedrockChat
 from langchain_community.retrievers import AmazonKendraRetriever
 from pydantic import EmailStr, Field
 from email_validator import validate_email, EmailNotValidError
 import requests
 import os
 
-#Bedrock client initialization
 bedrock = boto3.client(service_name='bedrock-runtime')
-#vector db init
-modeled = SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2")
-
 API_ENDPOINT = os.environ["apiEndpoint"]
 
 def setup_knowledge_base():
     """
     We assume that the product knowledge base is simply a text file.
     """
     llm = BedrockChat(
-        model_id='anthropic.claude-3-haiku-20240307-v1:0', 
+        model_id='anthropic.claude-3-5-sonnet-20240620-v1:0', 
         client=bedrock,
         model_kwargs={
             "temperature": 1,
@@ -133,7 +128,7 @@ def _run(self, input=""):
         }
         r = requests.post(url=url, json=body, timeout=300)
 
-        if r.json()["statusCode"] != 200:
+        if r.status_code != 200:
             return "Respond that our onboarding service is currently unavailable and to try again later."
 
         print(r.json())
@@ -142,18 +137,6 @@ def _run(self, input=""):
     def _arun(self, query: str):
         raise NotImplementedError("This tool does not support async")
 
-
-class ask_question(BaseTool):
-    name = "AskUser"
-    description = "Use this tool to ask something to the user. " \
-                    "It takes the question you want to ask as the input" \
-                  "It will return a question that you can ask"
-
-    def _run(self, question):
-        return "ask user " + question
-
-    def _arun(self, query: str):
-        raise NotImplementedError("This tool does not support async")
 
 class finish_onboarding(BaseTool):
     name = "SaveData"
@@ -193,14 +176,8 @@ def get_tools():
     # we only use one tool for now, but this is highly extensible!
     knowledge_base = setup_knowledge_base()
     tools = [
-        Tool(
-            name="ProductSearch",
-            func=knowledge_base.run,
-            description="useful for when you need to answer any question",
-        ),
-        #product_search
+        product_search(),
         email_validator(),
-        ask_question(),
         document_verification(),
         selfie_verification(),
         finish_onboarding()

api/llm/requirements.txt

@@ -1,14 +1,14 @@
 fastapi==0.109.1
 uvicorn==0.24.0.post1
-pydantic==2.5.2
+pydantic>=2.7.4,<3.0.0
 langchain==0.2.16
 langchain-community==0.2.5
+langchain-core>=0.2.40,<0.3.0
 torch==2.4.1
 torchvision==0.19.1
 sentence-transformers==2.2.2
-faiss-cpu==1.7.4
-email-validator==2.1.0
-requests==2.32.0
+email-validator==2.1.1
+requests==2.31.0
 python-multipart==0.0.7
 boto3==1.28.65
 botocore==1.31.65

infra/lib/infra-stack.ts

@@ -84,7 +84,7 @@ export class PennyInfraStack extends cdk.Stack {
 
   createLambda_verifyId = () => {
     this.verifyIdLambda = new lambda.Function(this, 'VerifyId', {
-      runtime: lambda.Runtime.PYTHON_3_8,
+      runtime: lambda.Runtime.PYTHON_3_11,
       code: lambda.Code.fromAsset('../api/lambdas'),
       handler: 'verify-id.main',
       timeout: Duration.seconds(10),
@@ -103,7 +103,7 @@ export class PennyInfraStack extends cdk.Stack {
 
   createLambda_verifyFace = () => {
     this.verifyFaceLambda = new lambda.Function(this, 'VerifyFace', {
-      runtime: lambda.Runtime.PYTHON_3_8,
+      runtime: lambda.Runtime.PYTHON_3_11,
       code: lambda.Code.fromAsset('../api/lambdas'),
       handler: 'verify-face.main',
       timeout: Duration.seconds(10),
@@ -115,15 +115,28 @@ export class PennyInfraStack extends cdk.Stack {
     this.verifyFaceLambda.addToRolePolicy(new iam.PolicyStatement(
       {
         effect: iam.Effect.ALLOW, 
-        actions: ['rekognition:CompareFaces', 's3:Get*', 's3:List*', 's3:Describe*'],
-        resources: ['*']
+        actions: ['s3:Get*', 's3:List*', 's3:Describe*'],
+        resources: [
+          `arn:aws:s3:::${this.idBucket.bucketName}`,
+          `arn:aws:s3:::${this.idBucket.bucketName}/*`
+        ]
+      }),
+    )
+
+    this.verifyFaceLambda.addToRolePolicy(new iam.PolicyStatement(
+      {
+        effect: iam.Effect.ALLOW, 
+        actions: ['rekognition:CompareFaces'],
+        resources: [
+          `*`,
+        ]
       }),
     )
   }
 
   createLambda_getAccount = () => {
     this.getAccountLambda = new lambda.Function(this, 'GetAccount', {
-      runtime: lambda.Runtime.PYTHON_3_8,
+      runtime: lambda.Runtime.PYTHON_3_11,
       code: lambda.Code.fromAsset('../api/lambdas'),
       handler: 'get-account.main',
       environment: {
@@ -142,7 +155,7 @@ export class PennyInfraStack extends cdk.Stack {
 
   createLambda_createAccount = () => {
     this.createAccountLambda = new lambda.Function(this, 'CreateAccount', {
-      runtime: lambda.Runtime.PYTHON_3_8,
+      runtime: lambda.Runtime.PYTHON_3_11,
       code: lambda.Code.fromAsset('../api/lambdas'),
       handler: 'create-account.main',
       environment: {
@@ -162,8 +175,10 @@ export class PennyInfraStack extends cdk.Stack {
     this.createAccountLambda.addToRolePolicy(new iam.PolicyStatement(
       {
         effect: iam.Effect.ALLOW, 
-        actions: ['ses:SendEmail'],
-        resources: ['*']
+        actions: ['ses:SendEmail', 'ses:SendRawEmail'],
+        resources: [
+          `arn:aws:ses:${this.region}:${this.account}:identity/*`,
+        ],
       }),
     )
   }
@@ -309,11 +324,40 @@ export class PennyInfraStack extends cdk.Stack {
       assumedBy: new iam.ServicePrincipal('kendra.amazonaws.com'),
     })
 
+    role.addToPolicy(new iam.PolicyStatement({
+      effect: iam.Effect.ALLOW,
+      actions: [
+        'cloudwatch:PutMetricData',
+      ],
+      resources: ['*'],
+      conditions: {
+        'StringEquals': {
+          'cloudwatch:namespace': 'AWS/Kendra'
+        }
+      }
+    }));
+
+    role.addToPolicy(new iam.PolicyStatement({
+      effect: iam.Effect.ALLOW,
+      actions: [
+        'logs:DescribeLogGroups',
+        'logs:CreateLogGroup',
+        'logs:PutLogEvents'
+      ],
+      resources: [
+        `arn:aws:logs:${this.region}:${this.account}:log-group:/aws/kendra/*`,
+        `arn:aws:logs:${this.region}:${this.account}:log-group:/aws/kendra/*:log-stream:*`
+      ]
+    }));
+
     role.addToPolicy(new iam.PolicyStatement(
       {
         effect: iam.Effect.ALLOW, 
-        actions: ['cloudwatch:PutMetricData', 'logs:DescribeLogGroups', 'logs:CreateLogGroup', 'logs:DescribeLogStreams', 'logs:CreateLogStream', 'logs:PutLogEvents', 's3:*'],
-        resources: ['*']
+        actions: ['s3:Get*', 's3:List*', 's3:Describe*'],
+        resources: [
+          `arn:aws:s3:::${this.catalogBucket.bucketName}`,
+          `arn:aws:s3:::${this.catalogBucket.bucketName}/*`
+        ]
       }),
     )
 
@@ -341,15 +385,15 @@ export class PennyInfraStack extends cdk.Stack {
   createLLM = () => {
     const vpc = new ec2.Vpc(this, 'LLMVpc', {
       subnetConfiguration: [
-        { subnetType: ec2.SubnetType.PUBLIC, cidrMask: 24, name: "public-" }
+        { subnetType: ec2.SubnetType.PUBLIC, cidrMask: 24, name: "public-" },
+        { subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS, cidrMask: 24, name: "private-" }
       ]
     })
 
     const sg = new ec2.SecurityGroup(this, 'LLMSG', {
       vpc: vpc,
       allowAllOutbound: true
     })
-    sg.addIngressRule(ec2.Peer.anyIpv4(), ec2.Port.tcp(80));
 
     const cluster = new ecs.Cluster(this, "LLMCluster", {
       vpc: vpc
@@ -365,8 +409,20 @@ export class PennyInfraStack extends cdk.Stack {
 
     taskDefinition.addToTaskRolePolicy(new iam.PolicyStatement({
       effect: iam.Effect.ALLOW,
-      resources: ['*'],
-      actions: ['bedrock:InvokeModel', 's3:PutObject', 'kendra:Retrieve']
+      resources: ['*'], 
+      actions: ['bedrock:InvokeModel']
+    }));
+
+    taskDefinition.addToTaskRolePolicy(new iam.PolicyStatement({
+      effect: iam.Effect.ALLOW,
+      resources: [`${this.idBucket.bucketArn}/*`],
+      actions: ['s3:PutObject']
+    }));
+
+    taskDefinition.addToTaskRolePolicy(new iam.PolicyStatement({
+      effect: iam.Effect.ALLOW,
+      resources: [this.kendraIndex.attrArn],
+      actions: ['kendra:Retrieve']
     }));
 
     const container = taskDefinition.addContainer("LLMContainer", {
@@ -394,9 +450,9 @@ export class PennyInfraStack extends cdk.Stack {
     this.llmService = new ecs.FargateService(this, 'LLMFargateService', { 
       cluster: cluster,
       taskDefinition: taskDefinition,
-      assignPublicIp: true,
+      assignPublicIp: false,
       securityGroups: [sg],
-      serviceName: 'LLMFargateService'
+      serviceName: 'LLMFargateService',
     })    
 
     const lb = new elbv2.ApplicationLoadBalancer(this, 'PennyALB', {