From c50e76ddacb61b0df6db2e407dee6e8b4b410688 Mon Sep 17 00:00:00 2001 From: Umang Mistry Date: Tue, 25 Jun 2024 18:03:25 -0500 Subject: [PATCH] Data pipeline ingestion from orchestrator & pulled changes from UAT --- .../src/lambda/prompt-orchestrator/index.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/textract-pipeline/src/lambda/prompt-orchestrator/index.py b/textract-pipeline/src/lambda/prompt-orchestrator/index.py index 378e5da..2f2698a 100644 --- a/textract-pipeline/src/lambda/prompt-orchestrator/index.py +++ b/textract-pipeline/src/lambda/prompt-orchestrator/index.py @@ -16,12 +16,17 @@ import traceback # Global configuration variables S3_REGION = "us-east-2" BEDROCK_REGION = "us-east-1" -S3_RESULT_BUCKET = "doczy-dev-infra-raw-data-ingestion" SNOWFLAKE_SCHEMA = "STG" MAX_PROMPT_LENGTH_PER_GROUP = 20000 # Adjust as needed LLM_SELECTED = "Claude 3 - Sonnet" SNOWFLAKE_CONN_SECRET_NAME = "doczy_dev_db_creds" +# Get variable ENVIRONMENT from env variables using os +env = os.environ.get('ENVIRONMENT') +# Get first character from env to get env name abbreviation +env = env[0] +# Get the bucket name based on the env that is linked with the storage integration +snowflake_ingestion_bucket = f' doczyai-use2-{env}-infra-s3-raw-data-ingestion' # Configure logging logger = logging.getLogger() @@ -415,7 +420,7 @@ def main(bucket, object_key, field_groups): df['Snippet'] = snippet_list df['New Page Number'] = page_no_list df['New Page Number'] = df['New Page Number'].apply(lambda x: re.search(r'\d+', x).group() if isinstance(x, str) and re.search(r'\d+', x) is not None else " ") - df['Revised Prompt'] = [question] * len(contract_list_f) + # df['Revised Prompt'] = [question] * len(contract_list_f) # Need to add batch_id to the df df['Batch_id'] = batch_id @@ -425,9 +430,17 @@ 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_{timestamp}.csv" + file_name = f"results_{batch_id}_{timestamp}.csv" s3_client.put_object(Bucket=bucket, Body=csv_buf.getvalue(), Key=f"final_output/{batch_id}/{file_name}") csv_buf = StringIO() + + # 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('{file_name}')" + cur.execute(load_data_query) + + # history.tail(1).to_csv(csv_buf, header=True, index=False) # csv_buf.seek(0) # timestamp = datetime.now().strftime("%Y%m%d%H%M%S")