diff --git a/lambda/receiver/lambda_function.py b/lambda/receiver/lambda_function.py index 9d62cf2..3bbfa17 100644 --- a/lambda/receiver/lambda_function.py +++ b/lambda/receiver/lambda_function.py @@ -44,6 +44,57 @@ def generate_output_json_path(src_folder,dest_folder, s3_object): if txt.lower().endswith(".pdf"): return dest_folder+txt[0:-4]+".json" +# Function to get Textract document analysis +def get_textract_document_detection(job_id, textract_client): + # Initialize an empty list to store blocks + all_blocks = [] + next_token = None + response = {} + flag = True + logger.info("Get document detection") + try: + while True: + if flag: + # Call the Textract API to get document analysis + response = textract_client.get_document_text_detection( + JobId=job_id + ) + flag = False + else: + # Call the Textract API to get document analysis + response = textract_client.get_document_text_detection( + JobId=job_id, + NextToken=next_token + ) + + job_status = response["JobStatus"] + logger.info("Job %s status is %s.", job_id, job_status) + + # Merge the blocks from the current response + all_blocks.extend(response.get('Blocks', [])) + + # Check if there are more blocks to retrieve + next_token = response.get('NextToken') + if not next_token: + logger.info("No more Textract response to retrieve") + break + + except ClientError: + logger.exception("Couldn't get data for job %s.", job_id) + raise + else: + # Remove unnecessary keys from the last response + last_response = response.copy() + last_response.pop('Blocks', None) + last_response.pop('ResponseMetadata', None) + logger.info("Removed 'ResponseMetadata' key") + + # Merge with {'Blocks': all_blocks} + final_response = {'Blocks': all_blocks} + final_response.update(last_response) + logger.info("Final Textract response is contructed") + return final_response + # Function to get Textract document analysis def get_textract_document_analysis(job_id, textract_client): # Initialize an empty list to store blocks @@ -51,7 +102,7 @@ def get_textract_document_analysis(job_id, textract_client): next_token = None response = {} flag = True - + logger.info("Get document analysis") try: while True: if flag: @@ -101,6 +152,7 @@ def upload_response_to_s3(response, bucket_name, object_key, s3_client): response_json = json.dumps(response) try: + tags = "env=dev" # Upload the JSON response to S3 s3_client.put_object( Bucket=bucket_name, @@ -154,11 +206,13 @@ def lambda_handler(event, context): OUTPUT_LOCATION = config_dict['FOLDER_LOCATIONS']['OUTPUT_LOCATION'].format(batch_id) PROCESSED_LOCATION = config_dict['FOLDER_LOCATIONS']['PROCESSED_LOCATION'].format(batch_id) UNPROCESSED_LOCATION = config_dict['FOLDER_LOCATIONS']['UNPROCESSED_LOCATION'].format(batch_id) + PROCESS_TYPE = str(config_dict['OTHERS']['PROCESS_TYPE']).upper() logger.info('STAGGING_LOCATION: ' + STAGING_LOCATION) logger.info('OUTPUT_LOCATION: ' + OUTPUT_LOCATION) logger.info('PROCESSED_LOCATION: ' + PROCESSED_LOCATION) logger.info('UNPROCESSED_LOCATION: ' + UNPROCESSED_LOCATION) + logger.info('PROCESS_TYPE: ' + str(PROCESS_TYPE)) # Process each message from the SQS event for record in event['Records']: @@ -176,12 +230,20 @@ def lambda_handler(event, context): # Check if the status is "SUCCEEDED" if message_body.get('Status') == 'SUCCEEDED': - # Call the function to get document analysis using Textract - document_analysis = get_textract_document_analysis(job_id, textract_client) + document = {} + if PROCESS_TYPE == "ANALYSIS": + + # Call the function to get document analysis using Textract + document = get_textract_document_analysis(job_id, textract_client) + + elif PROCESS_TYPE == "DETECTION": + + # Call the function to get document analysis using Textract + document = get_textract_document_detection(job_id, textract_client) # Save the document analysis response to S3 s3_object_key = generate_output_json_path(STAGING_LOCATION,OUTPUT_LOCATION, s3_object_name) - upload_response_to_s3(document_analysis, S3_BUCKET_NAME, s3_object_key, s3_client) + upload_response_to_s3(document, S3_BUCKET_NAME, s3_object_key, s3_client) # Construct the destination paths destination_path = PROCESSED_LOCATION + s3_object_name.replace(STAGING_LOCATION,"") diff --git a/lambda/sender/lambda_function.py b/lambda/sender/lambda_function.py index 5db2eb8..30351dc 100644 --- a/lambda/sender/lambda_function.py +++ b/lambda/sender/lambda_function.py @@ -4,6 +4,8 @@ from configparser import ConfigParser import logging import os from botocore.exceptions import ClientError +import json +from urllib.parse import unquote_plus # Initialize logger logger = logging.getLogger(__name__) @@ -40,8 +42,9 @@ def load_config_from_s3(bucket_name, file_key): # Function to move a file from source to destination in S3 def move_file_within_s3(source_bucket, source_key, destination_key): try: + tags = "env=dev" # Copy the file to the destination folder - s3_client.copy_object(Bucket=source_bucket, CopySource={'Bucket': source_bucket, 'Key': source_key}, Key=destination_key) + s3_client.copy_object(Bucket=source_bucket, CopySource={'Bucket': source_bucket, 'Key': source_key}, Key=destination_key, Tagging=f'{tags}') # Delete the file from the source folder s3_client.delete_object(Bucket=source_bucket, Key=source_key) @@ -70,12 +73,50 @@ def get_pdf_files_list_from_s3(source_bucket, source_folder): return file_list +def start_textract_detection_job( bucket_name, + document_file_name, + sns_topic_arn, + sns_role_arn, + job_tag,): + try: + # Define the parameters for the start_document_analysis API + start_document_detection_params = { + 'DocumentLocation': { + 'S3Object': { + 'Bucket': bucket_name, + 'Name': document_file_name + } + }, + 'ClientRequestToken': 'unique-token-'+str(generate_unix_timestamp()), # Use a unique token for each request + 'JobTag': job_tag, # Use a tag to identify your job + 'NotificationChannel': { + 'SNSTopicArn': sns_topic_arn, + 'RoleArn': sns_role_arn # Role to allow Textract service to notify SNS topic when response is ready + } + } + + logger.info('start_document_detection_params ' + str(start_document_detection_params)) + + # Send the request to start document detection + textract_response = textract_client.start_document_text_detection(**start_document_detection_params) + + job_id = textract_response["JobId"] + logger.info( + "Started text detection job %s on %s.", job_id, document_file_name + ) + except ClientError: + logger.exception("Couldn't detect text in %s.", document_file_name) + raise + else: + return job_id + def start_textract_analysis_job( bucket_name, document_file_name, analysis_feature_type, sns_topic_arn, sns_role_arn, + job_tag, ): try: @@ -89,7 +130,7 @@ def start_textract_analysis_job( }, 'FeatureTypes': analysis_feature_type, # Customize based on requirements 'ClientRequestToken': 'unique-token-'+str(generate_unix_timestamp()), # Use a unique token for each request - 'JobTag': 'healthcare-contract', # Use a tag to identify your job + 'JobTag': job_tag, # Use a tag to identify your job 'NotificationChannel': { 'SNSTopicArn': sns_topic_arn, 'RoleArn': sns_role_arn # Role to allow Textract service to notify SNS topic when response is ready @@ -150,58 +191,70 @@ def lambda_handler(event, context): SENDER_MAX_FILES = int(config_dict['OTHERS']['SENDER_MAX_FILES']) SNS_TOPIC_ARN = config_dict['RESOURCES']['SNS_TOPIC_ARN'].replace("{aws_region}",aws_region).replace("{aws_account_id}",aws_account_id) TEXTRACT_ROLE_ARN = config_dict['RESOURCES']['TEXTRACT_ROLE_ARN'].replace("{aws_account_id}",aws_account_id) # Textract IAM Role ARN to publish to SNS - + JOB_TAG = config_dict['OTHERS']['JOB_TAG'] + PROCESS_TYPE = str(config_dict['OTHERS']['PROCESS_TYPE']).upper() + logger.info('SOURCE_LOCATION: ' + SOURCE_LOCATION) logger.info('STAGING_LOCATION: ' + STAGING_LOCATION) logger.info('ANALYSIS_FEATURE_TYPE: ' + str(ANALYSIS_FEATURE_TYPE)) logger.info('SNS_TOPIC_ARN: ' + SNS_TOPIC_ARN) logger.info('TEXTRACT_ROLE_ARN: ' + TEXTRACT_ROLE_ARN) - logger.info('SENDER_MAX_FILES: ' + str(SENDER_MAX_FILES)) - - # List S3 Object & iterate (as per max files allowed) - files_list = get_pdf_files_list_from_s3(S3_BUCKET_NAME,SOURCE_LOCATION) - - if len(files_list): - # File count - file_count = 0 + logger.info('SENDER_MAX_FILES: ' + str(SENDER_MAX_FILES)) + logger.info('JOB_TAG: ' + str(JOB_TAG)) + logger.info('PROCESS_TYPE: ' + str(PROCESS_TYPE)) + + # File count + file_count = 0 + + # Process each message from the SQS event + for record in event['Records']: - for s3_file_key in files_list: + # Extract the message body from the record + record_body = json.loads(record['body']) + + #logger.info('Message Count: ', str(len(record_body['Records'])) ) + + for sqs_record in record_body['Records']: # Construct the source and destination paths - source_path = s3_file_key + source_path = unquote_plus(sqs_record['s3']['object']['key']) destination_path = STAGING_LOCATION + source_path.replace(SOURCE_LOCATION,"") # Move file to stagging move_file_within_s3(S3_BUCKET_NAME, source_path, destination_path) - # Start Textract analysis job - job_id = start_textract_analysis_job ( - S3_BUCKET_NAME, - destination_path, - ANALYSIS_FEATURE_TYPE, - SNS_TOPIC_ARN, - TEXTRACT_ROLE_ARN, - ) + job_id = "" + + if PROCESS_TYPE == "ANALYSIS": + # Start Textract analysis job + job_id = start_textract_analysis_job ( + S3_BUCKET_NAME, + destination_path, + ANALYSIS_FEATURE_TYPE, + SNS_TOPIC_ARN, + TEXTRACT_ROLE_ARN, + JOB_TAG, + ) + elif PROCESS_TYPE == "DETECTION": + # Start Textract detection job + job_id = start_textract_detection_job ( + S3_BUCKET_NAME, + destination_path, + SNS_TOPIC_ARN, + TEXTRACT_ROLE_ARN, + JOB_TAG, + ) file_count = file_count + 1 - logger.info(str(file_count) + '. ' + str(s3_file_key) + " Job Id: " + str(job_id)) - - if SENDER_MAX_FILES == file_count: - break + logger.info(str(file_count) + '. ' + str(source_path) + " Job Id: " + str(job_id)) - success_message = 'Total files sent to textract : '+ str(file_count) - logger.info(success_message) - return { - 'statusCode': 200, - 'body': success_message - } - else: - message = 'No files found' - logger.error(message) - return { - 'statusCode': 500, - 'body': message - } + success_message = 'Total files sent to textract : '+ str(file_count) + logger.info(success_message) + return { + 'statusCode': 200, + 'body': success_message + } + else: error_message = 'Incorrect value for ENVIRONMENT VARIABLES: PROPERTY_FILE_S3_PATH\r' + str(property_file_path)