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

import hashlib
import re
from typing import Callable
import logging

import src.constants.investment_values as investment_values
import src.investment.code_funcs as code_funcs
import src.prompts.investment_prompts as investment_prompts
import src.utils.llm_utils as llm_utils
import src.utils.string_utils as string_utils
from src import config
from src.constants.special_case_configs import SPECIAL_CASE_CONFIGS
from src.enums.delimiters import Delimiter
from src.prompts.investment_prompts import Field, FieldSet, EXHIBIT_LEVEL
from src.config import FIELD_JSON_PATH

import yaml # This is for experimental YAML parsing

METHODOLOGY_BREAKOUT_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="methodology_breakout"
).print_prompt_dict()
FS_BREAKOUT_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="fee_schedule_breakout"
).print_prompt_dict()


def get_exhibit_level_answers(exhibit_chunk, filename):
    exhibit_level_fields = FieldSet(relationship="one_to_n", field_type="exhibit_level", file_path=FIELD_JSON_PATH)

    if not exhibit_level_fields.contains_fields():
        return {}

    prompt = EXHIBIT_LEVEL(exhibit_chunk, exhibit_level_fields)
    llm_answer_raw = llm_utils.invoke_claude(
        prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, max_tokens=8192
    )
    llm_answer_final = string_utils.universal_json_load(llm_answer_raw)
    return llm_answer_final


def get_exhibit_header(exhibit_chunk, filename):
    """
    Extracts the exhibit header using an LLM prompt.

    Args:
        exhibit_chunk (str): A chunk of exhibit text to analyze.
        filename (str): The name of the file being processed.

    Returns:
        str: The extracted exhibit header.
    """
    prompt = investment_prompts.EXHIBIT_HEADER(exhibit_chunk[0:400])
    llm_answer_raw = llm_utils.invoke_claude(
        prompt, config.MODEL_ID_CLAUDE35_SONNET, filename
    )
    llm_answer_final = string_utils.extract_text_from_delimiters(
        llm_answer_raw, Delimiter.PIPE
    )
    return llm_answer_final

def identify_reimbursement_exhibits(exhibit_headers: dict, filename: str) -> list[str]:
    """Use LLM to identify which exhibits likely contain reimbursement terms.

    Args:
        exhibit_headers (dict): Dictionary of exhibit headers to analyze, keyed by 
            string page number, valued by the exhibit header text.
        filename (str): The name of the file being processed.

    Returns:
        list[str]: List of page numbers (as strings) that likely contain reimbursement terms.
    """
    prompt = investment_prompts.IDENTIFY_REIMBURSEMENT_EXHIBITS_PROMPT(exhibit_headers, yaml_output=True) # Trying YAML output
    logging.debug(f"""Identifying reimbursement exhibits for {filename} with prompt:
{prompt}""")
    llm_answer_raw = llm_utils.invoke_claude(
        prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 2000
    )
    yaml_content = string_utils.extract_text_from_delimiters(
        llm_answer_raw, Delimiter.TRIPLE_BACKTICK)
    
    # Check if the first line is just `yaml` and if so remove it
    lines = yaml_content.strip().split('\n')
    if lines[0].strip().lower() == 'yaml':
        yaml_content = '\n'.join(lines[1:])
    
    logging.debug(f"""LLM raw output for {filename}:
{llm_answer_raw}""")
    # print("Extracted YAML content:\n", yaml_content) # Debugging line
    logging.debug(f"""YAML content for {filename}:
{yaml_content}""")
    try: 
        # Parse the YAML output into a dictionary
        exhibits_dict = yaml.safe_load(yaml_content)

        # Return both EXPLICIT and IMPLICIT exhibits (exclude only 'NO')
        return [str(page) for page, value in exhibits_dict.items() 
                if value in ['EXPLICIT', 'IMPLICIT']]
    except yaml.YAMLError as e:
        logging.error(f"Error parsing YAML: {e}")
        logging.error(f"Raw LLM output: {llm_answer_raw}")
        return list(exhibit_headers.keys())  # Return all keys of the exhibit headers if parsing fails

