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

import logging

import pandas as pd
import src.codes.code_funcs as code_funcs
from src.investment import preprocessing_funcs, aarete_derived, dynamic_funcs, one_to_n_funcs, one_to_one_funcs, postprocess, postprocessing_funcs, preprocess, hybrid_smart_chunking_funcs, row_funcs, tin_npi_funcs
from src.utils import io_utils, logging_utils, string_utils
from constants.constants import Constants
from src import config
from src.prompts.fieldset import FieldSet, Field
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:
    # With this, all logging calls in this thread will now route to logs/{filename}.log
    # This includes logging from all called functions (preprocess, one_to_n_funcs, etc.)
    logging_utils.set_current_file(filename)
    logging.info(f"{datetime_str()} Processing {filename}...")

    # Set default values
    dynamic_one_to_one_fields = FieldSet()
    process_one_to_n = config.FIELDS in ['all', 'one_to_n']
    process_one_to_one = config.FIELDS in ['all', 'one_to_one']
    # Initialize default fallback for final results
    final_results = pd.DataFrame([{"FILE_NAME": 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)

    # ONE TO N PROCESSING
    one_to_n_results = pd.DataFrame()  # Initialize empty DataFrame for one_to_n_results
    if process_one_to_n:
        exhibit_chunk_mapping, 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_results = pd.DataFrame([{"FILE_NAME": filename}])  # Initialize here
        
        if string_utils.contains_reimbursement(contract_text):
            one_to_n_results, dynamic_one_to_one_fields, first_reimbursement_page = run_one_to_n_prompts(
                text_dict, exhibit_chunk_mapping, all_exhibit_headers, constants, filename
            )
            if not one_to_n_results.empty:
                one_to_n_results["FILE_NAME"] = filename
                one_to_n_results = postprocessing_funcs.generate_reimb_ids(one_to_n_results)
            logging.info(f"{datetime_str()} One to N Complete - {filename}")
        else:
            first_reimbursement_page = '1'
            logging.info(f"{datetime_str()} No Reimbursement Found, Skipping - {filename}")
        
        final_results = one_to_n_results  # Set as final results
    else:
        first_reimbursement_page = '1'
        logging.info(f"{datetime_str()} Fields not configured for One to N, Skipping - {filename}")

    # ONE TO ONE PROCESSING
    if process_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,
            first_reimbursement_page,
            constants,
        )
        one_to_one_results["FILE_NAME"] = filename
        logging.info(f"{datetime_str()} One to One Complete - {filename}")
        
        # Decide how to handle one_to_one results
        if not one_to_n_results.empty:
            # BOTH processed - merge into one_to_n
            final_results = row_funcs.merge_one_to_one_into_one_to_n(
                one_to_n_results, one_to_one_results, constants
            )
        else:
            # ONLY one_to_one processed - convert dict to DataFrame
            final_results = pd.DataFrame([one_to_one_results])
    else:
        logging.info(f"{datetime_str()} Fields not configured for One to One, Skipping - {filename}")

    # APPLY CODES IF ONE_TO_N WAS PROCESSED
    if not one_to_n_results.empty:
        results_with_code = code_funcs.code_breakout(final_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: str,
    contract_text: str,
    text_dict: dict[str, str],
    top_sheet_dict,
    dynamic_one_to_one_fields: FieldSet,
    first_reimbursement_page: str,
    constants: 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 HYBRID SMART CHUNKED PROMPTS ##################
    # RAG function loads retrieval questions internally and matches with investment_prompts.json
    hybrid_smart_chunked_answers_dict = hybrid_smart_chunking_funcs.run_hybrid_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, text_dict, first_reimbursement_page, constants, filename
    )
    ################## COMBINE ANSWERS ################
    # Prefer HSC answers over full_context when HSC value is not N/A
    for key, value in full_context_answers_dict.items():
        one_to_one_results[key] = value

    for key, value in hybrid_smart_chunked_answers_dict.items():
        if not string_utils.is_empty(value):
            one_to_one_results[key] = value
        elif key not in one_to_one_results:
            # Add HSC's N/A if field doesn't exist yet
            one_to_one_results[key] = value

    ################## ADD AD FIELDS ################
    one_to_one_results = one_to_one_funcs.get_aarete_derived_dates(one_to_one_results, text_dict, filename)
    
    ################## ADD SIGNATURE COUNT ##################
    one_to_one_results = one_to_one_funcs.add_signature_count(one_to_one_results, contract_text)

    ################## 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(text_dict: dict[str, str], 
                         exhibit_chunk_mapping: dict[str, list[str]], 
                         all_exhibit_headers: dict[str, str], 
                         constants: Constants, 
                         filename: str):

    one_to_n_results = []
    first_reimbursement_page = '1'
    previous_exhibit = None

    #################################### PROCESS EACH EXHIBIT SEPARATELY ####################################
    for exhibit_page, exhibit_page_nums in exhibit_chunk_mapping.items():

        exhibit_text_original = "\n".join([text_dict[page_num] for page_num in exhibit_page_nums])
        exhibit_header = all_exhibit_headers.get(exhibit_page, "")

        ############################### Get Exhibit Level ###############################
        exhibit_level_answers, dynamic_reimbursement_fields = (
            one_to_n_funcs.exhibit_level(
                exhibit_text_original,
                exhibit_header,
                exhibit_page,
                constants,
                filename,
            )
        ) # dict, FieldSet
        logging.debug(f"Exhibit Level Answers for {filename}, Page {exhibit_page}: {exhibit_level_answers}")

        # Store original values before inheritance (for tracking in previous_exhibit)
        original_exhibit_page_nums = exhibit_page_nums.copy()
        original_exhibit_text = exhibit_text_original
        original_dynamic_reimbursement_fields = dynamic_reimbursement_fields

        ############################### Check and Combine Exhibit Inheritance ###############################
        # IF dynamic_reimbursement_fields is empty FieldSet, and previous exhibit's dynamic_reimbursement_fields is NOT empty,
        # and the previous exhibit has NO reimbursement rows, then replace dynamic_reimbursement_fields with the previous
        # exhibit's dynamic_reimbursement_field values AND set exhibit_page_nums = exhibit_page_nums for exhibit N-1 + exhibit_page_nums for exhibit N
        should_inherit, exhibit_page_nums, dynamic_reimbursement_fields = (
            one_to_n_funcs.check_and_combine_exhibit_inheritance(
                previous_exhibit,
                exhibit_page_nums,
                exhibit_text_original,
                dynamic_reimbursement_fields,
            )
        )

        all_exhibit_reimbursements, all_exhibit_special_case = [], []
        #################################### PAGE-BY-PAGE REIMBURSEMENT PRIMARY ####################################
        for page_num in exhibit_page_nums:
            reimbursement_level_answers, special_case_answers, first_reimbursement_page = process_page_one_to_n(
                text_dict,
                exhibit_page_nums,
                exhibit_page,
                page_num,
                first_reimbursement_page,
                dynamic_reimbursement_fields,
                filename,
                constants
            )
            all_exhibit_reimbursements += reimbursement_level_answers
            all_exhibit_special_case += special_case_answers

        ################################ Combine Answers ###############################
        all_exhibit_rows = row_funcs.combine_one_to_n(
            exhibit_text_original,
            all_exhibit_reimbursements,
            all_exhibit_special_case,
            exhibit_level_answers,
            filename
        )  # returns list of dicts

        ################################ Mapping and Cleaning ###############################
        all_exhibit_rows = one_to_n_funcs.one_to_n_cleaning(all_exhibit_rows, exhibit_text_original, constants, filename)

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

        ################################ Store Current Exhibit for Next Iteration ###############################
        # Store original (non-combined) values for next exhibit to potentially inherit from
        # Determine if current exhibit had reimbursements
        had_reimbursements = len(all_exhibit_reimbursements) > 0 or len(all_exhibit_special_case) > 0
        
        previous_exhibit = {
            'exhibit_page': exhibit_page,
            'exhibit_page_nums': original_exhibit_page_nums,
            'exhibit_text_original': original_exhibit_text,
            'dynamic_reimbursement_fields': original_dynamic_reimbursement_fields,
            'exhibit_level_answers': exhibit_level_answers.copy(),
            'had_reimbursements': had_reimbursements
        }

    ################## 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, constants
    )
    
    ################## CONVERT TO DF ##################
    one_to_n_df = pd.DataFrame(one_to_n_results)

    return one_to_n_df, dynamic_one_to_one_fields, first_reimbursement_page


