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

Updated lambda handler & removed keys from config

Browse Source
This commit is contained in:
Umang Mistry
2024-06-28 17:45:28 -05:00
parent efb38bece9
commit 9268a55184
2 changed files with 104 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
textract-pipeline/src/lambda/prompt-orchestrator/config.py
+2 -12
View File
@@ -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'
textract-pipeline/src/lambda/prompt-orchestrator/index.py
+102 -10
View File
@@ -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)