doczyai-pipelines/fieldExtraction/src/bottom_up_funcs.py

import dict_operations
import postprocessing_funcs
import prompts
import config
import claude_funcs
import utils

import difflib

def run_bottom_up(filename, text_dict):
    """
    Processes the text of a document using a two-tiered Bottom Up approach to extract key financial and operational information.

    This function first runs BOTTOM_UP_PRIMARY and BOTTOM_UP_SECONDARY prompts, then processes these initial results to further 
    refine and structure them into a dictionary form. 

    Parameters:
        filename (str): The name of the file being processed, used to tag output data.
        text_dict (dict): A dictionary of text keyed by page numbers that contains the content to be analyzed.

    Returns:
        list of dict: A list of dictionaries with each dictionary containing refined and structured information
        from both primary and secondary Bottom Up analyses, all tagged with the filename.
    """

    # Bottom Up Primary - SERVICE, REIMBURSEMENT_FLAT_FEE, REIMBURSEMENT_RATE, METHODOLOGY
    answer_strings = run_bottom_up_primary(text_dict, 8000)
    answer_dicts = dict_operations.primary_string_to_dict(answer_strings, filename)
    answer_dicts_filtered = postprocessing_funcs.filter_service_column(answer_dicts)
    
    # Bottom Up Secondary
    results_dicts = run_bottom_up_secondary(answer_dicts_filtered, text_dict, 8000)

    for d in results_dicts:
        d['Filename'] = filename
    return results_dicts # List of dictionaries


def run_bottom_up_primary(text_dict, tokens):
    """
    Executes the primary Bottom Up processing for pages that contain reimbursement terms.

    This function scans through a dictionary of page texts, identifies pages containing reimbursement terms,
    and processes those pages using Claude to extract relevant information.

    Parameters:
        text_dict (dict): A dictionary containing text content of documents keyed by page numbers.
        tokens (int): The maximum number of tokens to use in the language model invocation for generating the response.

    Returns:
        dict: A dictionary where each key is a page number and the value is the response from the language model.
    """
    #chunk_dict = preprocess.chunk_text(text_dict)
    answer_dict = {}
    #for page_number in chunk_dict.keys():
        #if '%' in chunk_dict[page_number] or '$' in chunk_dict[page_number]:
            #prompt = prompts.BOTTOM_UP_PRIMARY(chunk_dict[page_number], config.CLIENT_NAME)
    for page_number in text_dict.keys():
        if utils.contains_reimbursement(text_dict, page_number):
            # Run Primary
            prompt = prompts.BOTTOM_UP_PRIMARY(text_dict[page_number], config.CLIENT_NAME)
            answer = claude_funcs.invoke_claude_3(prompt, max_tokens=tokens)
            answer_dict[page_number] = answer
    return answer_dict


def run_bottom_up_secondary(answer_dicts, text_dict, tokens):
    """
    Executes the secondary Bottom Up processing phase on the results obtained from the primary Bottom Up analysis.

    This function enhances the primary results with additional analyses based on configured conditions. Invokes 
    Claude 3 with tailored prompts to generate structured information that complements the initial results.

    Each piece of data processed possibly undergoes several rounds of checks and transformations, ensuring detailed and 
    comprehensive output.

    Parameters:
        answer_dicts (list of dict): Initial processed data from the primary Bottom Up analysis.
        text_dict (dict): Dictionary containing the original document text keyed by page numbers.
        tokens (int): Token limit for language model invocations.

    Returns:
        list of dict: A list of dictionaries containing enriched and finalized structured data from both the primary and 
        secondary analyses.
    """
    # New rows for lesser of
    temp_dicts = []
    for d in answer_dicts:     
        # Bottom Up Lesser
        if config.RUN_LESSER:
            lesser_object = run_bottom_up_lesser(d.copy(), text_dict.copy(), tokens)
            temp_dicts.append(lesser_object[1]) # Add original object from d
            if lesser_object[0]: temp_dicts.append(lesser_object[0]) # If lesser, add lesser object
        else:
            temp_dicts.append(d)
    
    # Add additional fields to each row
    final_dicts = []
    for d in temp_dicts:
        if d is not None:
            page_num = d['page_num']
            
            # Bottom Up Methodology
            if config.RUN_METHODOLOGY:
                prompt = prompts.BOTTOM_UP_METHODOLOGY(d)
                methodology_answer = claude_funcs.invoke_claude_3(prompt, max_tokens=100)
                d['REIMBURSEMENT_METHODOLOGY'] = methodology_answer
                
            # # Bottom Up FS
            if config.RUN_FS:
                prompt = prompts.BOTTOM_UP_FS(d)
                fs_answer = claude_funcs.invoke_claude_3(prompt, model_id=config.MODEL_ID_CLAUDE3_HAIKU, max_tokens=tokens)
                fs_dict = dict_operations.secondary_string_to_dict(fs_answer)
                d.update(fs_dict)

            # # Bottom Up Exception/Escalator
            if config.RUN_EXCEPTION:
                prompt = prompts.BOTTOM_UP_EXCEPT_ESC(d, text_dict[page_num])
                exc_answer = claude_funcs.invoke_claude_3(prompt, model_id=config.MODEL_ID_CLAUDE3_HAIKU, max_tokens=tokens)
                exc_dict = dict_operations.secondary_string_to_dict(exc_answer)
                d.update(exc_dict)

            # # Bottom Up Codes
            if config.RUN_CODES:
                prompt = prompts.BOTTOM_UP_CODES(d, text_dict[page_num])
                codes_answer = claude_funcs.invoke_claude_3(prompt, max_tokens=tokens)
                codes_dict = dict_operations.secondary_string_to_dict(codes_answer)
                d.update(codes_dict)

            final_dicts.append(d)
        
    return final_dicts


