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,"")