def get_reimbursement_primary(reimbursement_level_fields, exhibit_text, filename):
    field_prompts = reimbursement_level_fields.get_prompt_dict()
    prompt = investment_prompts.REIMBURSEMENT_LEVEL_PRIMARY(exhibit_text, field_prompts)
    logging.debug(f"""Running reimbursement primary prompt for {filename} with prompt:
{prompt}""")
    llm_answer_raw = llm_utils.invoke_claude(
        prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, max_tokens=25000
    )
    logging.debug(f"""LLM raw output for {filename}:
{llm_answer_raw}""")
    
    # Check for the special "no results" case
    if "NO_REIMBURSEMENT_TERMS_FOUND" in llm_answer_raw:
        return []  # Return empty list if no reimbursement terms found
    
    try:
        llm_answer_final = string_utils.universal_json_load(llm_answer_raw)
    except ValueError as e:
        logging.error(f"Error parsing LLM response: {e}")
        logging.error(f"Raw LLM output: {llm_answer_raw}")
        raise
    return llm_answer_final


def get_special_cases(
    reimbursement_primary_answers: list[dict], filename: str
) -> list[dict]:
    """
    Determines carveout indicators for reimbursement answers using LLM-based classification.

    Args:
        reimbursement_primary_answers (list[dict]): List of dictionaries containing reimbursement details.
        filename (str): The filename used for LLM processing.

    Returns:
        list[dict]: Updated list with carveout indicators and default indicators added.
        Fields Added: CARVEOUT_IND, DEFAULT_IND, CARVEOUT_CD
    """
    special_case_answers = []

    for answer_dict in reimbursement_primary_answers:
        contract_service_cd, contract_reimbursement_method = answer_dict.get(
            "SERVICE_TERM"
        ), answer_dict.get("REIMB_TERM")
        if contract_service_cd and contract_reimbursement_method:
            llm_answer_raw = llm_utils.invoke_claude(
                investment_prompts.SPECIAL_CASE_CHECK(
                    contract_service_cd, contract_reimbursement_method
                ),
                config.MODEL_ID_CLAUDE35_SONNET,
                filename,
            )
            llm_answer_final = string_utils.extract_text_from_delimiters(
                llm_answer_raw, Delimiter.PIPE
            )

            # Standard N/A response
            if llm_answer_final == "N/A":
                answer_dict["CARVEOUT_CD"] = "N/A"
                answer_dict["CARVEOUT_IND"] = "N"
                answer_dict["DEFAULT_IND"] = "N"
            # ALL special cases and carveouts
            elif any(
                [
                    carveout in llm_answer_final
                    for carveout in list(investment_values.VALID_CARVEOUTS.keys())
                    + list(investment_values.VALID_SPECIAL.keys())
                ]
            ):
                result_list = get_carveout_fields(
                    answer_dict, llm_answer_final, filename
                )  # should be a list of dicts
                special_case_answers.extend(result_list)
                continue  # skip the original append because we've just extended the list

            special_case_answers.append(
                answer_dict
            )  # this only happens if no special case was found (i.e. the "N/A" case)

    return special_case_answers


def process_carveout_or_special_case(
    answer_dict: dict,
    carveout_type: str,
    carveout_ind: bool,
    default_ind: bool,
    carveout_fields: list[dict],
    questions: dict,
    fs_breakout_questions: dict,
    filename: str,
    methodology_prompt_template: Callable,
    fee_schedule_prompt_template: Callable,
    term_key: str = None,
) -> None:
    """Process a carveout or special case by setting indicators and running the appropriate LLM prompts.

    Args:
        answer_dict (dict): Previous LLM answer dictionary to update.
        carveout_type (str): String indicating the type of carveout or special case.
        carveout_ind (bool): Boolean indicating if carveout_ind should be set to "Y" or "N".
        default_ind (bool): Boolean indicating if default_ind should be set to "Y" or "N".
        carveout_fields (list[dict]): Running list of dictionaries to append results to.
        questions (dict): List of questions to ask the LLM.
        fs_breakout_questions (dict): Fee schedule questions to ask the LLM.
        filename (str): Filename (used for LLM processing & logging).
        methodology_prompt_template (Callable): Template for methodology prompt.
        fee_schedule_prompt_template (Callable): Template for fee schedule prompt.
        term_key (str): Indicates if the term key should be included.
           If None, it will not be included.  
           If "methodology", the methodology only will be included.
           If "service and methodology", both will be included.

    Returns:
        None: No return value; the function modifies carveout_fields in place.
    """
    # set the carveout code and indicator
    answer_dict["CARVEOUT_CD"] = carveout_type
    answer_dict["CARVEOUT_IND"] = "Y" if carveout_ind else "N"
    # set the default indicator
    answer_dict["DEFAULT_IND"] = "Y" if default_ind else "N"

    result_list = methodology_breakout_single_row(
        reimbursement_primary_answer=answer_dict,
        methodology_breakout_questions=questions,
        fs_breakout_questions=fs_breakout_questions,
        filename=filename,
        methodology_prompt_template=methodology_prompt_template,
        fee_schedule_prompt_template=fee_schedule_prompt_template,
    )

    # Handle term key logic
    if term_key == "methodology":
        for result in result_list:
            result[f"{carveout_type}_TERMS"] = result.get("REIMB_TERM", "")
    elif term_key == "service and methodology":
        for result in result_list:
            result[f"{carveout_type}_TERMS"] = (
                f"{result.get('SERVICE_TERM', '')} - {result.get('REIMB_TERM', '')}"
            )

    carveout_fields.extend(
        result_list
    )  # in-place extends carveout_fields with the new results


