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

import logging

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.logging_utils as logging_utils
import src.utils.string_utils as string_utils
from constants.constants import Constants
from src import config
from src.prompts.fieldset import FieldSet
from src.utils.string_utils import datetime_str


def process_file(file_object, constants: Constants, run_timestamp):
    filename, contract_text = file_object

    # Set per-file logging context
    logging_utils.set_current_file(filename)
    logging.info(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, header, footer = preprocess.find_headers_and_footers(text_dict)
    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
    )
    logging.info(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
        logging.info(f"{datetime_str()} One to N Complete - {filename}")
    else:
        one_to_n_results = pd.DataFrame([{"FILE_NAME": filename}])
        dynamic_one_to_one_fields = FieldSet()
        logging.info(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,
        constants,
    )  # Return dict
    one_to_one_results["FILE_NAME"] = filename
    logging.info(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, constants
    )

    ################## CODES ##################
    results_with_code = code_funcs.code_breakout(merged_results, constants)
    final_results = code_funcs.grouper_breakout(results_with_code)
    logging.info(f"{datetime_str()} Codes Complete - {filename}")

    ################## POSTPROCESS ##################
    final_df = postprocess.postprocess(final_results, constants)
    logging.info(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")

    logging.info(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,
    constants,
):

    ################## 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, constants, 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, constants, 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], constants
    )

    ################## 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):

    one_to_n_results = []
    seen_pairs = set()

    for exhibit_page, exhibit_text in exhibit_dict.items():

        ############################### Initialize Fields ###############################
        exhibit_level_fields = FieldSet(
            config.FIELD_JSON_PATH, field_type="exhibit_level"
        )
        reimbursement_level_fields = FieldSet(
            config.FIELD_JSON_PATH, field_type="reimbursement_level"
        )
        dynamic_primary_fields = FieldSet(
            config.FIELD_JSON_PATH, field_type="dynamic_primary"
        )
        dynamic_reimb_info_fields = FieldSet(
            config.FIELD_JSON_PATH, field_type="dynamic_reimb_info"
        )
        dynamic_code_fields = FieldSet(
            config.FIELD_JSON_PATH, field_type="dynamic_code"
        )
        special_case_primary_fields = FieldSet(
            config.FIELD_JSON_PATH, field_type="special_case_primary"
        )

        ############################### Get Exhibit Header ###############################
        exhibit_header = one_to_n_funcs.get_exhibit_header(
            all_exhibit_headers, exhibit_page
        )

        ############################### Get Exhibit Level ###############################
        exhibit_level_answers, reimbursement_level_fields = (
            one_to_n_funcs.exhibit_level(
                exhibit_text,
                exhibit_header,
                exhibit_page,
                exhibit_level_fields,
                dynamic_primary_fields,
                dynamic_code_fields,
                dynamic_reimb_info_fields,
                reimbursement_level_fields,
                constants,
                filename,
            )
        )

        ############################### Get Reimbursement-Level ###############################
        reimbursement_primary_answers, special_case_primary_answers = (
            one_to_n_funcs.reimbursement_level(
                exhibit_text,
                reimbursement_level_fields,
                special_case_primary_fields,
                exhibit_level_answers["EXHIBIT_LESSER_OF_STATEMENT"],
                seen_pairs,
                exhibit_page,
                filename,
                constants,
            )
        )

        ################################ Get Breakouts ###############################
        reimbursement_level_answers, special_case_breakout_answers = (
            one_to_n_funcs.breakout(
                reimbursement_primary_answers,
                special_case_primary_answers,
                special_case_primary_fields,
                filename,
                constants,
            )
        )

        ################################ Combine Answers ###############################
        all_exhibit_rows = row_funcs.combine_one_to_n(
            exhibit_text,
            reimbursement_level_answers,
            exhibit_level_answers,
            special_case_breakout_answers,
            filename,
        )  # returns list of dicts

        ################################ Crosswalk Fields ################################
        all_exhibit_rows = aarete_derived.get_crosswalk_fields(
            all_exhibit_rows, constants
        )

        ################################ Determine LOB Relationship ################################
        all_exhibit_rows = one_to_n_funcs.get_lob_relationship(
            all_exhibit_rows, exhibit_text, filename
        )

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

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

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

        ################################ Add to Total ###############################
        one_to_n_results += all_exhibit_rows

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

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

    return one_to_n_df, dynamic_one_to_one_fields