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

import csv
import json
import os

import pandas as pd
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.investment_postprocessing_funcs as investment_postprocessing_funcs
import src.investment.preprocess as preprocess
import src.investment.smart_chunking_funcs as smart_chunking_funcs
import src.utils.io_utils as io_utils
import src.utils.string_utils as string_utils
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.investment.tin_npi_funcs import (merge_provider_info,
                                          reimbursement_tin_npi)
from src.prompts.investment_prompts import Field, FieldSet
from src.utils.string_utils import datetime_str


################## MERGE ##################
def merge_one_to_one_into_one_to_n(one_to_n_results, one_to_one_results):
    """
    Merges one_to_one_results into one_to_n_results.

    Args:
        one_to_n_results (pd.DataFrame): DataFrame with N rows and M columns, including FILE_NAME.
        one_to_one_results (dict): Dictionary with several keys, including FILE_NAME.

    Returns:
        pd.DataFrame: Updated one_to_n_results with merged values from one_to_one_results.
    """
    # Iterate over all key-value pairs in one_to_one_results
    for k, v in one_to_one_results.items():
        if k not in one_to_n_results.columns:
            # If k is not a column in one_to_n_results, create it and populate all rows with v
            one_to_n_results[k] = v
        else:
            # If k is a column in one_to_n_results, replace empty values with v
            one_to_n_results[k] = one_to_n_results[k].apply(
                lambda row_val: v if string_utils.is_empty(row_val) else row_val
            )
    return one_to_n_results


def process_file(file_object, all_dataset, 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, filename)
    exhibit_pages, exhibit_chunk_mapping = preprocess.one_to_n_exhibit_chunking(text_dict, 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, text_dict, exhibit_pages, exhibit_chunk_mapping, all_dataset) # Return df
        one_to_n_results['FILE_NAME'] = filename
        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 ##################
    final_results = merge_one_to_one_into_one_to_n(one_to_n_results, one_to_one_results)

    ################## POSTPROCESS ##################
    final_df = postprocess.postprocess(final_results)
    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 REGEX ##################
    regex_answers_dict = one_to_one_funcs.run_provider_info_fields(contract_text, text_dict, filename)
    one_to_one_results = {}
    one_to_one_results['PROV_INFO_JSON'] = json.dumps(regex_answers_dict)
    one_to_one_results = merge_provider_info(one_to_one_results, regex_answers_dict)

    # What if the TIN/NPI regexes don't find anything?  Add PROV_GROUP_NAME_FULL to full context
    if one_to_one_results.get("PROV_OTHER_NAME_FULL") == "NO_IDENTIFIERS_FOUND": # This is how we know the regex didn't find anything
        other_provider_name_field = Field.from_values(
            field_name="PROV_GROUP_NAME_FULL",
            relationship="one_to_one",
            field_type="full_context",
            prompt="What is the name of the group provider associated with this contract? Return the full name as it appears in the document."
        )
        one_to_one_fields.add_field(other_provider_name_field)
        one_to_one_results["PROV_OTHER_NAME_FULL"] = "UNKNOWN"
    
    ################## 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, text_dict, exhibit_pages, exhibit_chunk_mapping, all_dataset):

    ################## RUN PROMPTS ##################
    one_to_n_results = []
    for exhibit_page in exhibit_pages:
        exhibit_chunk = string_utils.get_exhibit_chunk(text_dict, exhibit_chunk_mapping, exhibit_page)
        if string_utils.contains_reimbursement(exhibit_chunk):
            
            ################## INITIALIZE FIELDS ##################
            reimbursement_level_fields = FieldSet(relationship="one_to_n", field_type="reimbursement_level", file_path=FIELD_JSON_PATH)

            ################## GET EXHIBIT HEADER ##################
            exhibit_header = one_to_n_funcs.get_exhibit_header(exhibit_chunk, filename)

            ################## GET REIMBURSEMENT TIN/NPI ##################
            tin_npi_answers, reimbursement_level_fields = reimbursement_tin_npi(exhibit_chunk, reimbursement_level_fields, filename)

            ################## GET EXHIBIT-LEVEL ANSWERS ##################
            exhibit_level_answers = one_to_n_funcs.get_exhibit_level_answers(exhibit_chunk, filename)
            exhibit_level_answers['EXHIBIT_PAGE'] = exhibit_page
            exhibit_level_answers['EXHIBIT_TITLE'] = exhibit_header

            ################## GET DYNAMIC-PRIMARY ANSWERS ##################
            dynamic_to_exhibit_level_answers, dynamic_to_reimbursement_level_fields = dynamic_funcs.get_dynamic_answers(exhibit_chunk, 
                                                                                                                        exhibit_header, 
                                                                                                                        filename,
                                                                                                                        FieldSet(relationship="one_to_n", field_type="dynamic", 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-CODE ANSWERS ##################
            code_to_exhibit_level_answers, code_to_reimbursement_level_fields  = dynamic_funcs.get_dynamic_answers(exhibit_chunk,
                                                                                                                   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_chunk, filename, reimbursement_level_fields, all_dataset, exhibit_page) # 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 = investment_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)

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