doczyai-pipelines/fieldExtraction/src/investment/file_processing.py

import pandas as pd

import src.codes.code_funcs as code_funcs
import src.investment.aarete_derived as aarete_derived
import src.investment.dynamic_funcs as dynamic_funcs
import src.investment.one_to_n_funcs as one_to_n_funcs
import src.investment.one_to_one_funcs as one_to_one_funcs
import src.investment.postprocess as postprocess
import src.investment.postprocessing_funcs as postprocessing_funcs
import src.investment.preprocess as preprocess
import src.investment.row_funcs as row_funcs
import src.investment.tin_npi_funcs as tin_npi_funcs
import src.utils.io_utils as io_utils
import src.utils.string_utils as string_utils
from constants.constants import Constants
from src import config
from src.config import FIELD_JSON_PATH
from src.investment.one_to_n_funcs import combine_one_to_n_answers, reimbursement_level
from src.prompts.investment_prompts import FieldSet
from src.utils.string_utils import datetime_str


def process_file(file_object, constants: Constants, run_timestamp):
    filename, contract_text = file_object
    print(f"{datetime_str()} Processing {filename}...")

    ################## PREPROCESS ##################
    contract_text = preprocess.clean_text(contract_text)
    text_dict, top_sheet_dict = preprocess.split_text(contract_text)
    text_dict = preprocess.clean_tables(
        text_dict, constants.EXHIBIT_HEADER_MARKERS, filename
    )
    exhibit_dict, all_exhibit_headers = preprocess.one_to_n_exhibit_chunking(
        text_dict, constants.EXHIBIT_HEADER_MARKERS, filename
    )
    print(f"{datetime_str()} Preprocessing Complete - {filename}")

    ################## ONE TO N ##################
    if string_utils.contains_reimbursement(contract_text):
        one_to_n_results, dynamic_one_to_one_fields = run_one_to_n_prompts(
            filename, exhibit_dict, all_exhibit_headers, constants
        )  # Return df
        one_to_n_results["FILE_NAME"] = filename
        one_to_n_results = postprocessing_funcs.generate_reimb_ids(
            one_to_n_results
        )  # Add reimb_id
        print(f"{datetime_str()} One to N Complete - {filename}")
    else:
        one_to_n_results = pd.DataFrame([{"FILE_NAME": filename}])
        dynamic_one_to_one_fields = FieldSet()
        print(f"{datetime_str()} No One to N Found, Skipping - {filename}")

    ################## ONE TO ONE ##################
    one_to_one_results = run_one_to_one_prompts(
        filename, contract_text, text_dict, top_sheet_dict, dynamic_one_to_one_fields
    )  # Return dict
    one_to_one_results["FILE_NAME"] = filename
    print(f"{datetime_str()} One to One Complete - {filename}")

    ################## MERGE ##################
    merged_results = row_funcs.merge_one_to_one_into_one_to_n(
        one_to_n_results, one_to_one_results
    )

    ################## CODES ##################
    final_results = code_funcs.code_breakout(merged_results, constants)

    ################## POSTPROCESS ##################
    final_df = postprocess.postprocess(final_results, constants)
    print(f"{datetime_str()} Postprocessing Complete - {filename}")

    ################## WRITE INDIVIDUAL ##################
    if config.WRITE_TO_S3:
        io_utils.write_s3(final_df, filename, run_timestamp, "individual")
    else:
        io_utils.write_local(final_df, filename, "", "individual")

    print(f"{datetime_str()} Writing Complete - {filename}")
    return final_df


def run_one_to_one_prompts(
    filename, contract_text, text_dict, top_sheet_dict, dynamic_one_to_one_fields
):

    ################## INITIALIZE FIELDS ##################
    one_to_one_fields = FieldSet(
        relationship="one_to_one", file_path=config.FIELD_JSON_PATH
    ).combine(dynamic_one_to_one_fields)

    ################## RUN PROVIDER INFO ##################
    one_to_one_results, one_to_one_fields = tin_npi_funcs.run_provider_info_fields(
        contract_text, one_to_one_fields, text_dict, filename
    )

    ################## RUN SMART CHUNKED PROMPTS ##################
    smart_chunked_answers_dict = one_to_one_funcs.run_smart_chunked_fields(
        one_to_one_fields, contract_text, filename, text_dict
    )

    ################## RUN FULL CONTEXT PROMPTS ##################
    full_context_answers_dict = one_to_one_funcs.run_full_context_fields(
        one_to_one_fields, contract_text, filename
    )

    ################## COMBINE ANSWERS ################
    one_to_one_results.update(smart_chunked_answers_dict)
    one_to_one_results.update(full_context_answers_dict)

    ################## Crosswalk Fields ##################
    one_to_one_results = aarete_derived.get_crosswalk_fields([one_to_one_results])

    ################## Fill NA Mapping ##################
    one_to_one_results = aarete_derived.fill_na_mapping(one_to_one_results)

    return one_to_one_results[0]


