Data pipeline ingestion from orchestrator & pulled changes from UAT

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