From c8ef3f36c705cb44bae1e1f19d4e7c4f7da9dd2e Mon Sep 17 00:00:00 2001 From: Umang Mistry Date: Fri, 21 Jun 2024 12:06:51 -0500 Subject: [PATCH] Changed SQS message parsing --- textract-pipeline/src/lambda/prompt-orchestrator/index.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/textract-pipeline/src/lambda/prompt-orchestrator/index.py b/textract-pipeline/src/lambda/prompt-orchestrator/index.py index bb260e2..86d9f63 100644 --- a/textract-pipeline/src/lambda/prompt-orchestrator/index.py +++ b/textract-pipeline/src/lambda/prompt-orchestrator/index.py @@ -115,8 +115,10 @@ def save_to_sf(dag_name, **kwargs): def lambda_handler(event, context): try: logger.info("Processing SQS event.") + records = json.loads(event['Records'][0]['body']) + logger.info(f"PROCESSING RECORDS:: {records}") # Process each message from the SQS event - for record in event['Records']: + for record in records['Records']: # Extract the message body from the record sqs_record = record['s3'] bucket_name = sqs_record['bucket']['name'] @@ -127,7 +129,7 @@ def lambda_handler(event, context): logger.info(f"Processing document: {document_id}") # Fetch batch_id from object tags - batch_id = get_s3_object_tags(bucket_name, object_key).get('batch_id') + batch_id = get_s3_object_tags(bucket_name, object_key).get('BatchId') logger.info(f"Batch ID: {batch_id}") # Query field groups from the database @@ -387,7 +389,7 @@ def main(bucket, object_key, field_groups): context = contents.decode("utf-8") # Fetch batch_id from object tags - batch_id = get_s3_object_tags(bucket, object_key).get('batch_id') + batch_id = get_s3_object_tags(bucket, object_key).get('BatchId') logger.info(f"Batch ID: {batch_id}") df = pd.DataFrame(columns=['Contract Name', 'Field Name', 'Raw value', 'New Extracted value', 'Confidence Level', 'Snippet', 'New Page Number', 'Revised Prompt'])