from PIL import Image import boto3 import fitz import re import io import os import json from llm import get_llm_client import tiktoken import base64 import time from s3Ops import ( read_file_from_s3 ) def pdf_to_image(pdf_key, input_bucket_name, output_bucket_name): s3 = boto3.client('s3') print(f"Reading file from {input_bucket_name}, with key {pdf_key}") # Load the PDF file into memory pdf_object = s3.get_object(Bucket=input_bucket_name, Key=pdf_key) pdf_content = pdf_object['Body'].read() # Extract file name from pdf_key file_name = pdf_key.split('/')[-1] output_file_name = file_name.rsplit('.', 1)[0] # Remove file extension image_list = [] # Open the PDF pdf_document = fitz.open(stream=pdf_content, filetype="pdf") for page_num in range(len(pdf_document)): page = pdf_document.load_page(page_num) scale = 1.8 pix = page.get_pixmap(matrix=fitz.Matrix(scale, scale)) # Convert pixmap to PNG bytes directly img_bytes = pix.tobytes("jpg") image_key = f"{output_file_name}/{output_file_name}_{page_num + 1}.jpg" image_list.append(image_key) # Upload the image to S3 s3.put_object(Bucket=output_bucket_name, Key=image_key, Body=img_bytes, ContentType='image/jpg') return image_list def get_images(image_key, bucket): s3 = boto3.client('s3') # Get the object from S3 response = s3.get_object(Bucket=bucket, Key=image_key) # Read the content of the file image_data = response['Body'].read() return image_data # for use with OCR/last pages def create_message(extracted_texts, prompt, json_template, last_images64): message = { "role": "user", "content": [ { "type": "text", "text": f"Prompt:\n{prompt}\n\n" f"JSON Template:\n```json\n{json_template}```\n\n" f"Text to analyze:\n" } ] } # Add extracted texts for index, text in enumerate(extracted_texts): message["content"].append({ "type": "text", "text": f"Page {index + 1} limited partnership agreement document:\n{text}\n" }) # Add a separator before images message["content"].append({ "type": "text", "text": "\nLast ten corresponding images to analyze:\n" }) # Add the last few images to the message for i, image_data in enumerate(last_images64): # Add the image message["content"].append({ "type": "image", "source": { "type": "base64", "media_type": "image/jpeg", "data": image_data } }) # Add a text description for each image message["content"].append({ "type": "text", "text": f"Image {i + 1} (corresponds to Page {len(extracted_texts) - len(last_images64) + i + 1})\n" }) return message def ask_llm(message): # increase timeout # (boto3/Config removed — the Gemini client manages its own timeouts) # initialize unified Gemini LLM client (.invoke_model() drop-in; see task/llm.py) bedrock = get_llm_client() # construct the request body for bedrock API body = json.dumps({ "max_tokens": 16384, "system": "You are a financial advisor specializing in extracting and evaluating key information from financial documents. We are discussing terms within a limited partnership agreement document.", "messages": [ message, {"role": "assistant", "content": '```json'} # Prefill here ], "anthropic_version": "bedrock-2023-05-31" }) # invoke the model response = bedrock.invoke_model(body=body, modelId="anthropic.claude-3-5-sonnet-20240620-v1:0") # parse the response response_body = json.loads(response.get("body").read()) # return the OCR result return response_body['content'][0]['text'] def detailed_token_count(message, model="cl100k_base"): enc = tiktoken.get_encoding(model) total_tokens = 0 breakdown = {} # Count tokens for the role role_tokens = len(enc.encode(message["role"])) total_tokens += role_tokens breakdown["role"] = role_tokens # Count tokens for each content item for i, item in enumerate(message["content"]): if item["type"] == "text": text_tokens = len(enc.encode(item["text"])) total_tokens += text_tokens breakdown[f"content_{i}"] = text_tokens # Add tokens for the "type" field itself type_tokens = len(enc.encode(item["type"])) total_tokens += type_tokens breakdown[f"type_{i}"] = type_tokens return total_tokens, breakdown def clean_json(text): # Find the start and end of the JSON object start = output.find('{') end = output.rfind('}') + 1 # Extract the JSON string json_str = output[start:end] # Parse the JSON string to ensure it's valid try: json.loads(json_str) except json.JSONDecodeError: return None # Return None if the JSON is invalid return json_str def add_backtack(text): # Find the position of the first '{' start_pos = text.find('{') # If there's no '{', return the original text if start_pos == -1: return text # Find the start of the line containing '{' line_start = text.rfind('\n', 0, start_pos) if line_start == -1: line_start = 0 else: line_start += 1 # Move past the newline character # Insert '```json\n' before the line containing '{' return text[:line_start] + '```json\n' + text[line_start:] def onboardFunds(lpa_file_path): input_bucket = os.getenv('S3_UPLOAD_BUCKET_NAME') output_bucket = input_bucket image_keys = pdf_to_image(lpa_file_path, input_bucket, output_bucket) images = [] for image_key in image_keys: image = get_images(image_key, output_bucket) images.append(image) print(image_keys[-6:-1]) os.environ['TESSDATA_PREFIX'] = 'RTC:pdf processing/eng.traineddata' extracted_texts2 = [] start_time = time.time() for i, image in enumerate(images, 1): # Open the image using PIL img = Image.open(io.BytesIO(image)).convert('L') # Append the extracted text to the list extracted_texts2.append(text) # Print status update every 10 images if i % 10 == 0: elapsed_time = time.time() - start_time print(f"Processed {i} images out of {len(images)} in {elapsed_time:.2f} seconds") # Calculate total time total_time = time.time() - start_time # Print final status after processing all images print(f"Finished processing all {len(images)} images in {total_time:.2f} seconds") print(f"Average time per image: {total_time/len(images):.2f} seconds") last_images64 = [] for image in images[-10:]: image = base64.b64encode(image).decode("utf-8") last_images64.append(image) print(len(extracted_texts2)) print(extracted_texts2[-2]) message = create_message(extracted_texts2, prompt2, json_template, last_images64) #print(f"\nFinal message structure:") #print(f"Number of content items: {len(message['content'])}") #for i, item in enumerate(message['content']): # print(f"Content {i} length: {len(item['text'])}") # Use the detailed token count function total_count, breakdown = detailed_token_count(message) print(f"\nTotal token count: {total_count}") #print("Token breakdown:") #for key, value in breakdown.items(): # print(f" {key}: {value}") outputs = ask_llm(message) print(outputs) if __name__ == "__main__": from dotenv import load_dotenv load_dotenv() onboardingId = "66ea163564e2f97a059160ef" os.environ["ONBOARDING_ID"] = onboardingId from initOnboarding import initialize_onboarding initialize_onboarding() found_files = {'lpa': '66c5e5c99ecbf552a05b84f9/20240917T235220Z/0-Please_DocuSign_CerraCap_II_LP_Limited_Partn.pdf', 'partner': '66c5e5c99ecbf552a05b84f9/20240917T235220Z/0-1-cerracap-ii-lp_2024-07_09_short_partner.xlsx', 'financials': '66c5e5c99ecbf552a05b84f9/20240917T235220Z/2-3-cerracap-ii-lp_2024-08-26_financials.xlsx', 'bankTransactions': '66c5e5c99ecbf552a05b84f9/20240917T235220Z/4-CerraCap_II__LP_bank_transactions_2016-01-01-2024-07-03.xlsx', 'journals': '66c5e5c99ecbf552a05b84f9/20240917T235220Z/5-cerracap-ii-lp_2024-07-09_journals-export.xlsx', 'fund_performance': '66c5e5c99ecbf552a05b84f9/20240917T235220Z/6-cerracap-ii-lp_2024-07-09_fund-performance-report.xlsx'} lpa_file_path = found_files["lpa"] onboardFunds(lpa_file_path)