Katon Minhas
23af15306b
Feature/new model testing * rename and set up * Setup * Get dynamic primary crosswalked answers * Generate stats for one file * Update dynamic primary evaluation * Multithreading * Set up reimbursement primary evaluation * UPdate output report generation * Update TN calculation * update poetry lock * reorder reimb primary * evaluate 3 models * Configure 4.5 * switch to debug logging * update type checking * Fix mypy * Resolve mypy * Merged main into feature/new-model-testing Approved-by: Karan Desai
329 lines
10 KiB
Python
329 lines
10 KiB
Python
import csv
|
|
import os
|
|
import sys
|
|
from collections import defaultdict
|
|
from datetime import datetime
|
|
|
|
import boto3
|
|
from botocore.config import Config
|
|
from dotenv import load_dotenv
|
|
|
|
# Please see the following confluence page for instructions on loading keys into your
|
|
# `.env`:
|
|
# https://aarete.atlassian.net/wiki/spaces/DoczyAI/pages/1283751954/Using+.env+for+AWS+keys
|
|
load_dotenv()
|
|
|
|
|
|
######################################## UTILS ###################################################
|
|
# def check_s3_bucket_exists(s3_client, bucket_name: str):
|
|
# try:
|
|
# s3_client.head_bucket(Bucket=bucket_name)
|
|
# print(f"{bucket_name} exists")
|
|
# except s3_client.exceptions.ClientError as e:
|
|
# error_code = e.response["Error"]["Code"]
|
|
# if error_code == "404":
|
|
# print(f"{bucket_name} does not exist")
|
|
# raise
|
|
# else:
|
|
# print(f"Error {error_code}")
|
|
# raise
|
|
# except Exception as e:
|
|
# raise
|
|
|
|
# return
|
|
|
|
from botocore.exceptions import NoCredentialsError
|
|
|
|
|
|
def check_s3_bucket_exists(s3_client, bucket_name: str):
|
|
try:
|
|
s3_client.head_bucket(Bucket=bucket_name)
|
|
print(f"{bucket_name} exists")
|
|
except NoCredentialsError:
|
|
print("AWS credentials not found")
|
|
return False
|
|
except s3_client.exceptions.ClientError as e:
|
|
error_code = e.response["Error"]["Code"]
|
|
if error_code == "404":
|
|
print(f"{bucket_name} does not exist")
|
|
raise
|
|
else:
|
|
print(f"Error {error_code}")
|
|
raise
|
|
except Exception as e:
|
|
raise
|
|
return True
|
|
|
|
|
|
######################################## GENERAL SETTINGS ########################################
|
|
TEST = False # True to run test prompt - just for testing model connection
|
|
TODAY = datetime.now().strftime("%Y%m%d")
|
|
|
|
MAX_CONTEXT_LENGTH = 100000
|
|
|
|
######################################## TRACKING ########################################
|
|
|
|
UPLOAD_FREQUENCY = 10
|
|
|
|
|
|
######################################## SYSTEM ARGUMENTS ########################################
|
|
# read_mode, run_mode, input_dir, output_dir, consolidated_output_dir
|
|
|
|
|
|
def get_value_from_env(key: str, default_value: str = None) -> str:
|
|
value = os.getenv(key, default_value)
|
|
|
|
if value is not None:
|
|
return value
|
|
|
|
raise KeyError(f"{key} not found in .env file")
|
|
|
|
|
|
def get_arg_value(arg_name, default):
|
|
args = [item for item in sys.argv if arg_name == item.split("=")[0]]
|
|
if len(args) == 0:
|
|
return default
|
|
elif len(args) == 1:
|
|
return args[0].split("=")[1]
|
|
else:
|
|
raise Exception(f"Too many argument values: {args}")
|
|
|
|
|
|
# Run config args
|
|
RUN_MODE = get_arg_value("run_mode", "ec2") # Valid: local or ec2
|
|
READ_MODE = get_arg_value("read_mode", "s3") # Valid: local or s3
|
|
BATCH_ID = get_arg_value("batch_id", "test_batch")
|
|
|
|
# Per-file logging settings (for debugging)
|
|
ENABLE_PER_FILE_LOGGING = get_arg_value("enable_per_file_logging", "False") == "True"
|
|
PER_FILE_LOG_LEVEL = get_arg_value(
|
|
"per_file_log_level", "DEBUG"
|
|
).upper() # DEBUG, INFO, WARNING, ERROR
|
|
PER_FILE_LOG_DIR = get_arg_value("per_file_log_dir", "logs")
|
|
|
|
ENABLE_RUNTIME_ROTATION = False # Set to True to enable multi-runtime failover
|
|
|
|
# Field args
|
|
FIELDS = get_arg_value("fields", "all") # Valid: one_to_one, one_to_n, all
|
|
FIELD_JSON_PATH = "src/prompts/investment_prompts.json"
|
|
|
|
# Client-specific args
|
|
STATE = [s.strip() for s in get_arg_value("state", "None").upper().split("-") if s]
|
|
CLIENT = [s.strip() for s in get_arg_value("client", "None").split("-") if s]
|
|
|
|
# File IO args
|
|
LOCAL_PATH = get_arg_value(
|
|
"input_dir", "data/test"
|
|
) # Valid: any valid path that contains txt files
|
|
S3_PREFIX = get_arg_value("s3_prefix", "")
|
|
S3_BUCKET = get_arg_value("s3_bucket", "")
|
|
OUTPUT_DIRECTORY = get_arg_value(
|
|
"output_dir", "output_individual"
|
|
) # Valid: any valid directory
|
|
TRACKING_FILE = get_arg_value("tracking_file", f"{BATCH_ID}_tracking.csv")
|
|
WRITE_TO_S3 = get_arg_value("write_to_s3", "True") == "True"
|
|
S3_OUTPUT_BUCKET = get_arg_value("s3_output_bucket", default="doczy-output")
|
|
|
|
# Filtering args
|
|
FILTER_ALREADY_PROCESSED = (
|
|
get_arg_value("filter", False) == "True"
|
|
) # Valid: True or False
|
|
|
|
# Multithreading args
|
|
MAX_WORKERS = int(get_arg_value("max_workers", 2)) # Valid: any integer
|
|
|
|
# default false for the regex based IRS chunking
|
|
IRS_REGEX_CHECK = get_arg_value("irs_regex_check", False) == "True"
|
|
|
|
# QCQA args - only need to be set when running qa_qc.py INDEPENDENTLY, not needed for at-scale runs
|
|
AC_DF = get_arg_value("ac_df", None)
|
|
B_DF = get_arg_value("b_df", None)
|
|
DF_READ_MODE = get_arg_value("df_read_mode", "s3")
|
|
|
|
|
|
######################################## FILE LOG ########################################
|
|
|
|
FILE_LOG_FIELDS = [
|
|
"filename",
|
|
"process_type",
|
|
"start_time",
|
|
"end_time",
|
|
"processing_time_seconds",
|
|
"delta_time_hours",
|
|
"status",
|
|
"error",
|
|
"memory_usage_mb",
|
|
"cpu_percent",
|
|
]
|
|
|
|
FILE_LOG_NAME = f"{BATCH_ID}_filelog.csv"
|
|
|
|
if not os.path.exists(FILE_LOG_NAME):
|
|
with open(FILE_LOG_NAME, "w", newline="", encoding="utf-8") as f:
|
|
writer = csv.DictWriter(f, fieldnames=FILE_LOG_FIELDS)
|
|
writer.writeheader()
|
|
|
|
|
|
######################################## CLIENT SETTINGS ########################################
|
|
CLIENT_NAME = get_arg_value("client_name", "Test-Client")
|
|
PROV_TYPE = None
|
|
|
|
######################################## I/O SETTINGS ########################################
|
|
|
|
CONSOLIDATED_OUTPUT_DIRECTORY = "output_consolidated"
|
|
REPORTING_OUTPUT_DIRECTORY = "output_reports"
|
|
|
|
UNPROCESSED_RESULTS_NAME = f"{BATCH_ID}-B-Unprocessed.csv"
|
|
AC_RESULTS_NAME = f"{BATCH_ID}-AC.csv"
|
|
B_RESULTS_NAME = f"{BATCH_ID}-B.csv"
|
|
TABLE_ANALYSIS_NAME = f"{BATCH_ID}-Table-Analysis.csv"
|
|
|
|
|
|
######################################## AWS SETTINGS #########################################
|
|
|
|
# RUN_MODE = "local" # for testing purposes only
|
|
|
|
# Bedrock
|
|
if RUN_MODE == "local":
|
|
AWS_ACCESS_KEY_ID = get_value_from_env("AWS_ACCESS_KEY_ID")
|
|
AWS_SECRET_ACCESS_KEY = get_value_from_env("AWS_SECRET_ACCESS_KEY")
|
|
AWS_SESSION_TOKEN = get_value_from_env("AWS_SESSION_TOKEN")
|
|
|
|
# Bedrock
|
|
config = Config(read_timeout=2000)
|
|
BEDROCK_RUNTIME = boto3.client(
|
|
service_name="bedrock-runtime",
|
|
region_name="us-east-1",
|
|
aws_access_key_id=AWS_ACCESS_KEY_ID,
|
|
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
|
|
aws_session_token=AWS_SESSION_TOKEN,
|
|
config=config,
|
|
)
|
|
|
|
# S3
|
|
S3_CLIENT = boto3.client(
|
|
"s3",
|
|
region_name="us-east-2",
|
|
aws_access_key_id=AWS_ACCESS_KEY_ID,
|
|
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
|
|
aws_session_token=AWS_SESSION_TOKEN,
|
|
)
|
|
|
|
elif RUN_MODE == "ec2":
|
|
# Bedrock
|
|
config = Config(
|
|
read_timeout=4000, max_pool_connections=60
|
|
) # Doubled read_timeout to 4000
|
|
BEDROCK_RUNTIME = boto3.client(
|
|
service_name="bedrock-runtime", region_name="us-east-1", config=config
|
|
)
|
|
|
|
# S3
|
|
S3_CLIENT = boto3.client("s3", region_name="us-east-2")
|
|
|
|
# Check that S3_OUTPUT_BUCKET exists and raise an error if not
|
|
if WRITE_TO_S3:
|
|
check_s3_bucket_exists(S3_CLIENT, S3_OUTPUT_BUCKET)
|
|
|
|
# Claude model IDs
|
|
MODEL_ID_CLAUDE3_HAIKU = "anthropic.claude-3-haiku-20240307-v1:0"
|
|
MODEL_ID_CLAUDE35_SONNET = "anthropic.claude-3-5-sonnet-20240620-v1:0"
|
|
MODEL_ID_CLAUDE2 = "anthropic.claude-instant-v1"
|
|
MODEL_ID_CLAUDE35_SONNET_V2 = "anthropic.claude-3-5-sonnet-20241022-v2:0"
|
|
MODEL_ID_CLAUDE37_SONNET = "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
|
|
MODEL_ID_CLAUDE4_SONNET = "us.anthropic.claude-sonnet-4-20250514-v1:0"
|
|
MODEL_ID_CLAUDE45_SONNET = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
|
|
|
# Model aliases for easy upgrades
|
|
MODEL_ALIASES = {
|
|
"sonnet_37" : MODEL_ID_CLAUDE37_SONNET,
|
|
"sonnet_4" : MODEL_ID_CLAUDE4_SONNET,
|
|
"sonnet_45": MODEL_ID_CLAUDE45_SONNET,
|
|
"sonnet_latest": MODEL_ID_CLAUDE4_SONNET,
|
|
"haiku_latest": MODEL_ID_CLAUDE3_HAIKU,
|
|
"legacy_sonnet": MODEL_ID_CLAUDE35_SONNET,
|
|
}
|
|
|
|
|
|
def resolve_model_id(model_identifier: str) -> str:
|
|
"""
|
|
Resolve a model identifier to its actual model ID.
|
|
Supports both direct model IDs and aliases.
|
|
"""
|
|
return MODEL_ALIASES.get(model_identifier, model_identifier)
|
|
|
|
|
|
######################################## CLAUDE AND PROMPT TRACKING ########################################
|
|
ENABLE_ANSWER_TRACKING = False
|
|
ANSWER_TRACKING_DICT = {}
|
|
|
|
######################################## POSTPROCESSING SETTINGS ########################################
|
|
FUZZY_MATCH_THRESHOLD = 0.8
|
|
FILE_NAME_COLUMN = "Contract Name"
|
|
BLANKS_THRESHOLD = 0.5
|
|
|
|
######################################## COMPLEX FLAG SETTINGS ########################################
|
|
COMPLEX_MAX_WORKERS = 50
|
|
RATE_THRESHOLD = 10
|
|
TABLE_THRESHOLD = 2
|
|
COMPLEX_OUTPUT_PATH = "complex_flag_output"
|
|
COMPLEX_OUTPUT_FILENAME = f"{TODAY}_complex_flag_test.csv"
|
|
|
|
TABLE_ROW_LIMIT = 10
|
|
|
|
######################################## EC2 ########################################
|
|
|
|
MODEL_STATS = defaultdict(
|
|
lambda: {
|
|
"requests": 0,
|
|
"tokens_sent": 0,
|
|
"tokens_received": 0,
|
|
"total_time": 0,
|
|
"claude2_requests": 0,
|
|
"claude3_requests": 0,
|
|
"claude35_requests": 0,
|
|
}
|
|
)
|
|
|
|
GLOBAL_STATS = {
|
|
"total_requests": 0,
|
|
"total_tokens_sent": 0,
|
|
"total_tokens_received": 0,
|
|
"total_time": 0,
|
|
"total_claude2_requests": 0,
|
|
"total_claude3_requests": 0,
|
|
"total_claude35_requests": 0,
|
|
}
|
|
|
|
|
|
def assume_role_and_get_runtime(account_id, role_name):
|
|
sts_client = boto3.client("sts")
|
|
role_arn = f"arn:aws:iam::{account_id}:role/{role_name}"
|
|
assumed_role = sts_client.assume_role(
|
|
RoleArn=role_arn, RoleSessionName=f"UATLambdaSession"
|
|
)
|
|
credentials = assumed_role["Credentials"]
|
|
runtime = boto3.client(
|
|
service_name="bedrock-runtime",
|
|
region_name="us-east-1",
|
|
aws_access_key_id=credentials["AccessKeyId"],
|
|
aws_secret_access_key=credentials["SecretAccessKey"],
|
|
aws_session_token=credentials["SessionToken"],
|
|
config=config,
|
|
)
|
|
return runtime
|
|
|
|
|
|
dev_account_id = "660131068782"
|
|
dev_role_name = "doczy-dev-bedrock-lambda-role"
|
|
|
|
prod_account_id = "211125720533"
|
|
prod_role_name = "doczy-prod-bedrock-lambda-role"
|
|
|
|
############## S3 BUCKETS FOR PDF FILES ##############
|
|
DOCZY_PDF_FILES_BUCKET_NAME = "doczy-pdf-files" # sample bucket name added. we can either use this or create a new one
|
|
DOCZY_PDF_FILES_BUCKET_S3_URL = f"https://{DOCZY_PDF_FILES_BUCKET_NAME}.s3.us-east-2.amazonaws.com/" # sample s3 url added. we can either use this or create a new one
|
|
|
|
############## VISION API SETTINGS ##############
|
|
ENABLE_VISION = get_arg_value("enable_vision", "False") == "True" # False by default
|