doczyai-pipelines/fieldExtraction/src/file_processing.py

import os
import pandas as pd
import time
import csv
import json

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

    ac_answers_dict = {}
    log_data = []
    chunk_data = []

    ################## INITIATE PROCESSING ##################
    filename, contract_text = file_object
    print(f"Processing AC for {filename}...")
    print(f"Total contract word count: {len(contract_text.split())}")

    ################## PREPROCESS - SMART CHUNK ##################
    text_dict, exhibit_pages, num_pages, ac_chunks = preprocess.preprocess(
        contract_text, filename, fields="ac"
    )
    print(f"AC Preprocessing Complete - {filename}")

    ################## RUN REGEX FUNCTIONALITY ##################
    if config.IRS_REGEX_CHECK == True:
        irs_answers = ac_funcs.tin_regex(filename, text_dict)
        ac_answers_dict.update(irs_answers)

    ################## RUN CHUNKED PROMPTS ##################

    ## Check `ac_funcs.create_prompt()` vs. `prompts.AC_*_template()`
    all_fields = set(prompts.AC_DICT.keys())

    field_groups = [
        field_group for field_group in keywords.GROUPED_KEYWORD_MAPPINGS.keys()
    ]

    for field_group in field_groups:
        keyword_dict = keywords.GROUPED_KEYWORD_MAPPINGS[field_group]["keywords"]
        fields = keywords.GROUPED_KEYWORD_MAPPINGS[field_group]["fields"]
        questions = {}
        for field in fields:
            questions[field] = prompts.AC_DICT[field]

        # Use chunk if available and different from full context, otherwise use full context
        if (
            field_group in ac_chunks
            and ac_chunks[field_group].strip()
            and ac_chunks[field_group] != contract_text
        ):
            context = ac_chunks[field_group]
            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"

        if len(fields) == 1:
            question = questions[
                fields[0]
            ]  # Take the one question for the single field included
            prompt = prompts.AC_SINGLE_FIELD_TEMPLATE(context, question)
        elif len(fields) > 1:
            prompt = prompts.AC_MULTI_FIELD_TEMPLATE(context, questions)
        else:
            raise ValueError("Field group has no defined fields")

        # TODO: Make sure prompt looks good for groups

        try:
            # field_group_answer_raw = {"test field": "test value"} ## uncomment this and comment below line to test keyword search
            field_group_answer_raw = claude_funcs.invoke_claude(
                prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 8192
            )

            if len(fields) == 1:
                field_group_answer_dict = {field: field_group_answer_raw}
            elif len(fields) > 1:
                field_group_answer_dict = json.loads(field_group_answer_raw)

            # TODO: Validate field names coming out of Claude, check against fields in group, make sure they're all there, make sure there aren't any extras
            for field in fields:
                field_answer = field_group_answer_dict[field]
                ac_answers_dict[field] = field_answer

            log_data.append(
                {
                    "Field": field,
                    "Prompt": f"Question: {questions}\n\nContext: {context}",  # Include full context
                    "Context Type": context_type,
                    "Response": field_answer,
                }
            )

            chunk_data.append(
                {
                    "Contract name": filename,
                    "Field group": field_group,
                    "Field list": fields,
                    "Methodology": ac_chunks[field_group + "_methodology"],
                    "Case sensitivity": ac_chunks[field_group + "_case"],
                    "Page list": ac_chunks[field_group + "_pages"],
                    "Chunk size": len(context),
                    "% Reduction": f"{100*(1-len(context)/len(contract_text)):0.2f}%",
                }
            )
        except Exception as e:
            print(f"Error processing field {field}: {type(e)}, {str(e)}")
            ac_answers_dict[field] = f"Error: {str(e)}"
            log_data.append(
                {
                    "Field": field,
                    "Prompt": f"Question: {questions}\n\nContext: {context}",  # Include full context
                    "Context Type": context_type,
                    "Response": f"Error: {str(e)}",
                }
            )
            chunk_data.append(
                {
                    "Contract name": filename,
                    "Field group": field_group,
                    "Field list": fields,
                    "Methodology": ac_chunks[field_group + "_methodology"],
                    "Case sensitivity": ac_chunks[field_group + "_case"],
                    "Page list": ac_chunks[field_group + "_pages"],
                    "Chunk size": len(context),
                    "% Reduction": f"{100*(1-len(context)/len(contract_text)):0.2f}%",
                }
            )

    ################## AC CONDITIONAL PROMPTS ##################
    # Non-Renewal - Days
    nrd_prompt = prompts.AC_SINGLE_FIELD_TEMPLATE(
        ac_answers_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_answers_dict["NON_RENEWAL_DAYS"] = nrd_answer
    ## TODO - add Notice Provider Name/Address here
    # Clean Up Contract Effective Date
    if (
        utils.is_empty(ac_answers_dict["CONTRACT_EFFECTIVE_DT"])
        and "meridian" in ac_answers_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_answers_dict["CONTRACT_EFFECTIVE_DT"] = date_answer

    ################## RUN FULL CONTEXT PROMPTS ##################
    # Process 'full_context' fields together

    smartly_chunked = []
    for field_groups in keywords.GROUPED_KEYWORD_MAPPINGS.keys():
        for fields in keywords.GROUPED_KEYWORD_MAPPINGS[field_groups]["fields"]:
            smartly_chunked.append(fields)
    full_context_fields = [
        field for field in prompts.AC_DICT if field not in smartly_chunked
    ]
    full_context_fields = [
        field for field in full_context_fields if field not in ac_answers_dict.keys()
    ]

    full_context_questions = {
        field: prompts.AC_DICT.get(field) for field in full_context_fields
    }
    if full_context_questions:
        full_context_prompt = ac_funcs.create_prompt(contract_text[0:min(100000, len(contract_text)-1)], question=full_context_questions)
        try:
            full_context_answers = ac_funcs.get_ac_answer(full_context_prompt, filename, fields=list(full_context_questions.keys()))
            ac_answers_dict.update(full_context_answers)
            
            for field, answer in full_context_answers.items():
                log_data.append({
                    'Field': field,
                    'Prompt': f"Question: {full_context_questions[field]}\n\nContext: {contract_text}",  # Include full context
                    'Context Type': 'Full Context (High Accuracy)',
                    'Response': answer
                })
        except Exception as e:
            print(f"Error processing high accuracy fields: {str(e)}")
            for field in full_context_questions.keys():
                ac_answers_dict[field] = f"Error: {str(e)}"
                log_data.append({
                    'Field': field,
                    'Prompt': f"Question: {full_context_questions[field]}\n\nContext: {contract_text}",  # Include full context
                    'Context Type': 'Full Context (High Accuracy)',
                    'Response': f"Error: {str(e)}"
                })

                
    ################## ENSURE ALL FIELDS ARE PRESENT ##################
    for field in all_fields:
        if field not in ac_answers_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_answers_dict.keys():
        print(key, ac_answers_dict[key])

    ac_df = pd.DataFrame([ac_answers_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)

    chunk_file_path = "chunk_log.csv"
    with open(chunk_file_path, "a", newline="", encoding="utf-8") as chunk_file:
        fieldnames = [
            "Contract name",
            "Field group",
            "Field list",
            "Methodology",
            "Case sensitivity",
            "Page list",
            "Chunk size",
            "% Reduction",
        ]
        writer = csv.DictWriter(chunk_file, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(chunk_data)

    print(f"Chunk log written to {chunk_file_path}")

    return ac_answers_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} ")