mirror of
https://github.com/x1xhlol/system-prompts-and-models-of-ai-tools.git
synced 2026-08-16 21:24:09 +00:00
250 lines
7.6 KiB
Python
250 lines
7.6 KiB
Python
import os
|
|
import boto3
|
|
|
|
from s3Ops import (
|
|
list_files_in_s3_folder
|
|
)
|
|
|
|
from utils.prompts import DOCUMENT_CLASSIFY_PROMPT
|
|
|
|
from backendAPIs import (
|
|
get_onboarding_status,
|
|
update_onboarding_status
|
|
)
|
|
|
|
bucket_name = os.getenv('S3_UPLOAD_BUCKET_NAME')
|
|
|
|
def document_classifier(record_details):
|
|
|
|
FINAL_DOCUMENT_CLASSIFY_PROMPT = DOCUMENT_CLASSIFY_PROMPT.format(document=record_details)
|
|
|
|
json_string = llm_bedrock.invoke(FINAL_DOCUMENT_CLASSIFY_PROMPT).content
|
|
|
|
print(json_string)
|
|
|
|
pattern = r'<output>(.*?)</output>'
|
|
match = re.search(pattern, json_string, re.DOTALL)
|
|
|
|
if match:
|
|
return match.group(1).strip()
|
|
else:
|
|
return ""
|
|
|
|
|
|
def get_pdf_details(key):
|
|
# Create a temporary file for the PDF
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf', dir='/tmp') as temp_file:
|
|
local_file_path = temp_file.name
|
|
# Download the file from S3
|
|
s3_client.download_file(bucket_name, key, local_file_path)
|
|
|
|
with pdfplumber.open(local_file_path) as pdf:
|
|
list_pages = []
|
|
|
|
for page in pdf.pages:
|
|
text = page.extract_text()
|
|
if text:
|
|
list_pages.append(text.strip())
|
|
|
|
pdf_string = "".join(list_pages)
|
|
|
|
if len(pdf_string.strip()) > 0:
|
|
# Clean up the temporary PDF file
|
|
os.remove(local_file_path)
|
|
return pdf_string
|
|
else: # start image processing
|
|
print("Starting image processing")
|
|
# Create a temporary directory for storing images
|
|
with tempfile.TemporaryDirectory(dir='/tmp') as temp_dir:
|
|
pdf_images_list = pdf_to_images(local_file_path, temp_dir)
|
|
result = read_images(pdf_images_list)
|
|
|
|
# Clean up the temporary PDF file
|
|
os.remove(local_file_path)
|
|
|
|
return result
|
|
|
|
def process_single_pdf(pdf_file):
|
|
pdf_content = get_pdf_details(key)
|
|
classification = document_classifier(pdf_content)
|
|
return pdf_file, classification
|
|
|
|
|
|
def pdf_classifier(files_list):
|
|
|
|
categories = {
|
|
"LPA(Limited Partner Agreement)": [],
|
|
"Capital Statements": [],
|
|
"Sub Docs(Subscription Document)": [],
|
|
"Side Letter": [],
|
|
"Capital Call Notices": [],
|
|
"Distribution Notices": [],
|
|
"Fund Reports": [],
|
|
"K1s": [],
|
|
"Other": []
|
|
}
|
|
|
|
with ThreadPoolExecutor(max_workers=10) as executor:
|
|
future_to_pdf = {executor.submit(process_single_pdf, pdf_file): pdf_file for pdf_file in pdf_files}
|
|
|
|
for future in as_completed(future_to_pdf):
|
|
pdf_file = future_to_pdf[future]
|
|
try:
|
|
pdf_file, classification = future.result()
|
|
|
|
# Determine which category the PDF belongs to
|
|
found_category = False
|
|
for category in categories.keys():
|
|
if category.lower() in classification.lower():
|
|
categories[category].append(pdf_file)
|
|
found_category = True
|
|
break
|
|
|
|
if not found_category:
|
|
categories["Other"].append(pdf_file)
|
|
|
|
except Exception as exc:
|
|
print(f'{pdf_file} generated an exception: {exc}')
|
|
categories["Other"].append(pdf_file)
|
|
|
|
return categories
|
|
|
|
def validate_input_files(file_list):
|
|
"""
|
|
Function to find specific files in a list based on naming patterns.
|
|
|
|
Args:
|
|
file_list (str): List of files to validate.
|
|
|
|
Returns:
|
|
dict: Dictionary containing file names that match the required patterns.
|
|
str: Error message if any required file is missing.
|
|
"""
|
|
# Patterns to look for
|
|
required_files = {
|
|
# "lpa": "_LP",
|
|
"partner": "_partner",
|
|
"financials": "_financials",
|
|
"bankTransactions": "_bank_transactions",
|
|
"journals": "_journals",
|
|
"fund_performance": "_fund-performance"
|
|
}
|
|
|
|
# Dictionary to hold found file names
|
|
found_files = {
|
|
# "lpa": "",
|
|
"partner": None,
|
|
"financials": None,
|
|
"bankTransactions": None,
|
|
"journals": None,
|
|
"fund_performance": None
|
|
}
|
|
|
|
pdf_files = []
|
|
|
|
try:
|
|
|
|
if file_list:
|
|
# Loop through files in the list
|
|
for file_name in file_list:
|
|
if required_files['partner'] in file_name:
|
|
found_files['partner'] = file_name
|
|
elif required_files['financials'] in file_name:
|
|
found_files['financials'] = file_name
|
|
elif required_files['bankTransactions'] in file_name:
|
|
found_files['bankTransactions'] = file_name
|
|
elif required_files['journals'] in file_name:
|
|
found_files['journals'] = file_name
|
|
elif required_files['fund_performance'] in file_name:
|
|
found_files['fund_performance'] = file_name
|
|
elif file_name.lower().endswith('.pdf'):
|
|
pdf_files.append(file_name)
|
|
|
|
else:
|
|
return(f"Error: No files found in folder.")
|
|
|
|
# if len(pdf_files) > 0:
|
|
# classified_pdf_files = pdf_classifier(pdf_files)
|
|
#
|
|
except FileNotFoundError:
|
|
return f"Error: Unable to find file.",""
|
|
|
|
# Check if all required files are found
|
|
missing_files = [key for key, value in found_files.items() if value is None]
|
|
|
|
if missing_files:
|
|
return f"Error: Missing required files - {', '.join(missing_files)}",""
|
|
|
|
# return found_files,classified_pdf_files
|
|
|
|
return found_files , pdf_files
|
|
|
|
|
|
def pretty_print_files(files_dict):
|
|
"""
|
|
Pretty print the dictionary containing file names.
|
|
|
|
Args:
|
|
files_dict (dict): Dictionary where the keys are the file categories, and the values are the file names.
|
|
"""
|
|
file_message = ""
|
|
for key, value in files_dict.items():
|
|
file_message += f"- {key.replace('_', ' ').capitalize()}: {value} \n"
|
|
|
|
return file_message
|
|
|
|
|
|
def run_file_validation():
|
|
|
|
bucket_name = os.getenv("S3_UPLOAD_BUCKET_NAME")
|
|
step_number = 1
|
|
# Array index starts from 0
|
|
step_number -= 1
|
|
|
|
|
|
print(f"Getting files from bucket {bucket_name}")
|
|
success_message = ""
|
|
error_message = ""
|
|
status = ""
|
|
|
|
try:
|
|
|
|
onboarding_id = os.getenv('ONBOARDING_ID')
|
|
print(f"Get onBoardingStatus for {onboarding_id}")
|
|
response = get_onboarding_status(onboarding_id)
|
|
# Get current onboarding status values.
|
|
onboarding_status = response['data']['data']
|
|
success_message = onboarding_status['steps'][step_number]['summary']
|
|
error_message = onboarding_status['steps'][step_number]['errorMessage']
|
|
status = onboarding_status['steps'][step_number]['status']
|
|
folder_path = onboarding_status['path']
|
|
|
|
print(f"Pulling files list from {bucket_name}/{folder_path}")
|
|
file_list = list_files_in_s3_folder(bucket_name, folder_path)
|
|
print("S3 files:", file_list)
|
|
found_files ,classified_files = validate_input_files(file_list)
|
|
|
|
print("Files found:", found_files)
|
|
|
|
if isinstance(found_files, dict):
|
|
success_message = pretty_print_files(found_files)
|
|
status = "COMPLETE"
|
|
#Validation failed
|
|
else:
|
|
error_message = found_files
|
|
status = "FAILED"
|
|
success_message = ""
|
|
|
|
response = update_onboarding_status(step_number, status, error_message, success_message)
|
|
if 'error' in response:
|
|
print(f"Failed to fetch onboarding status: Error: {response['error']}")
|
|
print(f"Status Code: {response['status_code']}")
|
|
return False , False
|
|
|
|
return found_files , classified_files
|
|
|
|
except Exception as e:
|
|
return False , False
|
|
|
|
|
|
|