system-prompts-and-models-o.../task/onboardFinancials.py
2026-06-18 21:23:01 +05:30

244 lines
8.5 KiB
Python

import pandas as pd
import json
import os
import re
from s3Ops import read_file_from_s3
from utils.prompts import ZIVE_SECURITY_TYPE_PROMPT
from backendAPIs import (
update_onboarding_status,
add_portfolio_company,
add_portfolio_company_investment
)
# LLM via the unified Gemini-only backend (see task/llm.py).
from llm import get_chat_model
model_parameter = {"temperature": 0, "max_tokens": 16384}
llm_bedrock = get_chat_model(model_kwargs=model_parameter)
def xlsx_to_df(xlsx_file, sheet_name):
# Read the Excel file, skipping the header rows
df = pd.read_excel(xlsx_file, sheet_name=sheet_name, header=5, na_values=[''])
print(df.columns)
investments = {}
current_investment = None
# Iterate through the rows
for index, row in df.iterrows():
if pd.notna(row['Investment legal name']):
# Start a new investment
current_investment = row['Investment legal name']
print(current_investment)
investments[current_investment] = {'securities': [], 'security_count': 0}
elif pd.notna(row['Security type']) and current_investment is not None:
security_type = row['Security type'].strip().lower()
if security_type != 'totals' and security_type != 'grand totals':
# Increment the security count for this investment
investments[current_investment]['security_count'] += 1
# Add security details to the current investment
security = {
'type': row['Security type'],
'investment_date': row['Investment date'],
'investment_number': investments[current_investment]['security_count']
}
investments[current_investment]['securities'].append(security)
return investments
# TODO: add the remaining asset type mapping
# def get_security_type(asset_type):
# asset_to_security_type = {
# "Common stock":"COMMON",
# "Preferred stock":"PREFERRED",
# "Warrants":"WARRANTS",
# "Convertible promissory note/SAFEs": "POST MONEY SAFE"
# }
#
# for asset, security_type in asset_to_security_type.items():
# if asset_type.lower() in asset.lower(): # Case-insensitive match
# return security_type
#
# return ""
def get_security_type(asset_type):
print(asset_type)
print("Calling LLM to get the security type")
FINAL_ZIVE_SECURITY_TYPE_PROMPT = ZIVE_SECURITY_TYPE_PROMPT.format(security_type=asset_type)
output = llm_bedrock.invoke(FINAL_ZIVE_SECURITY_TYPE_PROMPT).content
pattern = r'<output>(.*?)</output>'
match = re.search(pattern, output, re.DOTALL)
final_data = match.group(1).strip()
return final_data
def process_financial_records(investments_df):
"""
Process and insert users from an Excel file into MongoDB.
Args:
file_path (str): Path to the Excel file.
"""
entity_id = os.getenv('ENTITY_ID')
fund_id = os.getenv('FUND_ID')
# print(f"Header:\n {partners_df.head()}")
investments = []
# current_company = None
# portfoltio_company_id = None
# investment_number = 0
for investment, details in investments_df.items():
print(f"Adding company {investment}")
payload = {
"fundId": fund_id,
"companyName": investment
}
print(f"\nInvestment: {investment}")
response = add_portfolio_company(payload)
if 'error' in response:
print(f"Failed to fetch onboarding status: Error: {response['error']}")
print(f"Status Code: {response['status_code']}")
continue
portfoltio_company_id = response["data"]["data"]["_id"]
for security in details['securities']:
print(f" Security Type: {security['type']}")
print(f" Investment Date: {security['investment_date']}")
print(f" Investment Number: {security['investment_number']}")
investment = {
"fundId": fund_id,
"portfolioId": portfoltio_company_id,
"investmentNumber": security['investment_number'],
"investmentDate": security['investment_date'],
"securityType": get_security_type(security['type']),
"assetName": security['type'],
"asset": security['type']
}
print(investment)
response = add_portfolio_company_investment(investment)
if 'error' in response:
print(response)
print(f"Failed to fetch onboarding status: Error: {response['error']}")
print(f"Status Code: {response['status_code']}")
continue # Process next company investment record...
investments.append(investment)
# Iterate through the DataFrame and prepare partner records
# for index, row in financials_df.iterrows():
#
# if current_company != row['Investment']:
# current_company = row['Investment']
# investment_number = 0
#
# # make the call to add portfolio company
# print(f"Adding company {current_company}")
# payload = {
# "fundId": fund_id,
# "companyName": current_company
# }
# response = add_portfolio_company(payload)
# if 'error' in response:
# print(f"Failed to fetch onboarding status: Error: {response['error']}")
# print(f"Status Code: {response['status_code']}")
# break #Process next company record...
# portfoltio_company_id = response["data"]["data"]["_id"]
#
#
# investment_number += 1
# asset_type = row['Asset type'] if pd.notna(row['Asset type']) else None
#
# investment = {
# "fundId": fund_id,
# "portfolioId": portfoltio_company_id,
# "investmentNumber": investment_number,
# "investmentDate": row['Investment date'] if pd.notna(row['Investment date']) else None,
# "securityType": get_security_type(asset_type),
# "assetName": row['Asset'] if pd.notna(row['Asset']) else None,
# "asset": asset_type
# }
#
# response = add_portfolio_company_investment(investment)
# if 'error' in response:
# print(f"Failed to fetch onboarding status: Error: {response['error']}")
# print(f"Status Code: {response['status_code']}")
# break # Process next company investment record...
# investments.append(investment)
# print(f"Adding portforlio investment {index} {current_company}:{portfoltio_company_id}:\n {investment}")
# Test with smaller set of records
# if index == 1:
# break
return investments
def process_financials(file_path):
# Process partner data
step_number = 4
step_number -= 1
success_message = ""
error_message = ""
status = "IN-PROGRESS"
bucket_name = os.getenv('S3_UPLOAD_BUCKET_NAME')
print(f"Processing file: {file_path}")
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
partner_excel_file = read_file_from_s3(bucket_name, file_path)
sheet_name = 'Roll Forward'
financial_df = xlsx_to_df(partner_excel_file, sheet_name)
# print(financial_df.head())
financials = process_financial_records(financial_df)
if financials:
item_count = len(financials)
success_message = f"Onboarded {item_count} portfolio financial records."
status = "COMPLETE"
else:
error_message = "No transactions were onboaerded."
status = "FAILED"
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
# print(f"Investments list:\n {financials}")
# print(json.dumps(partners_dict, indent=2, sort_keys=True, cls=CustomJSONEncoder))
return True