def get_carveout_fields(
    answer_dict: dict, carveout_answer: str, filename: str
) -> list[dict]:
    """
    Updates the answer dictionary with carveout-related (and special case) fields based on the given carveout answer.

    Args:
        answer_dict (dict): Dictionary of answers to update.
        carveout_answer (str): The carveout code or description.
        filename (str): The name of the file being processed.

    Returns:
        list[dict]: The updated list of answer dictionaries with carveout fields populated.
    """

    carveout_fields = []  # List to hold the results

    if carveout_answer in SPECIAL_CASE_CONFIGS:
        carveout_config = SPECIAL_CASE_CONFIGS[carveout_answer]
        process_carveout_or_special_case(
            answer_dict=answer_dict,
            carveout_type=carveout_answer,
            carveout_ind=carveout_config["carveout_ind"],
            default_ind=carveout_config["default_ind"],
            carveout_fields=carveout_fields,
            questions=carveout_config["questions"],
            fs_breakout_questions=carveout_config["fs_questions"],
            filename=filename,
            methodology_prompt_template=carveout_config["methodology_prompt_template"],
            fee_schedule_prompt_template=carveout_config[
                "fee_schedule_prompt_template"
            ],
            term_key=carveout_config["term_key"],
        )

    return carveout_fields


def methodology_breakout_single_row(
    reimbursement_primary_answer: dict,
    methodology_breakout_questions: dict,
    fs_breakout_questions: dict,
    filename: str,
    methodology_prompt_template: Callable = investment_prompts.REIMB_TERM_BREAKOUT,
    fee_schedule_prompt_template: Callable = investment_prompts.FEE_SCHEDULE_BREAKOUT,
) -> list[dict]:
    """
    Runs methodology breakout prompt for a single generic (carveout or non-carveout) row, including a Fee Schedule breakout prompt for Fee Schedule rows.

    Args:
        reimbursement_primary_answer (dict): Single dictionary containing reimbursement details.
        methodology_breakout_questions (dict): Dictionary containing methodology breakout prompts.
        fs_breakout_questions (dict): Dictionary containing fee schedule breakout prompts.
        filename (str): The filename used for LLM processing.
        methodology_prompt_template (Callable): Function to generate the methodology prompt. Defaults to investment_prompts.REIMB_TERM_BREAKOUT.
        fee_schedule_prompt_template (Callable): Function to generate the fee schedule prompt. Defaults to investment_prompts.FEE_SCHEDULE_BREAKOUT.

    Returns:
        list[dict]: Updated list with methodology breakout details and fee schedule information.
    """

    answer_dict = reimbursement_primary_answer
    service = answer_dict.get("SERVICE_TERM")
    contract_reimbursement_method = answer_dict.get("REIMB_TERM")
    # Methodology Breakout: AARETE_DERIVED_REIMB_METHOD, REIMB_FEE_RATE, REIMB_PCT_RATE, LESSER_OF_IND, GREATER_OF_IND, AARETE_DERIVED_PROV_TYPE
    if contract_reimbursement_method:
        prompt = methodology_prompt_template(
                service, contract_reimbursement_method, methodology_breakout_questions
            )
        # print(prompt) # Debugging line
        llm_response = llm_utils.invoke_claude(
            prompt,
            config.MODEL_ID_CLAUDE35_SONNET,
            filename,
        )
        # print(llm_response) # Debugging line
        try:
            methodology_breakout_answers = string_utils.universal_json_load(
                llm_response
            )
        except ValueError as e:
            # print(f"universal_string_load failed for input {llm_response} with error {e}\nassigning methodology_breakout_answer to []")
            methodology_breakout_answers = []

        if isinstance(
            methodology_breakout_answers, dict
        ):  # if a single dict is returned, convert it to a list
            methodology_breakout_answers = [methodology_breakout_answers]
        elif not methodology_breakout_answers:
            return [
                answer_dict
            ]  # Return original in a list if no methodology breakout found

        results = []
        for methodology_breakout_dict in methodology_breakout_answers:
            DERIVED_REIMBURSEMENT_METHOD = methodology_breakout_dict.get(
                "AARETE_DERIVED_REIMB_METHOD", ""
            )
            # Fee Schedule Breakout
            if (
                contract_reimbursement_method
                and DERIVED_REIMBURSEMENT_METHOD
                and "Fee Schedule" in DERIVED_REIMBURSEMENT_METHOD
            ):
                fs_breakout_dict = string_utils.universal_json_load(
                    llm_utils.invoke_claude(
                        fee_schedule_prompt_template(
                            contract_reimbursement_method,
                            methodology_breakout_dict.get("REIMB_PCT_RATE"),
                            fs_breakout_questions,
                        ),
                        config.MODEL_ID_CLAUDE35_SONNET,
                        filename,
                    )
                )
            else:
                fs_breakout_dict = {
                    "FEE_SCHEDULE": "N/A",
                    "AARETE_DERIVED_FEE_SCHEDULE": "N/A",
                    "AARETE_DERIVED_FEE_SCHEDULE_VERSION": "N/A",
                }
            results.append(
                {**answer_dict, **methodology_breakout_dict, **fs_breakout_dict}
            )
        return results
    return [answer_dict]  # Return original in a list if no reimbursement method found


