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

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