def process_page_one_to_n(text_dict: dict[str, str], exhibit_page_nums: list[str], exhibit_page: str, page_num: str, first_reimbursement_page: str, dynamic_reimbursement_fields: FieldSet, filename: str, constants: Constants):

    page_text = text_dict[page_num]
    exhibit_text_simplified = preprocessing_funcs.simplify_exhibit(text_dict, exhibit_page_nums, page_num)
    
    ############################### Reimbursement Primary ###############################
    reimbursement_level_answers = one_to_n_funcs.reimbursement_level(
        page_text,
        constants,
        filename
    )
    
    if not reimbursement_level_answers:
        return [], [], first_reimbursement_page

    # Track the first page with reimbursements
    if first_reimbursement_page == '1':
        first_reimbursement_page = exhibit_page

    logging.debug(f"Reimbursement Primary Answers for {filename}, Page {page_num}: {reimbursement_level_answers}")

    ################################ Carveouts and Special Case ################################
    reimbursement_level_answers, special_case_answers = one_to_n_funcs.carveout_and_special_case(
        reimbursement_level_answers, constants, filename
    ) # Breaks the reimbursement answers into reimbursments (regular lines) and special cases (distributed across exhibit)

    ############################### Dynamic Assignment ###############################
    reimbursement_level_answers = dynamic_funcs.dynamic_assignment(reimbursement_level_answers, dynamic_reimbursement_fields, exhibit_text_simplified, page_num, constants, filename)

    ############################### Lesser of Distribution ###############################
    reimbursement_level_answers = one_to_n_funcs.lesser_of_distribution(reimbursement_level_answers, exhibit_text_simplified, page_num, constants, filename)
    
    ################################ Get Breakouts ###############################
    reimbursement_level_answers, special_case_answers = one_to_n_funcs.breakout(
        reimbursement_level_answers,
        special_case_answers,
        filename,
        constants,
    ) # list[dict[str, str]], list[dict[str, str]]
    
    return reimbursement_level_answers, special_case_answers, first_reimbursement_page