def get_methodology_breakout(
    reimbursement_primary_answers: list[dict], filename: str
) -> list[dict]:
    """
    Runs methodology breakout prompt for all non-carveout rows, including Fee Schedule breakout prompt when applicable

    Args:
        reimbursement_primary_answers (list[dict]): List of dictionaries containing reimbursement details.
        filename (str): The filename used for LLM processing.

    Returns:
        list[dict]: Updated list with methodology breakout details and fee schedule information.
        Fields Added:   AARETE_DERIVED_REIMB_METHOD,
                        REIMB_FEE_RATE,
                        REIMB_PCT_RATE,
                        FEE_SCHEDULE,
                        AARETE_DERIVED_FEE_SCHEDULE,
                        AARETE_DERIVED_FEE_SCHEDULE_VERSION
    """

    methodology_breakout_answers = []

    for answer_dict in reimbursement_primary_answers:
        if isinstance(answer_dict, list):  # if a nested list is found, flatten it
            # This shouldn't happen but is in for defensive programming purposes
            print(f"Found a nested list instead of a dict, flattening {answer_dict}")
            for nested_dict in answer_dict:
                _process_carveout_or_breakout(nested_dict, methodology_breakout_answers, filename)
        else:  # normal processing for dict items
            _process_carveout_or_breakout(answer_dict, methodology_breakout_answers, filename)

    return methodology_breakout_answers

def _process_carveout_or_breakout(answer_dict: dict, methodology_breakout_answers: list, filename: str):
    """Processes a single answer dictionary for methodology breakout.

    Args:
        answer_dict (dict): Dictionary containing reimbursement details.
        methodology_breakout_answers (list): List to append results to.
        filename (str): The filename used for LLM processing.

    Returns:
        bool: True if the item was processed, False if it was already seen.
    """
    if answer_dict.get("CARVEOUT_IND") == "Y":
        methodology_breakout_answers.append(answer_dict)
    else:  # If not a carveout, run normal methodology breakout
        methodology_breakout_answers.extend(
            methodology_breakout_single_row(
                answer_dict,
                METHODOLOGY_BREAKOUT_QUESTIONS,
                FS_BREAKOUT_QUESTIONS,
                filename,
            )
        )
    return True

