doczyai-pipelines/fieldExtraction/src/config.py

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")

######################################## COST LOG ########################################
COST_LOG_NAME = "costlog.csv"

COST_LOG_FIELDS = [
    "Filename",
    "Caller Function",
    "Input Prompt Length",
    "Input Prompt Tokens",
    "Input Cost",
    "Output Response Length",
    "Output Response Tokens",
    "Output Cost",
    "Total Cost",
]

if not os.path.exists(COST_LOG_NAME):
    with open(COST_LOG_NAME, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=COST_LOG_FIELDS)
        writer.writeheader()

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")

# Field args
FIELDS = get_arg_value("fields", "abc")  # Valid: ac, b, abc
FIELD_JSON_PATH = "src/prompts/investment_prompts.json"

# Client-specific args
STATE = get_arg_value("state", "None").upper()

# 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 = ""
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=2000, max_pool_connections=60)
    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"

######################################## 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"


######################################## 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"