mirror of
https://github.com/x1xhlol/system-prompts-and-models-of-ai-tools.git
synced 2026-08-16 21:24:09 +00:00
280 lines
8.7 KiB
Python
280 lines
8.7 KiB
Python
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)
|