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

import pandas as pd
import logging

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.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
    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 = 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)
    logging.info(f"{datetime_str()} Codes Complete - {filename}")

    final_results = code_funcs.grouper_breakout(results_with_code)

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

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

    if "OUTLIER_TERMS" in one_to_one_results:
        outlier_results = one_to_n_funcs.outlier_terms_breakout(
            one_to_one_results["OUTLIER_TERMS"], filename
        )
        one_to_one_results.update(outlier_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 ##################
        exhibit_level_answers, exhibit_lesser_of, exhibit_header = (
            one_to_n_funcs.get_exhibit_level(
                exhibit_text, exhibit_page, all_exhibit_headers, constants, filename
            )
        )

        ################# GET DYNAMIC ANSWERS ##################
        exhibit_level_answers, reimbursement_level_fields = one_to_n_funcs.get_dynamic(
            exhibit_text, exhibit_header, exhibit_level_answers, constants, filename
        )

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

    ################## Determine LOB Relationship ##################
    one_to_n_results = one_to_n_funcs.get_lob_relationship(
        one_to_n_results, exhibit_dict, filename
    )

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