diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md new file mode 100644 index 00000000..893dd7ea --- /dev/null +++ b/DEVELOPMENT.md @@ -0,0 +1,45 @@ + +### Project Structure within src/lex-gen-ai-demo-cdk: +``` +AWSLexKoiosBlogDemo/src/lex-gen-ai-demo-cdk +- app.py +- cdk.json +- endpoint_handler.py +- upload_file_to_s3.py +- shutdown_endpoint.py +- index-creation-docker-image/ + - index_creation_app.py + - Dockerfile + - index_creation_requirements.txt +- lex_gen_ai_demo_cdk_files/ + - __init__.py + - lex_gen_ai_demo_file_stack.py +- lex-gen-ai-demo-docker-image/ + - runtime_lambda_app.py + - Dockerfile + - runtime_lambda_requirements.txt +- requirements.txt +- source.bat +``` + +## Common Errors & Troubleshooting + +### "ValueError: Must setup local AWS configuration with a region supported by SageMaker." +Solution: You must set an aws region with `export AWS_DEFAULT_REGION=` + +### Error creating role +``` +botocore.exceptions.ClientError: An error occurred (AccessDenied) when calling the CreateRole operation: User: is not authorized to perform: iam:CreateRole on resource: because no identity-based policy allows the iam:CreateRole action +``` +Solution: you must ensure the Iam role you are using has sufficient permissions to create Iam roles + +### Error LexGenAIDemoFilesStack: fail: docker push exited with error code 1: tag does not exist +Issue: Error while building the image. Here are some common ones + +#### Error processing tar file(exit status 1): write /path/libcublas.so.11: no space left on device +Issue: Docker has run out of memory due to too many images +Solution: Delete unused images in the Docker application and then [prune docker](https://docs.docker.com/config/pruning/) in command line + +#### ConnectionResetError: [Errno 104] Connection reset by peer +Issue: Pip issue +Solution: Clear pip cache (`python3 -m pip cache purge`) and run again diff --git a/README.md b/README.md index 7f922047..25567239 100644 --- a/README.md +++ b/README.md @@ -1,17 +1,124 @@ -## My Project +# AWS Lex Conversational FAQ Demo -TODO: Fill this README out! +Demonstration of LLM integration into a lex bot using Lambda codehooks and a Sagemaker endpoint. -Be sure to: +![Diagram](diagram.png) -* Change the title in this README -* Edit your repository description on GitHub +### What resources will be created? +This CDK code will create the following: + - 1 Sagemaker endpoint hosting a model (falcon-7b-instruct on ml.g5.8xlarge by default but this is configurable) + - 1 Lex bot + - 2 S3 buckets (one for your uploaded source, one for the created index) + - 2 Lambda functions (one to ingest the source and create an image, one to be invoked as codehook during lambda and provide an FAQ answer when needed) + - 1 Event listener attached to an S3 bucket to call the index creation lambda automatically when a file is uploaded + - 2 Iam roles (one for the lex bot to call lambda, one for the lambdas to call sagemaker and S3) -## Security +## Requirements -See [CONTRIBUTING](CONTRIBUTING.md#security-issue-notifications) for more information. +### AWS setup +**Region** -## License +If you have not yet run `aws configure` and set a default region, you must do so, or you can also run `export AWS_DEFAULT_REGION=` -This library is licensed under the MIT-0 License. See the LICENSE file. +**Authorization** +You must use a role that has sufficient permissions to create Iam roles, as well as cloudformation resources + +### Python >3.7 +Make sure you have [python3](https://www.python.org/downloads/) installed at a version >=3.7.x + +### Docker +Make sure you have [Docker](https://www.docker.com/products/docker-desktop/) installed on your machine and running in the background + +### AWS CDK +Make sure you have the [AWS CDK](https://docs.aws.amazon.com/cdk/v2/guide/getting_started.html#getting_started_install) installed on your machine + + +## Setup + +### Set up virtual enviroment and gather packages + +``` +cd src/lex-gen-ai-demo-cdk-files +``` + +Install the required dependencies (aws-cdk-lib and constructs) into your Python environment +``` +pip install -r requirements.txt +``` + +### Gather and deploy resources with the CDK + +First synthesize, which executes the application, defines which resources will be created and translates this into a cloudformation template +``` +cdk synth +``` +Now bootstrap, which provisions the resources you'll use when deploying the application +``` +cdk bootstrap +``` +and deploy with +``` +cdk deploy LexGenAIDemoFilesStack +``` + +The deployment will create a lex bot and S3 buckets and will dockerize the code in the koios_cdk_files/koios-docker-image directory and push that image to ECR so it can run in Lambda. Dont worry if this step takes a long time while pushing to ECR, we are bundling up two docker images and uploading them so it will take some time. + +## Usage +Once all the resources are created after `cdk deploy` finishes running you must upload a .PDF or .txt file at least once so an index can be created. You can use our upload script `upload_file_to_s3.py path/to/your/file` or you can navigate to the S3 console and manually upload a file. On upload the ingestion lambda will read the file and create an embedding which it will upload to the other S3 bucket. Now that an embedding exists you can go to your bot and begin using it. If you want to update the embedding you can upload a new file and a new embedding will overwrite the old embedding. Once you have a new embedding you must restart the runtime lambda function for it to start using the new embedding. + +Note, the first time the embedding lambda and the runtime lambda are called the latency will be much slower as it must load resources and save them in the lambda enviroment. Once loaded these resources will stay in the enviroment as long as the ECR image is not deleted. This means your first request will be slow but after that it will speed up now that the resources are cached. + +### Uploading files +Now, you have to upload your source file so the indexing lambda can create an index for the runtime lambda function to use. You can use our script with any PDF or .txt file by running +``` +python3 upload_file_to_s3.py path/to/your/file +``` +or you can open the S3 bucket in the console and manually upload a file. On upload an index will automatically be generated. +Note: If you upload a large file, the index will be large and the S3 read time on cold start may become large. + +Once you've uploaded your file, wait a little for your index to be created and then you can go into the Lex console and test your bot (no need to build your bot unless you've made changes after creation). The first time you create an index and the first time you query the bot it will take a little longer (around 90 seconds) as we need to load models and cache them in the lambda-ECR enviroment but after they are cached there is no need to download them and latency will be much faster. These resources will remain cached as long as the ECR image is not deleted. Additionally for better cold start performance you can an instance for your runtime lambda function. There are directions to do so below. + +### Configurations + +🚨 **Remember to shut down your endpoint if you're done using it!** 🚨 + +We have provided a script to deactivate an endpoint and endpoint configuration with whatever name is in the endpoint creation script. To run: +``` +python3 shut_down_endpoint.py +``` + +#### Custom model and instance type configuration: + +The function `create_endpoint_from_HF_image()` is called in `app.py`. This function acceptst the following arguments: + - hf_model_id (required): For the purposes of the demo we have this set to [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b). You can find any https://huggingface.co/ and feed it in + - instance_type (optional, default is ml.g5.8xlarge): If you dont give an argument we'll use ml.g5.8xlarge. You can use any endpoint [sage instance type](https://aws.amazon.com/sagemaker/pricing/) + - endpoint_name (optional, default is whatever SAGEMAKER_ENDPOINT_NAME is set to in the file endpoint_handler.py): You can give your endpoint a custom name. It is recomended that you don't do this but if you do, you have to change it in the lamdba images (constant is called ENDPOINT_NAME in index_creation_app.py and runtime_lambda_app.py) + - number_of_gpu (optional, default is 1): Set this to any number of GPUs the hardware you chose allows. + + If you have in invalid configuration the endpoint will fail to create. You can see the specific error in the cloudwatch logs. If you fail creation you can run `python3 shut_down_endpoint.py` to clean up the endpoint but if you do so manually in the console **you must delete both the endpoint and the endpoint configuration** + +#### Fruther configuration +If you would like to further configure the endpoint you can change the specific code in `endpoint_handler.py` + +The LLM is hosted on a sagemaker endpoint and deployed as a sagemaker [HuggingFaceModel](https://sagemaker.readthedocs.io/en/stable/frameworks/huggingface/sagemaker.huggingface.html). We are also using a huggingface model image. You can read more about it [here](https://aws.amazon.com/blogs/machine-learning/announcing-the-launch-of-new-hugging-face-llm-inference-containers-on-amazon-sagemaker/). For further model configuration you can read about sagemaker model deployments [here](https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints-deployment.html). + +For our indexing and retrieval we are using [llama-index](https://github.com/jerryjliu/llama_index). If you would like to configure the index retriever you can do so in the `runtime_lambda_app.py` file in the `VectorIndexRetriever` object on line 70. If you want to update index creation you can update the constants defined at the top of the index creation and runtime lambdas (`index_creation_app.py` and `runtime_lambda_app.py`). Make sure to familiarize yourself with [llama-index terms](https://gpt-index.readthedocs.io/en/latest/guides/tutorials/terms_definitions_tutorial.html) and the [llama-index prompthelper](https://gpt-index.readthedocs.io/en/latest/reference/service_context/prompt_helper.html) for best results. + +### Tips for best results + +**Keep your lambda perpetually warm by provisioning an instance for the runtime lambda (lex-codehook-fn)** + +Go to Lambda console > select the function lex-codehook-fn + +Versions > Publish new version + +Under this version + - Provisioned Concurrency > set value to 1 + - Permissions > Resource based policy statements > Add Permissions > AWS Service > Other, your-policy-name, lexv2.amazonaws.com, your-lambda-arn, lamdba:InvokeFunction + +Go to your Lex Bot (LexGenAIDemoBotCfn) + +Aliases > your-alias > your-language > change lambda function version or alias > change to your-version + +This will keep an instance running at all times and keep your lambda ready so that you wont have cold start latency. This will cost a bit extra (https://aws.amazon.com/lambda/pricing/) so use thoughfully. \ No newline at end of file diff --git a/diagram.png b/diagram.png new file mode 100644 index 00000000..689c0a1f Binary files /dev/null and b/diagram.png differ diff --git a/src/lex-gen-ai-demo-cdk/app.py b/src/lex-gen-ai-demo-cdk/app.py new file mode 100644 index 00000000..10a393b3 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/app.py @@ -0,0 +1,16 @@ + +import aws_cdk as cdk + +from lex_gen_ai_demo_cdk_files.lex_gen_ai_demo_cdk_files_stack import LexGenAIDemoFilesStack +from create_web_crawler_lambda import LambdaStack +from endpoint_handler import create_endpoint_from_HF_image + +# create_endpoint_from_HF_image(hf_model_id, instance_type="ml.g5.8xlarge", endpoint_name=SAGEMAKER_ENDPOINT_NAME, number_of_gpu=1) +# You can run with no arguments to get default values of google/flan-t5-xxl on ml.g5.8xlarge, or pass in your own arguments +create_endpoint_from_HF_image(hf_model_id="tiiuae/falcon-7b-instruct") + +app = cdk.App() +filestack = LexGenAIDemoFilesStack(app, "LexGenAIDemoFilesStack") +web_crawler_lambda_stack = LambdaStack(app, 'LexGenAIDemoFilesStack-Webcrawler') + +app.synth() diff --git a/src/lex-gen-ai-demo-cdk/cdk.json b/src/lex-gen-ai-demo-cdk/cdk.json new file mode 100644 index 00000000..095c4b11 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/cdk.json @@ -0,0 +1,51 @@ +{ + "app": "python3 app.py", + "watch": { + "include": [ + "**" + ], + "exclude": [ + "README.md", + "cdk*.json", + "requirements*.txt", + "source.bat", + "**/__init__.py", + "python/__pycache__", + "tests" + ] + }, + "context": { + "@aws-cdk/aws-lambda:recognizeLayerVersion": true, + "@aws-cdk/core:checkSecretUsage": true, + "@aws-cdk/core:target-partitions": [ + "aws", + "aws-cn" + ], + "@aws-cdk-containers/ecs-service-extensions:enableDefaultLogDriver": true, + "@aws-cdk/aws-ec2:uniqueImdsv2TemplateName": true, + "@aws-cdk/aws-ecs:arnFormatIncludesClusterName": true, + "@aws-cdk/aws-iam:minimizePolicies": true, + "@aws-cdk/core:validateSnapshotRemovalPolicy": true, + "@aws-cdk/aws-codepipeline:crossAccountKeyAliasStackSafeResourceName": true, + "@aws-cdk/aws-s3:createDefaultLoggingPolicy": true, + "@aws-cdk/aws-sns-subscriptions:restrictSqsDescryption": true, + "@aws-cdk/aws-apigateway:disableCloudWatchRole": true, + "@aws-cdk/core:enablePartitionLiterals": true, + "@aws-cdk/aws-events:eventsTargetQueueSameAccount": true, + "@aws-cdk/aws-iam:standardizedServicePrincipals": true, + "@aws-cdk/aws-ecs:disableExplicitDeploymentControllerForCircuitBreaker": true, + "@aws-cdk/aws-iam:importedRoleStackSafeDefaultPolicyName": true, + "@aws-cdk/aws-s3:serverAccessLogsUseBucketPolicy": true, + "@aws-cdk/aws-route53-patters:useCertificate": true, + "@aws-cdk/customresources:installLatestAwsSdkDefault": false, + "@aws-cdk/aws-rds:databaseProxyUniqueResourceName": true, + "@aws-cdk/aws-codedeploy:removeAlarmsFromDeploymentGroup": true, + "@aws-cdk/aws-apigateway:authorizerChangeDeploymentLogicalId": true, + "@aws-cdk/aws-ec2:launchTemplateDefaultUserData": true, + "@aws-cdk/aws-secretsmanager:useAttachedSecretResourcePolicyForSecretTargetAttachments": true, + "@aws-cdk/aws-redshift:columnId": true, + "@aws-cdk/aws-stepfunctions-tasks:enableEmrServicePolicyV2": true, + "@aws-cdk/aws-ec2:restrictDefaultSecurityGroup": true, + "@aws-cdk/aws-apigateway:requestValidatorUniqueId": true + } + } \ No newline at end of file diff --git a/src/lex-gen-ai-demo-cdk/create_web_crawler_lambda.py b/src/lex-gen-ai-demo-cdk/create_web_crawler_lambda.py new file mode 100644 index 00000000..5aca2aae --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/create_web_crawler_lambda.py @@ -0,0 +1,36 @@ +from aws_cdk import ( + Duration, Stack, + aws_lambda as lambda_, + aws_s3 as s3, + aws_iam as iam +) + +from constructs import Construct + +class LambdaStack(Stack): + + def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None: + super().__init__(scope, construct_id, **kwargs) + # Iam role for lambda to invoke sagemaker + web_crawl_lambda_cfn_role = iam.Role(self, "Cfn-gen-ai-demo-web-crawler", + assumed_by=iam.ServicePrincipal("lambda.amazonaws.com") + ) + web_crawl_lambda_cfn_role.add_managed_policy(iam.ManagedPolicy.from_aws_managed_policy_name("AmazonS3FullAccess")) + web_crawl_lambda_cfn_role.add_to_policy( + iam.PolicyStatement( + actions=[ + "logs:CreateLogGroup", + "logs:CreateLogStream", + "logs:PutLogEvents" + ], + resources=["*"] + ) + ) + # Lambda function + lambda_function= lambda_.DockerImageFunction(self, "web-crawler-docker-image-CFN", + function_name="WebCrawlerLambda", + code=lambda_.DockerImageCode.from_image_asset("web-crawler-docker-image"), + role=web_crawl_lambda_cfn_role, + memory_size=1024, + timeout=Duration.minutes(5) + ) diff --git a/src/lex-gen-ai-demo-cdk/endpoint_handler.py b/src/lex-gen-ai-demo-cdk/endpoint_handler.py new file mode 100644 index 00000000..9fee1ab1 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/endpoint_handler.py @@ -0,0 +1,99 @@ +import json +import boto3 +import time +from sagemaker.huggingface import get_huggingface_llm_image_uri +from sagemaker.huggingface import HuggingFaceModel + +# get image from huggingface +llm_image = get_huggingface_llm_image_uri( + "huggingface", + version="0.8.2" +) + +assume_role_policy_document = json.dumps({ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Principal": { + "Service": [ + "sagemaker.amazonaws.com", + "ecs.amazonaws.com" + ] + }, + "Action": "sts:AssumeRole" + } + ] +}) + +# editable to whatever you want your endpoint and role to be. You can use an existing role or a new one +# IMPORTANT: make sure your lambda endpoint name in lambda_app.py is consisitent if you change it here +SAGEMAKER_IAM_ROLE_NAME = 'Sagemaker-Endpoint-Creation-Role' +SAGEMAKER_ENDPOINT_NAME = "huggingface-pytorch-sagemaker-endpoint" + +# Create role and give sagemaker permissions +def get_iam_role(role_name=SAGEMAKER_IAM_ROLE_NAME): + iam_client = boto3.client('iam') + + try: + role = iam_client.get_role(RoleName=role_name) + role_arn = role['Role']['Arn'] + print(f"Role {role_arn} found!") + return role_arn + + except: + role_arn = iam_client.create_role( + RoleName=SAGEMAKER_IAM_ROLE_NAME, + AssumeRolePolicyDocument=assume_role_policy_document + )['Role']['Arn'] + + time.sleep(10) # give the policy some time to properly create + + response = iam_client.attach_role_policy( + PolicyArn='arn:aws:iam::aws:policy/AmazonSageMakerFullAccess', + RoleName=SAGEMAKER_IAM_ROLE_NAME, + ) + print(f"Creating {role_arn}") + time.sleep(20) # give iam time to let the role create + return role_arn + + +# Define Model and Endpoint configuration parameter + +health_check_timeout = 300 +trust_remote_code = True + +# Create sagemaker endpoint, default values are flan t5 xxl in a g5.8xl instance +def create_endpoint_from_HF_image(hf_model_id, instance_type="ml.g5.8xlarge", endpoint_name=SAGEMAKER_ENDPOINT_NAME, number_of_gpu=1): + sagemaker_client = boto3.client('sagemaker') + + try: # check if endpoint already existst + sagemaker_client.describe_endpoint(EndpointName=SAGEMAKER_ENDPOINT_NAME) + print(f"Endpoint with name {SAGEMAKER_ENDPOINT_NAME} found!") + return + + except: + print(f"Creating endpoint with model{hf_model_id} on {instance_type}...") + + # create HuggingFaceModel with the image uri + llm_model = HuggingFaceModel( + role=get_iam_role(), + image_uri=llm_image, + env={ + 'HF_MODEL_ID': hf_model_id, + 'SM_NUM_GPUS': json.dumps(number_of_gpu), + 'HF_MODEL_TRUST_REMOTE_CODE': json.dumps(trust_remote_code) + } + ) + + # Deploy model to an endpoint + # https://sagemaker.readthedocs.io/en/stable/api/inference/model.html#sagemaker.model.Model.deploy + llm = llm_model.deploy( + endpoint_name=endpoint_name, + initial_instance_count=1, + instance_type=instance_type, + # volume_size=400, # If using an instance with local SSD storage, volume_size must be None, e.g. p4 but not p3 + container_startup_health_check_timeout=health_check_timeout # 10 minutes to be able to load the model + ) + + print(f"\nEndpoint created ({endpoint_name})") diff --git a/src/lex-gen-ai-demo-cdk/index-creation-docker-image/Dockerfile b/src/lex-gen-ai-demo-cdk/index-creation-docker-image/Dockerfile new file mode 100644 index 00000000..6fadfb43 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/index-creation-docker-image/Dockerfile @@ -0,0 +1,14 @@ +FROM public.ecr.aws/lambda/python:3.8 + +COPY index_creation_requirements.txt . +RUN pip3 install -r index_creation_requirements.txt --target "${LAMBDA_TASK_ROOT}" + +# Copy function code +COPY *.py ${LAMBDA_TASK_ROOT} + +# Set the CMD to your handler (could also be done as a parameter override outside of the Dockerfile) +CMD [ "index_creation_app.handler" ] + +# Set cache to a location lambda can write to +ENV TRANSFORMERS_CACHE="/tmp/TRANSFORMERS_CACHE" + diff --git a/src/lex-gen-ai-demo-cdk/index-creation-docker-image/index_creation_app.py b/src/lex-gen-ai-demo-cdk/index-creation-docker-image/index_creation_app.py new file mode 100644 index 00000000..03716455 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/index-creation-docker-image/index_creation_app.py @@ -0,0 +1,126 @@ +import boto3 +import json +from pathlib import Path + +import logging +from langchain.llms.base import LLM +from typing import Optional, List, Mapping, Any +import os +from llama_index import ( + LangchainEmbedding, + GPTVectorStoreIndex, + LLMPredictor, + ServiceContext, + Document, + PromptHelper, + download_loader +) + +from langchain.embeddings import HuggingFaceEmbeddings + +import logging +from botocore.exceptions import ClientError + +logger = logging.getLogger() +logger.setLevel(logging.INFO) + +ACCOUNT_ID = boto3.client('sts').get_caller_identity().get('Account') +INDEX_BUCKET = "lexgenaistack-created-index-bucket-"+ACCOUNT_ID +S3_BUCKET = "lexgenaistack-source-materials-bucket-"+ACCOUNT_ID +ENDPOINT_NAME = "huggingface-pytorch-sagemaker-endpoint" +DELIMITER = "\n\n\n" +LOCAL_INDEX_LOC = "/tmp/index_files" + +def handler(event, context): + event_record = event['Records'][0] + if event_record['eventName'] == "ObjectCreated:Put": + if ".txt" in event_record['s3']['object']['key'].lower() or ".pdf" in event_record['s3']['object']['key'].lower(): + source_material_key = event_record['s3']['object']['key'] + logger.info(f"Source file {source_material_key} found") + else: + logger.error("INVALID FILE, MUST END IN .TXT or .PDF") + return + else: + logger.error("NON OBJECTCREATION INVOCATION") + return + + s3_client = boto3.client('s3') + try: + s3_client.download_file(S3_BUCKET, source_material_key, "/tmp/"+source_material_key) + logger.info(f"Downloaded {source_material_key}") + except ClientError as e: + logger.error(e) + return "ERROR READING FILE" + + if ".pdf" in source_material_key.lower(): + PDFReader = download_loader("PDFReader", custom_path="/tmp/llama_cache") + loader = PDFReader() + documents = loader.load_data(file=Path("/tmp/"+source_material_key)) + else: + with open("/tmp/"+source_material_key) as f: + text_list = f.read().split(DELIMITER) + logger.info(f"Reading text with delimiter {repr(DELIMITER)}") + documents = [Document(t) for t in text_list] + + # define prompt helper + max_input_size = 400 # set maximum input size + num_output = 50 # set number of output tokens + max_chunk_overlap = 0 # set maximum chunk overlap + prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap) + + # define our LLM + llm_predictor = LLMPredictor(llm=CustomLLM()) + embed_model = LangchainEmbedding(HuggingFaceEmbeddings(cache_folder="/tmp/HF_CACHE")) + service_context = ServiceContext.from_defaults( + llm_predictor=llm_predictor, prompt_helper=prompt_helper, embed_model=embed_model, + ) + + index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context) + index.storage_context.persist(persist_dir=LOCAL_INDEX_LOC) + + for file in os.listdir(LOCAL_INDEX_LOC): + s3_client.upload_file(LOCAL_INDEX_LOC+"/"+file, INDEX_BUCKET, file) # ASSUMES IT CAN OVERWRITE, I.E. S3 OBJECT LOCK MUST BE OFF + + logger.info("Index successfully created") + return + +def call_sagemaker(prompt, endpoint_name=ENDPOINT_NAME): + payload = { + "inputs": prompt, + "parameters": { + "do_sample": False, + # "top_p": 0.9, + "temperature": 0.1, + "max_new_tokens": 200, + "repetition_penalty": 1.03, + "stop": ["\nUser:", "<|endoftext|>", ""] + } + } + + sagemaker_client = boto3.client("sagemaker-runtime") + payload = json.dumps(payload) + response = sagemaker_client.invoke_endpoint( + EndpointName=endpoint_name, ContentType="application/json", Body=payload + ) + response_string = response["Body"].read().decode() + return response_string + +def get_response_sagemaker_inference(prompt, endpoint_name=ENDPOINT_NAME): + resp = call_sagemaker(prompt, endpoint_name) + resp = json.loads(resp)[0]["generated_text"][len(prompt):] + return resp + +class CustomLLM(LLM): + model_name = "tiiuae/falcon-7b-instruct" + + def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: + response = get_response_sagemaker_inference(prompt, ENDPOINT_NAME) + return response + + @property + def _identifying_params(self) -> Mapping[str, Any]: + return {"name_of_model": self.model_name} + + @property + def _llm_type(self) -> str: + return "custom" \ No newline at end of file diff --git a/src/lex-gen-ai-demo-cdk/index-creation-docker-image/index_creation_requirements.txt b/src/lex-gen-ai-demo-cdk/index-creation-docker-image/index_creation_requirements.txt new file mode 100644 index 00000000..91b261b8 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/index-creation-docker-image/index_creation_requirements.txt @@ -0,0 +1,6 @@ +transformers==4.25.1 +langchain +llama-index==0.6.20 +sentence-transformers +pypdf +typing_extensions \ No newline at end of file diff --git a/src/lex-gen-ai-demo-cdk/lex-gen-ai-demo-docker-image/Dockerfile b/src/lex-gen-ai-demo-cdk/lex-gen-ai-demo-docker-image/Dockerfile new file mode 100644 index 00000000..25be1f60 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/lex-gen-ai-demo-docker-image/Dockerfile @@ -0,0 +1,14 @@ +FROM public.ecr.aws/lambda/python:3.8 + +COPY runtime_lambda_requirements.txt . +RUN pip3 install -r runtime_lambda_requirements.txt --target "${LAMBDA_TASK_ROOT}" + +# Copy function code +COPY *.py ${LAMBDA_TASK_ROOT} + +# Set the CMD to your handler (could also be done as a parameter override outside of the Dockerfile) +CMD [ "runtime_lambda_app.handler" ] + +# Set cache to a location lambda can write to +ENV TRANSFORMERS_CACHE="/tmp/TRANSFORMERS_CACHE" + diff --git a/src/lex-gen-ai-demo-cdk/lex-gen-ai-demo-docker-image/runtime_lambda_app.py b/src/lex-gen-ai-demo-cdk/lex-gen-ai-demo-docker-image/runtime_lambda_app.py new file mode 100644 index 00000000..467c3d81 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/lex-gen-ai-demo-docker-image/runtime_lambda_app.py @@ -0,0 +1,176 @@ +import boto3 +from botocore.exceptions import ClientError +import logging +import json +import os +from typing import Optional, List, Mapping, Any +from langchain.llms.base import LLM +from llama_index import ( + LangchainEmbedding, + PromptHelper, + ResponseSynthesizer, + LLMPredictor, + ServiceContext, + Prompt, +) + +from langchain.embeddings import HuggingFaceEmbeddings +from llama_index.query_engine import RetrieverQueryEngine +from llama_index.retrievers import VectorIndexRetriever +from llama_index.vector_stores.types import VectorStoreQueryMode +from llama_index import StorageContext, load_index_from_storage + +s3_client = boto3.client('s3') + +logger = logging.getLogger() +logger.setLevel(logging.INFO) + +ENDPOINT_NAME = "huggingface-pytorch-sagemaker-endpoint" +OUT_OF_DOMAIN_RESPONSE = "I'm sorry, but I am only able to give responses regarding the source topic" +INDEX_WRITE_LOCATION = "/tmp/index" +ACCOUNT_ID = boto3.client('sts').get_caller_identity().get('Account') +INDEX_BUCKET = "lexgenaistack-created-index-bucket-"+ACCOUNT_ID +RETRIEVAL_THRESHOLD = 0.4 + +# define prompt helper +max_input_size = 400 # set maximum input size +num_output = 50 # set number of output tokens +max_chunk_overlap = 0 # set maximum chunk overlap +prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap) + + +def handler(event, context): + + # lamda can only write to /tmp/ + initialize_cache() + + # define our LLM + llm_predictor = LLMPredictor(llm=CustomLLM()) + embed_model = LangchainEmbedding(HuggingFaceEmbeddings(cache_folder="/tmp/HF_CACHE")) + service_context = ServiceContext.from_defaults( + llm_predictor=llm_predictor, prompt_helper=prompt_helper, embed_model=embed_model, + ) + + ### Download index here + if not os.path.exists(INDEX_WRITE_LOCATION): + os.mkdir(INDEX_WRITE_LOCATION) + try: + s3_client.download_file(INDEX_BUCKET, "docstore.json", INDEX_WRITE_LOCATION + "/docstore.json") + s3_client.download_file(INDEX_BUCKET, "index_store.json", INDEX_WRITE_LOCATION + "/index_store.json") + s3_client.download_file(INDEX_BUCKET, "vector_store.json", INDEX_WRITE_LOCATION + "/vector_store.json") + + # load index + storage_context = StorageContext.from_defaults(persist_dir=INDEX_WRITE_LOCATION) + index = load_index_from_storage(storage_context, service_context=service_context) + logger.info("Index successfully loaded") + except ClientError as e: + logger.error(e) + return "ERROR LOADING/READING INDEX" + + retriever = VectorIndexRetriever( + service_context=service_context, + index=index, + similarity_top_k=5, + vector_store_query_mode=VectorStoreQueryMode.DEFAULT, # doesn't work with simple + alpha=0.5, + ) + + # configure response synthesizer + synth = ResponseSynthesizer.from_args( + response_mode="simple_summarize", + service_context=service_context + ) + + query_engine = RetrieverQueryEngine(retriever=retriever, response_synthesizer=synth) + query_input = event["inputTranscript"] + + try: + answer = query_engine.query(query_input) + if answer.source_nodes[0].score < RETRIEVAL_THRESHOLD: + answer = OUT_OF_DOMAIN_RESPONSE + except: + answer = OUT_OF_DOMAIN_RESPONSE + + response = generate_lex_response(event, {}, "Fulfilled", answer) + jsonified_resp = json.loads(json.dumps(response, default=str)) + return jsonified_resp + +def generate_lex_response(intent_request, session_attributes, fulfillment_state, message): + intent_request['sessionState']['intent']['state'] = fulfillment_state + return { + 'sessionState': { + 'sessionAttributes': session_attributes, + 'dialogAction': { + 'type': 'Close' + }, + 'intent': intent_request['sessionState']['intent'] + }, + 'messages': [ + { + "contentType": "PlainText", + "content": message + } + ], + 'requestAttributes': intent_request['requestAttributes'] if 'requestAttributes' in intent_request else None + } + +# define prompt template +template = ( + "We have provided context information below. \n" + "---------------------\n" + "CONTEXT1:\n" + "{context_str}\n\n" + "CONTEXT2:\n" + "CANNOTANSWER" + "\n---------------------\n" + 'Given this context, please answer the question if answerable based on on the CONTEXT1 and CONTEXT2: "{query_str}"\n; ' # otherwise specify it as CANNOTANSWER +) +my_qa_template = Prompt(template) + +def call_sagemaker(prompt, endpoint_name=ENDPOINT_NAME): + payload = { + "inputs": prompt, + "parameters": { + "do_sample": False, + # "top_p": 0.9, + "temperature": 0.1, + "max_new_tokens": 200, + "repetition_penalty": 1.03, + "stop": ["\nUser:", "<|endoftext|>", ""] + } + } + + sagemaker_client = boto3.client("sagemaker-runtime") + payload = json.dumps(payload) + response = sagemaker_client.invoke_endpoint( + EndpointName=endpoint_name, ContentType="application/json", Body=payload + ) + response_string = response["Body"].read().decode() + return response_string + +def get_response_sagemaker_inference(prompt, endpoint_name=ENDPOINT_NAME): + resp = call_sagemaker(prompt, endpoint_name) + resp = json.loads(resp)[0]["generated_text"][len(prompt):] + return resp + +class CustomLLM(LLM): + model_name = "tiiuae/falcon-7b-instruct" + + def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: + response = get_response_sagemaker_inference(prompt, ENDPOINT_NAME) + return response + + @property + def _identifying_params(self) -> Mapping[str, Any]: + return {"name_of_model": self.model_name} + + @property + def _llm_type(self) -> str: + return "custom" + +def initialize_cache(): + if not os.path.exists("/tmp/TRANSFORMERS_CACHE"): + os.mkdir("/tmp/TRANSFORMERS_CACHE") + + if not os.path.exists("/tmp/HF_CACHE"): + os.mkdir("/tmp/HF_CACHE") \ No newline at end of file diff --git a/src/lex-gen-ai-demo-cdk/lex-gen-ai-demo-docker-image/runtime_lambda_requirements.txt b/src/lex-gen-ai-demo-cdk/lex-gen-ai-demo-docker-image/runtime_lambda_requirements.txt new file mode 100644 index 00000000..3105a486 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/lex-gen-ai-demo-docker-image/runtime_lambda_requirements.txt @@ -0,0 +1,4 @@ +transformers==4.30.0 +langchain +llama-index==0.6.20 +sentence-transformers \ No newline at end of file diff --git a/src/lex-gen-ai-demo-cdk/lex_gen_ai_demo_cdk_files/__init__.py b/src/lex-gen-ai-demo-cdk/lex_gen_ai_demo_cdk_files/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/lex-gen-ai-demo-cdk/lex_gen_ai_demo_cdk_files/lex_gen_ai_demo_cdk_files_stack.py b/src/lex-gen-ai-demo-cdk/lex_gen_ai_demo_cdk_files/lex_gen_ai_demo_cdk_files_stack.py new file mode 100644 index 00000000..05746f9e --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/lex_gen_ai_demo_cdk_files/lex_gen_ai_demo_cdk_files_stack.py @@ -0,0 +1,130 @@ +from aws_cdk import ( + Duration, App, Stack, CfnResource, + aws_lex as lex, + aws_s3 as s3, + aws_s3_notifications as s3n, + aws_s3_deployment as s3deploy, + aws_iam as iam, + aws_lambda as lambda_ +) + +from constructs import Construct + +class LexGenAIDemoFilesStack(Stack): + + def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None: + super().__init__(scope, construct_id, **kwargs) + + # Iam role for bot to invoke lambda + lex_cfn_role = iam.Role(self, "CfnLexGenAIDemoRole", + assumed_by=iam.ServicePrincipal("lexv2.amazonaws.com") + ) + lex_cfn_role.add_managed_policy(iam.ManagedPolicy.from_aws_managed_policy_name("AWSLambdaExecute")) + + # Iam role for lambda to invoke sagemaker + lambda_cfn_role = iam.Role(self, "CfnLambdaGenAIDemoRole", + assumed_by=iam.ServicePrincipal("lambda.amazonaws.com") + ) + lambda_cfn_role.add_managed_policy(iam.ManagedPolicy.from_aws_managed_policy_name("AmazonSageMakerFullAccess")) + lambda_cfn_role.add_managed_policy(iam.ManagedPolicy.from_aws_managed_policy_name("AmazonS3FullAccess")) + + # will append account id to this string to avoid in region collisions + source_bucket_name = "lexgenaistack-source-materials-bucket-" + index_bucket_name = "lexgenaistack-created-index-bucket-" + + # S3 Buckets for materials to index and for the resulting indexes + source_bucket = s3.Bucket(self, "SourceMatBucketID-CFN", + bucket_name=source_bucket_name+lex_cfn_role.principal_account, + block_public_access=s3.BlockPublicAccess.BLOCK_ALL, + encryption=s3.BucketEncryption.S3_MANAGED, + enforce_ssl=True, + versioned=True) + + index_bucket = s3.Bucket(self, "IndexBucket-CFN", + bucket_name=index_bucket_name+lex_cfn_role.principal_account, + block_public_access=s3.BlockPublicAccess.BLOCK_ALL, + encryption=s3.BucketEncryption.S3_MANAGED, + enforce_ssl=True, + versioned=True) + + # create lambda image for on demand index creation + read_source_and_build_index_function = lambda_.DockerImageFunction(self, "read-source-and-build-index-function-CFN", function_name="read-source-and-build-index-fn", + code=lambda_.DockerImageCode.from_image_asset("index-creation-docker-image"), + role=lambda_cfn_role, + memory_size=10240, + timeout=Duration.minutes(5) + ) + source_bucket.add_event_notification(s3.EventType.OBJECT_CREATED, s3n.LambdaDestination(read_source_and_build_index_function)) + + # create image of lex-gen-ai-demo-docker-image, push to ECR and into a lambda function + runtime_function = lambda_.DockerImageFunction(self, "CFN-runtime-fn", function_name="lex-codehook-fn", + code=lambda_.DockerImageCode.from_image_asset("lex-gen-ai-demo-docker-image"), + role=lambda_cfn_role, + memory_size=10240, + timeout=Duration.minutes(5) + ) + runtime_function.grant_invoke(iam.ServicePrincipal("lexv2.amazonaws.com")) + + ### BOT SETUP + + # alias settings, where we define the lambda function with the ECR container with our LLM dialog code (defined in the lex-gen-ai-demo-docker-image directory) + # test bot alias for demo, create a dedicated alias for serving traffic + bot_alias_settings = lex.CfnBot.TestBotAliasSettingsProperty( + bot_alias_locale_settings=[lex.CfnBot.BotAliasLocaleSettingsItemProperty( + bot_alias_locale_setting=lex.CfnBot.BotAliasLocaleSettingsProperty( + enabled=True, + code_hook_specification=lex.CfnBot.CodeHookSpecificationProperty( + lambda_code_hook=lex.CfnBot.LambdaCodeHookProperty( + code_hook_interface_version="1.0", + lambda_arn=runtime_function.function_arn + ) + ) + ), + locale_id="en_US" + )]) + + # lambda itself is tied to alias but codehook settings are intent specific + initial_response_codehook_settings = lex.CfnBot.InitialResponseSettingProperty( + code_hook=lex.CfnBot.DialogCodeHookInvocationSettingProperty( + enable_code_hook_invocation=True, + is_active=True, + post_code_hook_specification=lex.CfnBot.PostDialogCodeHookInvocationSpecificationProperty() + ) + ) + + # placeholder intent to be missed for this demo + placeholder_intent = lex.CfnBot.IntentProperty( + name="placeHolderIntent", + initial_response_setting=initial_response_codehook_settings, + sample_utterances=[lex.CfnBot.SampleUtteranceProperty( + utterance="utterance" + )] + ) + + fallback_intent = lex.CfnBot.IntentProperty( + name="FallbackIntent", + parent_intent_signature="AMAZON.FallbackIntent", + initial_response_setting=initial_response_codehook_settings, + fulfillment_code_hook=lex.CfnBot.FulfillmentCodeHookSettingProperty( + enabled=True, + is_active=True, + post_fulfillment_status_specification=lex.CfnBot.PostFulfillmentStatusSpecificationProperty() + ) + ) + + # Create actual Lex Bot + cfn_bot = lex.CfnBot(self, "LexGenAIDemoCfnBot", + data_privacy={"ChildDirected":"false"}, + idle_session_ttl_in_seconds=300, + name="LexGenAIDemoBotCfn", + role_arn=lex_cfn_role.role_arn, + bot_locales=[lex.CfnBot.BotLocaleProperty( + locale_id="en_US", + nlu_confidence_threshold=0.4, + intents=[placeholder_intent, fallback_intent]) + ], + test_bot_alias_settings = bot_alias_settings, + auto_build_bot_locales=True + ) + + diff --git a/src/lex-gen-ai-demo-cdk/requirements.txt b/src/lex-gen-ai-demo-cdk/requirements.txt new file mode 100644 index 00000000..b9360e07 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/requirements.txt @@ -0,0 +1,4 @@ +aws-cdk-lib==2.80.0 +constructs>=10.0.0,<11.0.0 +sagemaker==2.163.0 +boto3 \ No newline at end of file diff --git a/src/lex-gen-ai-demo-cdk/shut_down_endpoint.py b/src/lex-gen-ai-demo-cdk/shut_down_endpoint.py new file mode 100644 index 00000000..844755f3 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/shut_down_endpoint.py @@ -0,0 +1,26 @@ +import boto3 +from botocore.exceptions import ClientError + +from endpoint_handler import SAGEMAKER_ENDPOINT_NAME + + +sagemaker_client = boto3.client('sagemaker') + + +try: + # verify endpoint exists + endpoint = sagemaker_client.describe_endpoint(EndpointName=SAGEMAKER_ENDPOINT_NAME) + print(f"Endpoint {endpoint['EndpointName']} found, shutting down") + + try: # delete both endpoint and configuration + sagemaker_client.delete_endpoint( + EndpointName=SAGEMAKER_ENDPOINT_NAME + ) + sagemaker_client.delete_endpoint_config( + EndpointConfigName=SAGEMAKER_ENDPOINT_NAME + ) + print(f"Endpoint {SAGEMAKER_ENDPOINT_NAME} shut down") + except ClientError as e: + print(e) +except: + print(f"Endpoint {SAGEMAKER_ENDPOINT_NAME} does not exist in account {boto3.client('sts').get_caller_identity().get('Account')}") \ No newline at end of file diff --git a/src/lex-gen-ai-demo-cdk/source.bat b/src/lex-gen-ai-demo-cdk/source.bat new file mode 100644 index 00000000..5cfcddd0 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/source.bat @@ -0,0 +1,13 @@ +@echo off + +rem The sole purpose of this script is to make the command +rem +rem source .venv/bin/activate +rem +rem (which activates a Python virtualenv on Linux or Mac OS X) work on Windows. +rem On Windows, this command just runs this batch file (the argument is ignored). +rem +rem Now we don't need to document a Windows command for activating a virtualenv. + +echo Executing .venv\Scripts\activate.bat for you +.venv\Scripts\activate.bat \ No newline at end of file diff --git a/src/lex-gen-ai-demo-cdk/upload_file_to_s3.py b/src/lex-gen-ai-demo-cdk/upload_file_to_s3.py new file mode 100644 index 00000000..7020b5a8 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/upload_file_to_s3.py @@ -0,0 +1,37 @@ +import sys +import boto3 +from botocore.exceptions import ClientError +import logging + +ACCOUNT_ID = boto3.client('sts').get_caller_identity().get('Account') +S3_BUCKET = "lexgenaistack-source-materials-bucket-"+ACCOUNT_ID +s3_client = boto3.client("s3") + +def main(): + if len(sys.argv) == 1: + print(f"[ERROR] You must specify file to upload") + elif len(sys.argv) == 2: + filepath = sys.argv[1] + upload(filepath) + elif len(sys.argv) == 3: + filepath = sys.argv[2] + upload(filepath) + else: + print("[ERROR] Too many arguments, only include /path/to/your/file") + + +def upload(filepath): + if filepath[-4:].lower() == '.txt' or filepath[-4:].lower() == '.pdf': + print(f"Uploading file at {filepath}") + try: + upload_name = filepath.split("/")[-1].replace(" ","").replace("/","") + s3_client.upload_file(filepath, S3_BUCKET, upload_name) + print(f"Successfully uploaded file at {filepath}, creating index...") + except ClientError as e: + logging.error(e) + else: + print("[ERROR] File must be txt or PDF") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/lex-gen-ai-demo-cdk/web-crawler-docker-image/Dockerfile b/src/lex-gen-ai-demo-cdk/web-crawler-docker-image/Dockerfile new file mode 100644 index 00000000..c7e2b358 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/web-crawler-docker-image/Dockerfile @@ -0,0 +1,11 @@ +FROM public.ecr.aws/lambda/python:3.8 + +COPY web_crawler_requirements.txt . +RUN pip3 install -r web_crawler_requirements.txt --target "${LAMBDA_TASK_ROOT}" + +# Copy function code +COPY *.py ${LAMBDA_TASK_ROOT} + +# Set the CMD to your handler (could also be done as a parameter override outside of the Dockerfile) +CMD [ "web_crawler_app.handler" ] + diff --git a/src/lex-gen-ai-demo-cdk/web-crawler-docker-image/web_crawler_app.py b/src/lex-gen-ai-demo-cdk/web-crawler-docker-image/web_crawler_app.py new file mode 100644 index 00000000..25af7799 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/web-crawler-docker-image/web_crawler_app.py @@ -0,0 +1,125 @@ +import boto3 +import requests +import html2text +from typing import List +import re +import logging +import json +import traceback + +logger = logging.getLogger() +logger.setLevel(logging.INFO) + + +def find_http_urls_in_parentheses(s: str, prefix: str = None): + pattern = r'\((https?://[^)]+)\)' + urls = re.findall(pattern, s) + + matched = [] + if prefix is not None: + for url in urls: + if str(url).startswith(prefix): + matched.append(url) + else: + matched = urls + + return list(set(matched)) # remove duplicates by converting to set, then convert back to list + + + +class EZWebLoader: + + def __init__(self, default_header: str = None): + self._html_to_text_parser = html2text + if default_header is None: + self._default_header = {"User-agent":"Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/47.0.2526.80 Safari/537.36"} + else: + self._default_header = default_header + + def load_data(self, + urls: List[str], + num_levels: int = 0, + level_prefix: str = None, + headers: str = None) -> List[str]: + + logging.info(f"Number of urls: {len(urls)}.") + + if headers is None: + headers = self._default_header + + documents = [] + visited = {} + for url in urls: + q = [url] + depth = num_levels + for page in q: + if page not in visited: #prevent cycles by checking to see if we already crawled a link + logging.info(f"Crawling {page}") + visited[page] = True #add entry to visited to prevent re-crawling pages + response = requests.get(page, headers=headers).text + response = self._html_to_text_parser.html2text(response) #reduce html to text + documents.append(response) + if depth > 0: + #crawl linked pages + ingest_urls = find_http_urls_in_parentheses(response, level_prefix) + logging.info(f"Found {len(ingest_urls)} pages to crawl.") + q.extend(ingest_urls) + depth -= 1 #reduce the depth counter so we go only num_levels deep in our crawl + else: + logging.info(f"Skipping {page} as it has already been crawled") + logging.info(f"Number of documents: {len(documents)}.") + return documents + +ACCOUNT_ID = boto3.client('sts').get_caller_identity().get('Account') +S3_BUCKET = "lexgenaistack-source-materials-bucket-" + ACCOUNT_ID +FILE_NAME = 'web-crawl-results.txt' + + +def handler(event, context): + url = "http://www.zappos.com/general-questions" + depth = 1 + level_prefix = "https://www.zappos.com/" + + if event is not None: + if "url" in event: + url = event["url"] + if "depth" in event: + depth = int(event["depth"]) + if "level_prefix" in event: + level_prefix = event["level_prefix"] + + # crawl the website + try: + logger.info(f"Crawling {url} to depth of {depth}...") + loader = EZWebLoader() + documents = loader.load_data([url], depth, level_prefix) + doc_string = json.dumps(documents, indent=1) + logger.info(f"Crawling {url} to depth of {depth} succeeded") + except Exception as e: + # If there's an error, print the error message + logging.error(f"An error occurred during the crawl of {url}.") + exception_traceback = traceback.format_exc() + logger.error(exception_traceback) + return { + "status": 500, + "message": exception_traceback + } + # save the results for indexing + try: + # Use the S3 client to write the string to S3 + s3 = boto3.client('s3') + s3.put_object(Body=doc_string, Bucket=S3_BUCKET, Key=FILE_NAME) + success_msg = f'Successfully put {FILE_NAME} to {S3_BUCKET}' + logging.info(success_msg) + return { + "status": 200, + "message": success_msg + } + except Exception as e: + # If there's an error, print the error message + exception_traceback = traceback.format_exc() + logger.error(exception_traceback) + return { + "status": 500, + "message": exception_traceback + } diff --git a/src/lex-gen-ai-demo-cdk/web-crawler-docker-image/web_crawler_requirements.txt b/src/lex-gen-ai-demo-cdk/web-crawler-docker-image/web_crawler_requirements.txt new file mode 100644 index 00000000..388c7bea --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/web-crawler-docker-image/web_crawler_requirements.txt @@ -0,0 +1,4 @@ +requests +html2text +accelerate +boto3 diff --git a/src/lex-gen-ai-demo-cdk/web_crawl.py b/src/lex-gen-ai-demo-cdk/web_crawl.py new file mode 100644 index 00000000..c2d596c5 --- /dev/null +++ b/src/lex-gen-ai-demo-cdk/web_crawl.py @@ -0,0 +1,43 @@ +import boto3 +import argparse +import json + + +def invoke_lambda(url=None, depth="1", level_prefix=None): + client = boto3.client('lambda') + + # Prepare the payload + payload = {} + if url is not None: + payload["url"] = url + if depth is not None: + payload["depth"] = depth + if level_prefix is not None: + payload["level_prefix"] = level_prefix + + try: + response = client.invoke( + FunctionName='WebCrawlerLambda', + InvocationType='RequestResponse', + LogType='Tail', + # The payload must be a JSON-formatted string + Payload=json.dumps(payload) + ) + + # The response from Lambda will be a JSON string, so you need to parse it + result = response['Payload'].read().decode('utf-8') + + print("Response: " + result) + + except Exception as e: + print(e) + + +# Parse command-line arguments +parser = argparse.ArgumentParser() +parser.add_argument('--url', type=str, help='The URL to process.', required=False, default=None) +parser.add_argument('--depth', type=int, help='The depth of the crawl.', required=False, default="1") +parser.add_argument('--level_prefix', type=str, help='The prefix that any links must contain to crawl.', required=False, default=None) +args = parser.parse_args() + +invoke_lambda(args.url, args.depth, args.level_prefix)