def get_service_breakout(answer_dicts, filename):
    """
    Runs service breakout prompt.

    Args:
        answer_dicts (list[dict]): List of dictionaries containing reimbursement results
        filename (str): The filename used for LLM processing.

    Returns:
        list[dict]: Updated list Prov Type info
        Fields Added:   AARETE_DERIVED_PROV_TYPE

    Add future fields that require both Service and Methodology here
    """

    service_breakout_questions = FieldSet(
        file_path=config.FIELD_JSON_PATH, field_type="service_breakout"
    ).get_prompt_dict()
    already_seen = {}
    service_breakout_answers = []
    for answer_dict in answer_dicts:
        service, methodology = answer_dict.get("SERVICE_TERM"), answer_dict.get(
            "REIMB_TERM"
        )
        service_methodology = service + methodology
        if service_methodology not in already_seen:
            service_breakout_dict = string_utils.universal_json_load(
                llm_utils.invoke_claude(
                    investment_prompts.SERVICE_METHODOLOGY_BREAKOUT(
                        service, methodology, service_breakout_questions
                    ),
                    config.MODEL_ID_CLAUDE35_SONNET,
                    filename,
                )
            )
            already_seen[service_methodology] = service_breakout_dict
        else:
            service_breakout_dict = already_seen[service_methodology]

        service_breakout_answers.append({**answer_dict, **service_breakout_dict})

    return service_breakout_answers


def combine_one_to_n_answers(
    exhibit_level_answers: dict,
    reimbursement_level_answers: list[dict],
    tin_npi_answers: dict,
) -> list[dict]:
    """
    Combines exhibit-level answers with each dictionary in the reimbursement-level answers.

    Args:
        exhibit_level_answers (dict): A dictionary containing exhibit-level key-value pairs.
        reimbursement_level_answers (list[dict]): A list of dictionaries containing reimbursement-level answers.
        tin_npi_answers (dict): A dictionary containing exhibit-level key-value pairs, or nothing, depending on regex dynamic assignment

    Returns:
        list[dict]: A new list of dictionaries where each dictionary combines exhibit-level answers
                    with the respective reimbursement-level answers.
    """

    # Add exhibit-level answers to each reimbursement dictionary
    combined_answers = [
        {**reimbursement_dict, **tin_npi_answers, **exhibit_level_answers}
        for reimbursement_dict in reimbursement_level_answers
    ]

    return combined_answers


def reimbursement_level(
    exhibit_text, filename, reimbursement_level_fields, all_dataset, exhibit_page: str
):
    """
    Processes reimbursement-level fields.
    Currently, runs REIMBURSEMENT_LEVEL_PRIMARY prompt only. Future secondary prompts will go in this function

    Args:
        exhibit_text (str): The text data to process.
        filename (str): The name of the file
        reimbursement_level_fields (FieldSet): A FieldSet object containing field prompts.
        all_dataset (dict): # TODO: add docstring here
        exhibit_page (str): The page number of the exhibit being processed (for a unique identifier)

    Returns:
        dict: The parsed LLM response as a list of dictionaries with unique identifiers
    """
    reimbursement_primary_answers = get_reimbursement_primary(
        reimbursement_level_fields, exhibit_text, filename
    )  # Returns list of dicts

    if reimbursement_primary_answers:  # If any reimbursements found
        carveout_answers = get_special_cases(reimbursement_primary_answers, filename)


        # # Later, we will only run methodology breakout for those without certain carveout indicators. We will run carveout-specific breakouts for the others
        # # not required now as get_special_cases will already have run the methodology breakout for all cases
        # methodology_breakout_answers = get_methodology_breakout(
        #     carveout_answers, filename
        # )

        code_breakout_answers = code_funcs.get_code_breakout(
            carveout_answers, filename, all_dataset
        )
        # print(f"{datetime_str()} code_breakout_answers complete for exhibit page {exhibit_page} of {filename}...")

        return code_breakout_answers
    else:
        return []