def run_bottom_up_lesser(d, text_dict, tokens=4000):
    """
    Processes specific clauses or conditions within the document text, such as 'lesser of' or 'greater of',
    using a language model to analyze and extract relevant data from the identified text segment.

    This function takes a dictionary representing initial results and a text dictionary, uses the page number from the initial
    results to locate the relevant text, and generates a prompt for language model analysis. Based on the response, it
    determines whether 'lesser of' or 'greater of' language is indicated, updates the results, and returns a new dictionary
    for the updated data or retains the original if no such language is present.

    Parameters:
        d (dict): A dictionary containing initial processing results from previous analyses, including a 'page_num' key.
        text_dict (dict): A dictionary of text keyed by page numbers that contains the content to be analyzed.
        tokens (int): Token limit for language model invocations, default set to 8000.

    Returns:
        tuple: A tuple containing two dictionaries, the first for any identified 'lesser of' or 'greater of' data and the second
               containing the original or updated data depending on the presence of such language.
    """

    prompt = prompts.BOTTOM_UP_LESSER(d, text_dict[d['page_num']])
    lesser_of_answer = claude_funcs.invoke_claude_3(prompt, max_tokens=tokens)
    lesser_of_dict_list = dict_operations.primary_string_to_dict({d['page_num'] : lesser_of_answer}, d['Filename'])
    # print(lesser_of_dict_list)
    # If lesser of is Y
    if len(lesser_of_dict_list) > 1:
        lesser_of_dict = get_lesser_of_dict(lesser_of_dict_list, d)
    else:
        d['LESSER_OF_LANGUAGE_IND'] = 'N' 
        d['GREATER_OF_LANGUAGE_IND'] = 'N'
        return ({}, d)
    
    lesser_of_dict['page_num'] = d['page_num']
    if lesser_of_dict['LESSER_OF_LANGUAGE_IND'] == 'Y' or lesser_of_dict['GREATER_OF_LANGUAGE_IND'] == 'Y':
        lesser_of_dict['SERVICE'] = d['SERVICE']
        return (lesser_of_dict, d)
    else:
        d['LESSER_OF_LANGUAGE_IND'] = 'N' 
        d['GREATER_OF_LANGUAGE_IND'] = 'N'
        return ({}, d)



def get_least_similar(dict_list, original_dict, field):
    # print(f"Running similarity for {field}")
    
    min_similarity = float('inf')
    least_similar_dict = None
    for dictionary in dict_list:
        # Extract the methodology text
        field_text = dictionary[field]
        # print(f"Field text: {field_text}")
        
        # Compute similarity using difflib
        similarity = difflib.SequenceMatcher(None, str(field_text), str(original_dict[field])).ratio()
        # print(f"Simlarity: {similarity}")
        
        # If the similarity is less than the current minimum, update the minimum and the corresponding dictionary
        if similarity < min_similarity:
            min_similarity = similarity
            least_similar_dict = dictionary
    
    return least_similar_dict

def get_lesser_of_dict(dict_list, original_dict):

    unique_rates = list({d['REIMBURSEMENT_RATE'] for d in dict_list})
    unique_fees = list({d['REIMBURSEMENT_FLAT_FEE'] for d in dict_list})

    # If the rates are different, return the lesser of dict with a different rate than original 
    if len(unique_rates) > 1:
        least_similar_dict = get_least_similar(dict_list, original_dict, 'REIMBURSEMENT_RATE')
    elif len(unique_fees) > 1:
        least_similar_dict = get_least_similar(dict_list, original_dict, 'REIMBURSEMENT_FLAT_FEE')
    else:
        least_similar_dict = get_least_similar(dict_list, original_dict, 'FULL_METHODOLOGY')

    return least_similar_dict