doczyai-pipelines/fieldExtraction/src/file_processing.py

import os
import pandas as pd
import time
import csv

import config
import keywords
import preprocess
import preprocessing_funcs
import table_funcs
import bottom_up_funcs
import top_down_funcs
import conditional_funcs
import postprocess
import merge_funcs
import utils
import ac_funcs
import prompts
import valid
import claude_funcs


def run_ac_prompts(file_object):

    ################## INITIATE PROCESSING ##################
    filename, contract_text = file_object
    print(f"Processing AC for {filename}...")
    
    ################## PREPROCESS ##################
    text_dict, exhibit_pages, num_pages, ac_chunks = preprocess.preprocess(
        contract_text, filename, fields="ac"
    )
    print(f"AC Preprocessing Complete - {filename}")

    ################## RUN FULL CONTEXT PROMPTS ##################
    ac_dict = {}

    full_context_fields = [
        field for field in prompts.AC_DICT if field not in keywords.KEYWORD_MAPPINGS
    ]

    full_context_questions = {
        field: prompts.AC_DICT.get(field) for field in full_context_fields
    }

    if full_context_questions:
        full_context_prompt = prompts.AC_MULTI_FIELD_TEMPLATE(contract_text[0 : min(100000, len(contract_text) - 1)], full_context_questions)

        try:
            full_context_answers = ac_funcs.get_ac_answer(
                full_context_prompt,
                filename,
                fields=list(full_context_questions.keys()),
            )
            ac_dict.update(full_context_answers)

        except Exception as e:
            print(f'Exception in AC Full Context : {e}')

    ################## RUN CHUNKED PROMPTS ##################
    all_fields = set(prompts.AC_DICT.keys())
    chunked_fields = [field for field in keywords.KEYWORD_MAPPINGS]

    for field in chunked_fields:
        if field in prompts.AC_DICT:
            question = prompts.AC_DICT[field]
        else:
            print(f"Field {field} not found in prompts. Skipping...")
            continue

        # Use chunk if available and different from full context, otherwise use full context
        if (
            field in ac_chunks
            and ac_chunks[field].strip()
            and ac_chunks[field] != contract_text
        ):
            context = ac_chunks[field]
            context = context[0 : min(100000, len(context)) - 1]
            context_type = "Smart Chunking"
        else:
            context = contract_text[0 : min(100000, len(contract_text) - 1)]
            context_type = "Full Context"

        prompt = prompts.AC_SINGLE_FIELD_TEMPLATE(context, question)

        try:
            field_answer = claude_funcs.invoke_claude(
                prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 8192
            )
            ac_dict[field] = field_answer
        except Exception as e:
            print(f"Error processing field {field}: {str(e)}")

    ################## AC CONDITIONAL PROMPTS ##################
    # Non-Renewal - Days
    nrd_prompt = prompts.AC_SINGLE_FIELD_TEMPLATE(
        ac_dict["NON_RENEWAL_LANGUAGE"], prompts.AC_DICT["NON_RENEWAL_DAYS"]
    )
    nrd_answer = claude_funcs.invoke_claude(
        nrd_prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 8192
    )
    ac_dict["NON_RENEWAL_DAYS"] = nrd_answer

    ## TODO - add Notice Provider Name/Address here

    # Clean Up Contract Effective Date
    if (
        utils.is_empty(ac_dict["CONTRACT_EFFECTIVE_DT"])
        and "meridian" in ac_dict["PAYER_NAME"].lower()
    ):
        date_prompt = prompts.AC_EFFECTIVE_DATE_CLEANUP(
            contract_text[0 : min(100000, len(contract_text)) - 1]
        )
        date_answer = claude_funcs.invoke_claude(
            date_prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 124
        )
        ac_dict["CONTRACT_EFFECTIVE_DT"] = date_answer

    ################## ENSURE ALL FIELDS ARE PRESENT ##################
    for field in all_fields:
        if field not in ac_dict:
            print(f"Field not found in results. {field}")

    ################## CREATE OUTPUT DIRECTORIES ##################
    base_filename = os.path.splitext(filename)[0].strip()
    output_dir = os.path.join(config.OUTPUT_DIRECTORY, base_filename)
    os.makedirs(output_dir, exist_ok=True)

    ################## POSTPROCESS ##################
    for key in ac_dict.keys():
        print(key, ac_dict[key])

    ac_df = pd.DataFrame([ac_dict])
    ac_df = postprocess.ac_postprocess(ac_df, filename, num_pages)

    ################## WRITE TO OUTPUT ##################
    ac_df.to_csv(os.path.join(output_dir, config.AC_RESULTS_NAME), index=False)

    return ac_dict


def run_b_prompts(file_object):
    ################## INITIATE PROCESSING ##################
    filename, contract_text = file_object
    print(f"Processing B for {filename}...")

    ################## PREPROCESS ##################
    text_dict, exhibit_pages, num_pages, ac_chunks = preprocess.preprocess(
        contract_text, filename, fields="b"
    )
    print(f"B Preprocessing Complete - {filename}")

    ################## RUN BOTTOM UP PROMPTS ##################
    bu_results = bottom_up_funcs.run_bottom_up(
        filename, text_dict
    )  # Returns list of dictionaries
    print(f"B Bottom Up Complete - {filename}")

    ################## RUN TOP DOWN PROMPTS ##################
    combined_results = top_down_funcs.run_top_down(
        filename, text_dict, bu_results, exhibit_pages
    )  # Returns list of dictionaries
    print(f"B Top Down Complete - {filename}")

    ################## RUN CONDITIONAL PROMPTS ##################
    conditional_results = conditional_funcs.run_conditional(
        combined_results, text_dict, filename
    )  # List of dictionaries
    print(f"B Conditional Prompts Complete - {filename}")

    ################## CREATE OUTPUT DIRECTORIES ##################
    base_filename = os.path.splitext(filename)[0].strip()
    output_dir = os.path.join(config.OUTPUT_DIRECTORY, base_filename)
    os.makedirs(output_dir, exist_ok=True)

    ################## WRITE UNPROCESSED OUTPUT ##################
    combined_df = pd.DataFrame(conditional_results)
    combined_df.to_csv(
        os.path.join(output_dir, config.UNPROCESSED_RESULTS_NAME), index=False
    )

    ################## RUN POSTPROCESSING ##################
    final_df = postprocess.b_postprocess(filename, combined_df, num_pages)
    print(f"B Postprocessing Complete - {filename} ")

    ################## WRITE FINAL ##################
    final_df.to_csv(os.path.join(output_dir, config.B_RESULTS_NAME), index=False)
    print(f"B Output Complete - {filename} ")