def run_one_to_n_prompts(filename, exhibit_dict, all_exhibit_headers, constants):

    ################## RUN PROMPTS ##################
    one_to_n_results = []
    seen_pairs = (
        set()
    )  # Accumulator for tracking duplicate pairs of (`SERVICE_TERM`, `REIMB_TERM`)

    for exhibit_page in exhibit_dict.keys():

        exhibit_text = exhibit_dict[exhibit_page]

        ################# GET EXHIBIT-LEVEL ANSWERS ##################
        if exhibit_page in all_exhibit_headers.keys():
            exhibit_header = all_exhibit_headers[exhibit_page]
        elif (
            ".".join(exhibit_page.split(".")[:-1]) in all_exhibit_headers.keys()
        ):  # Take the parent header if it exists
            exhibit_header = all_exhibit_headers[".".join(exhibit_page.split(".")[:-1])]
        else:
            exhibit_header = "No Header Found"

        exhibit_level_answers = one_to_n_funcs.get_exhibit_level_answers(
            exhibit_text, filename
        )
        exhibit_level_answers["EXHIBIT_PAGE"] = exhibit_page
        exhibit_level_answers["EXHIBIT_TITLE"] = exhibit_header

        # Extract exhibit lesser-of statement for passing to reimbursement_level
        exhibit_lesser_of = exhibit_level_answers.get(
            "EXHIBIT_LESSER_OF_STATEMENT", "N/A"
        )

        ################## INITIALIZE FIELDS ##################
        reimbursement_level_fields = FieldSet(
            relationship="one_to_n",
            field_type="reimbursement_level",
            file_path=FIELD_JSON_PATH,
        )

        ################## GET DYNAMIC-PRIMARY ANSWERS ##################
        dynamic_to_exhibit_level_answers, dynamic_to_reimbursement_level_fields = (
            dynamic_funcs.get_dynamic_primary(
                exhibit_text,
                exhibit_header,
                filename,
                FieldSet(
                    relationship="one_to_n",
                    field_type="dynamic_primary",
                    file_path=config.FIELD_JSON_PATH,
                ),
            )
        )
        exhibit_level_answers.update(dynamic_to_exhibit_level_answers)
        reimbursement_level_fields.combine(
            dynamic_to_reimbursement_level_fields, inplace=True
        )

        ################## GET DYNAMIC-REIMB-INFO ANSWERS ##################
        (
            reimb_info_to_exhibit_level_answers,
            reimb_info_to_reimbursement_level_fields,
        ) = dynamic_funcs.get_dynamic_answers(
            exhibit_text,
            exhibit_header,
            filename,
            FieldSet(
                relationship="one_to_n",
                field_type="dynamic_reimb_info",
                file_path=config.FIELD_JSON_PATH,
            ),
        )
        exhibit_level_answers.update(reimb_info_to_exhibit_level_answers)
        reimbursement_level_fields.combine(
            reimb_info_to_reimbursement_level_fields, inplace=True
        )

        ################## GET DYNAMIC-CODE ANSWERS ##################
        code_to_exhibit_level_answers, code_to_reimbursement_level_fields = (
            dynamic_funcs.get_dynamic_answers(
                exhibit_text,
                exhibit_header,
                filename,
                FieldSet(
                    relationship="one_to_n",
                    field_type="dynamic_code",
                    file_path=config.FIELD_JSON_PATH,
                ),
            )
        )
        exhibit_level_answers.update(code_to_exhibit_level_answers)
        reimbursement_level_fields.combine(
            code_to_reimbursement_level_fields, inplace=True
        )

        ################## GET REIMBURSEMENT-LEVEL ANSWERS (INCLUDING DYNAMIC) ##################
        reimbursement_level_answers = reimbursement_level(
            exhibit_text,
            filename,
            reimbursement_level_fields,
            exhibit_page,
            seen_pairs,
            exhibit_lesser_of=exhibit_lesser_of,
            constants=constants,
        )  # Return list of dictionaries

        ################# COMBINE ANSWERS ##################
        full_answer_dict = combine_one_to_n_answers(
            exhibit_level_answers, reimbursement_level_answers, tin_npi_answers={}
        )
        one_to_n_results += full_answer_dict

    ################## Crosswalk Fields ##################
    one_to_n_results = aarete_derived.get_crosswalk_fields(one_to_n_results)

    ################## Fill NA Mapping ##################
    one_to_n_results = aarete_derived.fill_na_mapping(one_to_n_results)

    ################## Update LOB for Duals ##################
    one_to_n_results = postprocessing_funcs.update_lob_for_duals(one_to_n_results)

    ################## Add N/A Dynamic or Exhibit to One-to-One ##################
    dynamic_one_to_one_fields = dynamic_funcs.get_dynamic_one_to_one_fields(
        one_to_n_results
    )

    ################### Split REIMB_DATES ##################
    one_to_n_results = one_to_n_funcs.split_reimb_dates(one_to_n_results, filename)

    ################## CONVERT TO DF ##################
    one_to_n_df = pd.DataFrame(one_to_n_results)

    return one_to_n_df, dynamic_one_to_one_fields