diff --git a/textract-pipeline/src/lambda/prompt-orchestrator/config.py b/textract-pipeline/src/lambda/prompt-orchestrator/config.py index 02baf9f..0ee313b 100644 --- a/textract-pipeline/src/lambda/prompt-orchestrator/config.py +++ b/textract-pipeline/src/lambda/prompt-orchestrator/config.py @@ -75,30 +75,20 @@ VALID_COLUMNS = ['Filename', 'Corrected_PROGRAM', 'Corrected_NETWORK'] -# AWS Keys -AWS_ACCESS_KEY_ID="ASIAZTMXAXNXBHQKW32L" -AWS_SECRET_ACCESS_KEY="IozsUqPsyXu2klzNMU42CDZ1ZtGz8Yo29mhuJRaQ" -AWS_SESSION_TOKEN="IQoJb3JpZ2luX2VjEDsaCXVzLWVhc3QtMiJHMEUCIFkt/LwV1X3UmAyNi25xrDIL+mdaV7eV3U03VUQPc0o/AiEA4R54Gyd4d/55N5P0a84zzHiPUuejjyEMF6ARAk0i2fgqjAMI5P//////////ARAAGgw2NjAxMzEwNjg3ODIiDC+Kl0HOGnW+Vct/FyrgAg+t9ssFLXSj0VFXfVud6jLJnIAChoXSikj59wmU1hX3vcsN699t4UT5vfAjIwZ3emEPrfaKnuvcdC4ev2V4fT3IxRS8FLu6eUKqImF8ZkATIxg0nYpPqx+92CBvcT1gjLv4Y/1oorb0GxLcTiZLBY6IrCbrr3bDyexU+fgqoGlVFhaqfVPotvXy5+s0RbFWF7wQ0IanZ4eX5ZTFUXZMJh0LBYSTIuas6oDF2z/k/De48NKOWEZCmVxIG/ykyIbBujSICetKKX1qGtB74NrxSfhERGfof4DjZl5q4SUEouTRVb5aGk779lqdKjTXjmGyWyryzUu9jJfqyZiC+ymcIliK4aTn4W88/prqx3p1mq+/T1e5uMhtXi/9JURT3P/avOVz34j5ohggUxqvl2TBASKaUGLrOIBoJzoMevgodoWy1TleJe7xJWBOS5nq5knxMOreHjuImgNofmiRuSNbayowpqTzswY6pgFwPrVOCxqQyxXwNsCEB27me72Ez3splxi+JYiYyaxkANoftb6d1/JyDBprnsact7hvD3mPVUhjdvsh92o5+F1lR0usiyCqrvUrXgpYKdBh6jIJh2kcNDSYwoGMHI2kVb2G53GWe1Vb06cS79Ei3QVurxVe7SVYVqd+bpicuYDdcrvdzCeF2E+9P9qEQa1/4B6ei+0O+lao4WBjvy7ZyuAlwaUUgsw0" # File Paths LOCAL_PATH = 'data/test/' # Replace with local # S3 Settings S3_CLIENT = boto3.client('s3', - region_name="us-east-2", - aws_access_key_id=AWS_ACCESS_KEY_ID, - aws_secret_access_key=AWS_SECRET_ACCESS_KEY, - aws_session_token=AWS_SESSION_TOKEN + region_name="us-east-2" ) BUCKET = "doczy-dev-infra-textract" PREFIX = "batches/batch_1/contract-text-file/" # replace with s3 path # Bedrock Settings BEDROCK_RUNTIME = boto3.client(service_name="bedrock-runtime", - region_name="us-east-1", - aws_access_key_id=AWS_ACCESS_KEY_ID, - aws_secret_access_key=AWS_SECRET_ACCESS_KEY, - aws_session_token=AWS_SESSION_TOKEN + region_name="us-east-1" ) MODEL_ID_CLAUDE3_HAIKU = 'anthropic.claude-3-haiku-20240307-v1:0' MODEL_ID_CLAUDE3_SONNET = 'anthropic.claude-3-sonnet-20240229-v1:0' diff --git a/textract-pipeline/src/lambda/prompt-orchestrator/index.py b/textract-pipeline/src/lambda/prompt-orchestrator/index.py index e447663..9df42d7 100644 --- a/textract-pipeline/src/lambda/prompt-orchestrator/index.py +++ b/textract-pipeline/src/lambda/prompt-orchestrator/index.py @@ -12,6 +12,7 @@ import logging from urllib.parse import unquote_plus from botocore.exceptions import ClientError import traceback +import file_processing # Global configuration variables S3_REGION = "us-east-2" @@ -117,6 +118,12 @@ def save_to_sf(dag_name, **kwargs): return payload +def parse_field_groups(field_groups): + # Split the string by comma and strip any surrounding whitespace + groups = [group.strip() for group in field_groups.split(',')] + # Convert the list to a tuple + return tuple(groups) + def lambda_handler(event, context): try: logger.info("Processing SQS event.") @@ -136,17 +143,75 @@ def lambda_handler(event, context): # Fetch batch_id from object tags 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 - prompt_config_query = f"SELECT DISTINCT GROUP_ID FROM DOCUMENT_LOGS WHERE DOCUMENT_ID = '{document_id}'" - # field_group_df = read_from_db(prompt_config_query) - # field_groups = tuple(field_group_df['GROUP_ID'].unique()) - field_groups = ('A','C') - logger.info(f"Field groups: {field_groups}") - # Prepare and process the prompt - main(bucket_name, object_key, field_groups) + # Get file name from s3 bucket object key + try: + logger.info(f"OBJECT KEY:: {object_key}") + # Get file name from s3 bucket object key and remove file extension + file_name = object_key.split('/')[-1].split('.')[0] + cursor = snowflake_conn.cursor() + result = cursor.execute(f"select group_id from stg.document_logs where batch_id= '{batch_id}' and document_id='{file_name}' limit 1;") + logger.info(f"QUERY BEING EXECUTED: select group_id from stg.document_logs where batch_id= '{batch_id}' and file_name='{file_name}' limit 1;") + + + for rec in result: + logger.info(f"QUERY RESULT::: {rec[0]}") + field_groups = rec[0] + + # logger.info(f"QUERY RESULT::: {cursor.fetchone()} & {cursor.fetchone()[0]}&fetching all {cursor.fetchall()}") + # field_groups = cursor.fetchone()[0] + # logger.info(f"RESULT FIELD GROUP:: {field_groups}") + # Check if field group is a tuple, if not convert it to tuple + field_groups = parse_field_groups(field_groups) + logger.info(f"RESULT FIELD GROUP after tuple conversion:: {field_groups}") + + except Exception as e: + logger.error(f"Error fetching field groups: {e} stopping execution") + exit() + field_groups = ('A','C') + + + # field_groups = ('A','C') + # field_groups = 'B' + # logger.info(f"Field groups: {field_groups}") + if field_groups == ('B',): + data = s3_client.get_object(Bucket=bucket_name, Key=object_key) + contents = data['Body'].read() + context = contents.decode("utf-8") + # Get the file name from the s3 object prefix + filename = object_key.split('/')[-1] + process_b_fields(filename, context, bucket_name,batch_id) + # Prepare and process the prompt if field groups are A and C or Just A or just C + elif field_groups == ('A','C'): + main(bucket_name, object_key, field_groups) + # Process scenario where all field groups are processed + elif field_groups == ('A','B','C'): + main(bucket_name, object_key, ('A','C')) + data = s3_client.get_object(Bucket=bucket_name, Key=object_key) + contents = data['Body'].read() + context = contents.decode("utf-8") + # Get the file name from the s3 object prefix + filename = object_key.split('/')[-1] + process_b_fields(filename, context, bucket_name,batch_id) + elif field_groups == ('A',): + main(bucket_name, object_key, field_groups) + elif field_groups == ('C',): + main(bucket_name, object_key, field_groups) + elif field_groups == ('B','A') or field_groups == ('B','C') or field_groups == ('A','B') or field_groups == ('C','B'): + main(bucket_name, object_key, ('A','C')) + data = s3_client.get_object(Bucket=bucket_name, Key=object_key) + contents = data['Body'].read() + context = contents.decode("utf-8") + # Get the file name from the s3 object prefix + filename = object_key.split('/')[-1] + process_b_fields(filename, context, bucket_name,batch_id) + else: + logger.error(f"Unsupported field groups: {field_groups}") + return { + 'statusCode': 500, + 'body': f"Unsupported field groups: {field_groups}" + } return { 'statusCode': 200, @@ -159,6 +224,33 @@ def lambda_handler(event, context): 'body': f"{str(e)} \n {traceback.format_exc()}" } + +def process_b_fields(filename, context, bucket, batch_id): + # Create tuple from filename and context + item = (filename, context) + df = file_processing.process_file(item) + # Adding batch_id to the last column of df + df['Batch_id'] = batch_id + csv_buf = StringIO() + df.to_csv(csv_buf, header=True, index=False) + csv_buf.seek(0) + timestamp = datetime.now().strftime("%Y%m%d%H%M%S") + file_name = f"results_{batch_id}_B_{timestamp}.csv" + s3_client.put_object(Bucket=bucket, Body=csv_buf.getvalue(), Key=f"final_output/{batch_id}/{file_name}") + csv_buf = StringIO() + df.to_csv(csv_buf, header=True, index=False) + csv_buf.seek(0) + + # Saving data to snowflake raw data ingestion bucket & calling stored proc + s3_client.put_object(Bucket=snowflake_ingestion_bucket, Body=csv_buf.getvalue(), Key=f"doczy_pipeline_output/{file_name}") + # cur = snowflake_conn.cursor() + # load_data_query = f"CALL LOAD_DOCZY_PIPELINE_RAW_OUTPUT_B('{file_name}')" + # cur.execute(load_data_query) + # log df shape as output + logger.info(f"B fields Processed file: {filename} with shape: {df.shape}") + + + def get_filename_from_path(full_path): return os.path.basename(full_path) @@ -434,7 +526,7 @@ def main(bucket, object_key, field_groups): df.to_csv(csv_buf, header=True, index=False) csv_buf.seek(0) timestamp = datetime.now().strftime("%Y%m%d%H%M%S") - file_name = f"results_{batch_id}_{timestamp}.csv" + file_name = f"results_{batch_id}_AC_{timestamp}.csv" s3_client.put_object(Bucket=bucket, Body=csv_buf.getvalue(), Key=f"final_output/{batch_id}/{file_name}") csv_buf = StringIO() df.to_csv(csv_buf, header=True, index=False)