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'(.*?)' 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