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.hybrid_smart_chunking_funcs as hybrid_smart_chunking_funcs
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:
    # 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)
    text_dict = preprocess.clean_tables(
        text_dict, constants.EXHIBIT_HEADER_MARKERS, filename
    )

    # ONE TO N PROCESSING
    one_to_n_results = pd.DataFrame()  # Initialize empty DataFrame for one_to_n_results
    if process_one_to_n:
        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_results = pd.DataFrame([{"FILE_NAME": filename}])  # Initialize here
        
        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
            )
            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:
            logging.info(f"{datetime_str()} No Reimbursement Found, Skipping - {filename}")
        
        final_results = one_to_n_results  # Set as final results
    else:
        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,
            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,
    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 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, 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)
    
    ################## 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,
            )
        )
        logging.debug(f"Exhibit Level Answers for {filename}, Page {exhibit_page}: {exhibit_level_answers}")

        ############################### 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,
            )
        )
        logging.debug(f"Reimbursement Primary Answers for {filename}, Page {exhibit_page}: {reimbursement_primary_answers}")

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

        ################################ Mapping and Cleaning ###############################
        all_exhibit_rows = one_to_n_funcs.one_to_n_cleaning(all_exhibit_rows, exhibit_text, constants, 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