doczyai-pipelines/src/config.py

import csv
import os
import re
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", None)  # BATCH_ID is mandatory

# Logging settings
LOG_LEVEL = get_arg_value(
    "log_level", "INFO"
).upper()  # Standard logging level for console output

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

# Specific field or field group to extract (comma-separated)
# Valid: "all", "claim_type", or specific field names like "CLAIM_TYPE_CD,CONTRACT_TITLE"
# "claim_type" is a shortcut for CLAIM_TYPE_CD + AARETE_DERIVED_CLAIM_TYPE_CD + CONTRACT_TITLE
SPECIFIC_FIELDS = get_arg_value("specific_fields", "all")

# Field groups - predefined sets of related fields
FIELD_GROUPS = {
    "claim_type": [
        "CLAIM_TYPE_CD",
        "AARETE_DERIVED_CLAIM_TYPE_CD",
        "CONTRACT_TITLE",
    ],
    "dates": [
        "EFFECTIVE_DT",
        "AARETE_DERIVED_EFFECTIVE_DT",
        "TERMINATION_DT",
        "AARETE_DERIVED_TERMINATION_DT",
    ],
    "provider": [
        "PROV_GROUP_TIN",
        "PROV_GROUP_NPI",
        "PROV_GROUP_NAME_FULL",
        "PROV_OTHER_TIN",
        "PROV_OTHER_NPI",
        "PROV_OTHER_NAME_FULL",
    ],
}


def get_specific_fields_list():
    """Get the list of specific fields to extract based on config."""
    if SPECIFIC_FIELDS == "all":
        return None  # None means extract all fields

    # Check if it's a field group
    if SPECIFIC_FIELDS in FIELD_GROUPS:
        return FIELD_GROUPS[SPECIFIC_FIELDS]

    # Otherwise, treat as comma-separated list of field names
    return [f.strip() for f in SPECIFIC_FIELDS.split(",") if f.strip()]


# Client-specific prompts
CLOVER_PROMPTS_PATH = "src/prompts/clover_prompts.json"
BCBS_PROMISE_PROMPTS_PATH = "src/prompts/bcbs_promise_prompts.json"
CHC_PROMPTS_PATH = "src/prompts/chc_prompts.json"

RAG_RETRIEVAL_QUESTIONS_PATH = "src/prompts/rag_retrieval_questions.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 ########################################

OUTPUT_BASE = "outputs"

# Determine client-specific output directory
# saas → outputs/saas, clover → outputs/clients/clover, etc.
_client_name = CLIENT[0].lower() if CLIENT and CLIENT[0] != "None" else "saas"
if _client_name == "saas":
    CONSOLIDATED_OUTPUT_DIRECTORY = f"{OUTPUT_BASE}/saas"
else:
    CONSOLIDATED_OUTPUT_DIRECTORY = f"{OUTPUT_BASE}/clients/{_client_name}"

# Pipeline-specific output directories
SAAS_OUTPUT_DIRECTORY = f"{OUTPUT_BASE}/saas"
QC_QA_OUTPUT_DIRECTORY = f"{OUTPUT_BASE}/qc_qa"
PARENT_CHILD_OUTPUT_DIRECTORY = f"{OUTPUT_BASE}/parent_child"
INDIVIDUAL_OUTPUT_DIRECTORY = f"{OUTPUT_BASE}/individual"
REPORTING_OUTPUT_DIRECTORY = f"{OUTPUT_BASE}/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, max_pool_connections=60)
    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
# Note: Claude 2 support has been removed - only Claude 3+ models are supported
MODEL_ID_CLAUDE35_SONNET = "anthropic.claude-3-5-sonnet-20240620-v1:0"
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_ID_CLAUDE45_HAIKU = "us.anthropic.claude-haiku-4-5-20251001-v1:0"

