jacob-mathison/doczyai-pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Added detection capability and sender orchectration

Browse Source
This commit is contained in:
Pankaj Katariya
2024-02-23 15:50:03 +05:30
parent 5d90d22d07
commit c182135ebc
2 changed files with 157 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lambda/receiver/lambda_function.py
+66 -4
View File
@@ -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,"")
lambda/sender/lambda_function.py
+91 -38
View File
@@ -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)