# Model aliases for easy upgrades
MODEL_ALIASES = {
    "sonnet_latest": MODEL_ID_CLAUDE4_SONNET,
    "haiku_latest": MODEL_ID_CLAUDE45_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

######################################## TABLE SPLITTING SETTINGS ########################################
# Enable table splitting feature
TABLE_SPLITTING_ENABLED = get_arg_value("table_splitting_enabled", "True") == "True"
# Maximum cells (rows × columns) per table sub-page
MAX_TABLE_CELLS_PER_SUBPAGE = int(get_arg_value("max_table_cells_per_subpage", "40"))
# Minimum rows to trigger splitting
MIN_TABLE_ROWS_TO_SPLIT = int(get_arg_value("min_table_rows_to_split", "20"))
# Table splitting method (currently only "cells" is supported)
TABLE_SPLITTING_METHOD = get_arg_value("table_splitting_method", "cells")

######################################## OUTPUT FILE SPLITTING SETTINGS ########################################
# Maximum rows per split file (files are split to keep filenames together)
MAX_ROWS_PER_SPLIT = int(get_arg_value("max_rows_per_split", "70000"))

######################################## EC2 ########################################


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

############## PARENT CHILD CONFIG ##############
PERFORM_PARENT_CHILD_MAPPING = get_arg_value("perform_pc", "True") == "True"
DOCZY_OUTPUT_FOR_PC = get_arg_value(
    "doczy_output_for_pc", ""
)  # Pass either S3 URI path, or local path

############## DOCUMENT TYPE CLASSIFICATION SETTINGS ##############
PERFORM_DTC = (
    get_arg_value("perform_dtc", "False") == "True"
)  # Enable document type classification pre-filter
DTC_PROMPTS_JSON_PATH = "src/prompts/document_classification_prompts.json"
DTC_MAX_PAGES_TO_CHECK = (
    5  # Number of pages to check per document for contract detection
)
DTC_OUTPUT_FILE = (
    f"{BATCH_ID}_Document_Classification_Report.csv"  # Output CSV file path
)
DTC_JSON_OUTPUT_FOLDER = "dtc_json_results"  # Folder for individual JSON results

############## PRE-DOCZY REPORT SETTINGS ##############
PERFORM_PRE_DOCZY = (
    get_arg_value("perform_pre_doczy", "False") == "True"
)  # Enable Pre-Doczy classification + keyword report

############## POST-DOCZY VALIDATION REPORT SETTINGS ##############
PERFORM_POST_DOCZY = (
    get_arg_value("perform_post_doczy", "False") == "True"
)  # Enable post-Doczy validation report
PERFORM_POST_DOCZY_EXCEL = (
    get_arg_value("perform_post_doczy_excel", "False") == "True"
)  # Enable Excel-only post-Doczy report (no raw text files needed)
POST_DOCZY_JSON_OUTPUT_FOLDER = (
    "post_doczy_json_results"  # Folder for individual JSON results
)
DOCZY_RESULTS_PATH = get_arg_value(
    "doczy_results_path", ""
)  # Path to Doczy extraction output CSV/Excel (local path or S3 URI)

############## DASHBOARD POSTPROCESSING ##############
# Run dashboard postprocessing (comma-separated list format). Default: CC only.
# Set to True via run_dashboard arg when dashboard output is needed.
RUN_DASHBOARD_POSTPROCESSING = get_arg_value("run_dashboard", "False") == "True"

############## QC/QA VALIDATION SETTINGS ##############
# Enable QC/QA validation in pipeline
ENABLE_QC_QA = get_arg_value("enable_qc_qa", "True") == "True"

############## ADD_AARETE_DERIVED_PAYER_NAME SETTINGS ############
# STATE_FLAG controls whether payer names are treated as state-specific.
#
# True  → Preserve state distinctions.
#         Example: "Molina Healthcare of Texas" and
#                  "Molina Healthcare of Utah" remain separate.
#
# False → Collapse regional state suffixes into national brand.
#         Example: "Molina Healthcare of Texas" and
#                  "Molina Healthcare of Utah" both become "Molina Healthcare".
#         Leading state-based brand names (e.g., "Oklahoma Complete Health")
#         are preserved.
STATE_FLAG = True