doczyai-pipelines/fieldExtraction/src/codes/code_funcs.py

import os
import re

import pandas as pd

import src.config as config
import src.prompts.prompt_templates as prompt_templates
import src.utils.llm_utils as llm_utils
import src.utils.string_utils as string_utils
from constants.constants import Constants
from constants.delimiters import Delimiter
from src.prompts.fieldset import FieldSet

def clean_service(service, constants: Constants) -> str:
    """
    Cleans the service string by removing unnecessary terms and formatting it for further processing.

    Args:
        service (str): The service name or identifier to be cleaned.
    Returns:
        str: The cleaned service string. If the service is empty or contains only stop words, it returns an empty string.
    """

    if string_utils.is_empty(service):
        return ""

    service = service.strip().upper()

    # Map common values
    for acronym, full_form in constants.SYNONYM_MAP.items():
        service = re.sub(rf"\b{re.escape(acronym)}\b", full_form, service)

    # Remove other values
    for term in constants.REMOVAL_LIST:
        service = re.sub(rf"\b{re.escape(term)}\b", "", service).strip()
        service = re.sub(r"\s+", " ", service)
    service = service.replace(" - ", "-")

    if any([v in service for v in ["UNLISTED", "UNCATEGORIZED", "NON-LISTED"]]):
        return "UNLISTED"

    # Check if all remaining words are stop words
    remaining_words = service.split()
    if all(word in constants.STOP_WORD_LIST for word in remaining_words):
        return ""

    return service


def get_regex_answers(service, code_answer_dict):
    """
    **UNUSED** Extracts CPT and Revenue codes from the service string using regex patterns.
    This function may be used in the future to save tokens for simple code extractions. It can be used if there is only one code match.
    If there are multiple matches, using this function risks losing information if we do not have a regex string for EVERY explicit code we need.

    Args:
        service (str): The service name or identifier.
        code_answer_dict (dict): Dictionary to store the extracted codes.
    Returns:
        dict: A dictionary containing the extracted codes. If no codes are found, returns an empty dictionary.
    """

    # Define the regex pattern for CPT codes and ranges
    # Only accept single code or single code ranges - multiple codes may have complexities that need to be prompted for
    cpt_pattern = (
        r"\b(?:[A-Z]\d{4}|\d{5})(?:\s*(?:-|through)\s*(?:[A-Z]\d{4}|\d{5}))?\b"
    )
    cpt_matches = re.findall(cpt_pattern, service)
    if len(cpt_matches) == 1:
        code_answer_dict["PROCEDURE_CD"] = cpt_matches[0].replace("through", "-")

    if "rev" in service.lower():
        rev_pattern = r"\b(?:\d{3}|\d{2}X)(?:\s*(?:-|through)\s*(?:\d{3}|\d{2}X))?\b"
        rev_matches = re.findall(rev_pattern, service)
        if len(rev_matches) > 0:
            code_answer_dict["REVENUE_CD"] = "|".join(rev_matches)

    return code_answer_dict


def code_explicit(service, filename):
    """
    Processes a service string and invokes a language model to generate a response
    based on the provided primary questions. The response is then parsed into a
    dictionary format.
    Args:
        service (str): The service name or identifier. If empty, an empty dictionary
                       is returned.
        code_primary_questions (str): The primary questions or prompts to be used
                                       for generating the response.
        filename (str): The name of the file associated with the operation, used
                        for logging or tracking purposes.
    Returns:
        dict: A dictionary containing the parsed response from the language model.
              Returns an empty dictionary if the service string is empty.
    """

    # Prompt
    code_primary_questions = FieldSet(
        file_path=config.FIELD_JSON_PATH, field_type="code_primary_breakout"
    ).print_prompt_dict()
    claude_answer_raw = llm_utils.invoke_claude(
        prompt_templates.CODE_EXPLICIT(service, code_primary_questions),
        "sonnet_latest",
        filename,
    )
    code_answer_dict = string_utils.universal_json_load(claude_answer_raw)
    return code_answer_dict


def code_category(service, proc_category, hcpcs_level2_mapping, filename):
    """
    Determines if the service explicitly names a code category, such as "A-Codes" or "J-Codes", and returns the corresponding codes.
    Args:
        service (str): The service name or identifier.
        proc_category (list): List of procedure categories to check against.
        filename (str): The name of the file associated with the operation, used for logging or tracking purposes.
    Returns:
        dict: A dictionary containing the procedure codes and their descriptions. If no codes are found, returns an empty dictionary.
    If the service does not explicitly name a code category, it returns an empty dictionary.
    """

    code_category_list = [v.split(":")[1].strip() for v in proc_category if ":" in v]

    matching_values = []
    for category in code_category_list:
        matching_values += [
            value
            for key, value in hcpcs_level2_mapping.items()
            if key.startswith(category)
        ]
    claude_answer_raw = llm_utils.invoke_claude(
        prompt_templates.CODE_CATEGORY(service, matching_values),
        "sonnet_latest",
        filename,
    )
    claude_answer_final = string_utils.universal_json_load(
        claude_answer_raw
    )  # Returns list

    code_answer_dict = {"PROCEDURE_CD": [], "PROCEDURE_CD_DESC": []}
    for answer in claude_answer_final:
        if "0000" in answer:
            code_answer_dict["PROCEDURE_CD"].append(answer)
        elif answer in hcpcs_level2_mapping.values():
            code = list(hcpcs_level2_mapping.keys())[
                list(hcpcs_level2_mapping.values()).index(answer)
            ]
            code_answer_dict["PROCEDURE_CD"].append(code)
            code_answer_dict["PROCEDURE_CD_DESC"].append(answer)

    return code_answer_dict


def code_implicit_special(service, filename):
    """
    Allow certain Service classes to be handled as special cases, such as Drugs, Vaccines, Surgery, and PT/OT/ST.
    This mapping and prompt can be customized, as clients may have different requirements for these categories that are not captured by our Proc Code levels

    Args:
        service (str): The service name or identifier.
        filename (str): The name of the file associated with the operation, used for logging or tracking purposes.
    Returns:
        dict: A dictionary containing the implicit codes found for the service. If no codes are found, returns an empty dictionary.
    """

    claude_answer_raw = llm_utils.invoke_claude(
        prompt_templates.CODE_IMPLICIT_SPECIAL(service), "sonnet_latest", filename
    )
    claude_answer_final = string_utils.extract_text_from_delimiters(
        claude_answer_raw, Delimiter.PIPE
    )

    special_case_mapping = {
        "Drugs": "J0000-J9999",
        "Vaccines": "J0000-J9999|90471‑90474|90620‑90621|90633|90647‑90648|90651|90670|90672|90680‑90681|90686|90696|90698|90700|90707|90710|90713|90714",
        "Surgery": "10004-69990",
        "PT/OT/ST": "92507-92508|92526|97014|97110|97112|97116|97150|97161-97168|97530|97535",
        "PT": "PT Codes TBD",
        "OT": "OT Codes TBD",
        "ST": "ST Codes TBD",
    }
    code_answer_dict = {}
    if claude_answer_final in special_case_mapping:
        code_answer_dict["PROCEDURE_CD"] = special_case_mapping[claude_answer_final]
        code_answer_dict["PROCEDURE_CD_DESC"] = claude_answer_final

    return code_answer_dict


def get_embedding_levels(level_suffix, implicit_run_dict, constants: Constants):
    """
    Retrieves the embedding levels and their corresponding mappings for a given level suffix.
    Args:
        level_suffix (int): The level suffix to determine which mappings to use.
        implicit_run_dict (dict): Dictionary indicating which codes to run implicitly.
    Returns:
        tuple: A tuple containing:
            - embedding_levels (list): List of embedding levels to use.
            - cpt_mapping (dict): Mapping for CPT codes.
            - hcpcs_mapping (dict): Mapping for HCPCS codes.
            - rev_mapping (dict): Mapping for Revenue codes.
    """
    if level_suffix == 1:
        cpt_mapping = constants.CPT_LEVEL1_MAPPING
        hcpcs_mapping = constants.HCPCS_LEVEL1_MAPPING
        rev_mapping = constants.REV_LEVEL1_MAPPING
        embedding_levels = []
        if implicit_run_dict["PROCEDURE_CD"]:
            embedding_levels += [
                f"cpt_level{level_suffix}",
                f"hcpcs_level{level_suffix}",
            ]
        if implicit_run_dict["REVENUE_CD"]:
            embedding_levels += [f"rev_level{level_suffix}"]

    elif level_suffix == 2:
        cpt_mapping = constants.CPT_LEVEL2_MAPPING
        hcpcs_mapping = constants.HCPCS_LEVEL2_MAPPING
        rev_mapping = constants.REV_MAPPING

        embedding_levels = []
        if implicit_run_dict["PROCEDURE_CD"]:
            embedding_levels += [
                f"cpt_level{level_suffix}",
                f"hcpcs_level{level_suffix}",
            ]
        if implicit_run_dict["REVENUE_CD"]:
            embedding_levels += [f"rev"]

    return embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping


def get_match_list(service, embedding_levels, constants: Constants, top_k: int):
    """
    Retrieves the best matches for a given service from multiple embedding levels using FAISS index search.
    Args:
        service (str): The service name or identifier.
        embedding_levels (list): List of embedding levels to search in.
        top_k (int): The number of top matches to retrieve.
    Returns:
        tuple: A tuple containing:
            - match_list (list): List of matched service descriptions.
            - highest_similarity (float): The highest similarity score among the matches.
    """
    target_vec = constants.EMBEDDING_MODEL.encode(
        [service], normalize_embeddings=True
    ).astype("float32")
    match_list = []
    highest_similarity = 0
    for level in embedding_levels:
        embedding = constants.get_embedding(level)

        similarity_scores, match_indices = embedding.index.search(target_vec, top_k)
        if max(similarity_scores[0]) > highest_similarity:
            highest_similarity = max(similarity_scores[0])

        for idx, score in zip(match_indices[0], similarity_scores[0]):
            match_list.append(embedding.choices[idx])

    return match_list, highest_similarity


def code_implicit_rag(service, implicit_run_dict, filename, constants):
    """
    Processes a service string to find implicit codes using RAG (Retrieval-Augmented Generation) methodology.
    Args:
        service (str): The service name or identifier.
        implicit_run_dict (dict): Dictionary indicating which codes to run implicitly.
        filename (str): The name of the file associated with the operation, used for logging or tracking purposes.
    Returns:
        dict: A dictionary containing the implicit codes found for the service. If no codes are found, returns an empty dictionary.
    """

    # Get Embedding Embeddings and Mappings for Levels 1 and 2
    level_dicts = []
    for level_suffix in [1, 2]:
        # Get Mappings and Embeddings
        embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping = (
            get_embedding_levels(level_suffix, implicit_run_dict, constants)
        )
        if not embedding_levels:
            return {}

        # Get Best Matches
        match_list, highest_similarity = get_match_list(
            service, embedding_levels, constants, top_k=5
        )
        level_dicts.append(
            {
                "level_suffix": level_suffix,
                "match_list": match_list,
                "highest_similarity": highest_similarity,
            }
        )

    # Sort level_answers by highest_similarity in descending order
    level_dicts.sort(key=lambda x: x["highest_similarity"], reverse=True)

    # Run the Implicit prompt for each level, starting from the highest similarity
    for level_dict in level_dicts:
        embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping = (
            get_embedding_levels(
                level_dict["level_suffix"], implicit_run_dict, constants
            )
        )

        # Run Prompt
        claude_answer_raw = llm_utils.invoke_claude(
            prompt_templates.CODE_IMPLICIT(service, level_dict["match_list"]),
            "sonnet_latest",
            filename,
        )

        try:
            claude_answer_final = string_utils.universal_json_load(claude_answer_raw)
        except Exception as e:
            return {"CODE_METHODOLOGY": e}

        if not claude_answer_final:
            continue

        # Populate answers, if any
        code_answer_dict = {}
        proc_codes, rev_codes = [], []
        proc_descs, rev_descs = [], []
        for description in claude_answer_final:
            if description == "INVALID_SERVICE":
                code_answer_dict["CODE_METHODOLOGY"] = "Generic - Prompt"
                continue

            if description in cpt_mapping.values():
                matching_codes = [
                    str(key) for key, val in cpt_mapping.items() if val == description
                ]
                if matching_codes:
                    proc_codes.append("|".join(matching_codes))
                    proc_descs.append(description)
            if description in hcpcs_mapping.values():
                matching_codes = [
                    str(key) for key, val in hcpcs_mapping.items() if val == description
                ]
                if matching_codes:
                    proc_codes.append("|".join(matching_codes))
                    proc_descs.append(description)
            if description in rev_mapping.values():
                matching_codes = [
                    str(key) for key, val in rev_mapping.items() if val == description
                ]
                if matching_codes:
                    rev_codes.append("|".join(matching_codes))
                    rev_descs.append(description)

        # Combine answers
        if proc_codes:
            code_answer_dict["PROCEDURE_CD"] = "|".join(proc_codes)
            code_answer_dict["PROCEDURE_CD_DESC"] = "|".join(proc_descs)
            code_answer_dict["CODE_METHODOLOGY"] = (
                f"Implicit - Level {level_dict["level_suffix"]}"
            )
        if rev_codes:
            code_answer_dict["REVENUE_CD"] = "|".join(rev_codes)
            code_answer_dict["REVENUE_CD_DESC"] = "|".join(rev_descs)
            code_answer_dict["CODE_METHODOLOGY"] = (
                f"Implicit - Level {level_dict["level_suffix"]}"
            )

        if code_answer_dict:
            return code_answer_dict

    return {}


def code_last_check(service, filename):
    """
    For Services that did not match any codes, this function categorizes the reason as "Specific" or "Generic".

    Args:
        service (str): The service term.
        filename (str): The name of the file associated with the operation.
    Returns:
        str: "Specific" if the service is a specific case that requires manual intervention,
             "Generic" if the service is a generic case that does not match any codes.
    """

    claude_answer_raw = llm_utils.invoke_claude(
        prompt_templates.CODE_LAST_CHECK(service), "sonnet_latest", filename
    )
    claude_answer_final = string_utils.extract_text_from_delimiters(
        claude_answer_raw, Delimiter.PIPE
    )
    return claude_answer_final


def fill_bill_type(service, answer_dict, BILL_TYPE_MAPPING, BILL_TYPE_REVERSE_MAPPING):
    """
    Fills the BILL_TYPE_CD and BILL_TYPE_CD_DESC fields in the answer dictionary.

    Args:
        service (str): The service term.
        answer_dict (dict): The answer dictionary to update.

    Returns:
        dict: The updated answer dictionary.
    """

    valid_bill_type = sorted(list(set(BILL_TYPE_MAPPING.values())))
    claude_answer_raw = llm_utils.invoke_claude(
        prompt_templates.FILL_BILL_TYPE(service, valid_bill_type), "sonnet_latest", ""
    )
    claude_answer_final = string_utils.universal_json_load(claude_answer_raw)

    if not claude_answer_final:
        return answer_dict

    bill_codes, bill_descs = [], []
    for description in claude_answer_final:
        if description in BILL_TYPE_REVERSE_MAPPING:
            bill_codes.append(BILL_TYPE_REVERSE_MAPPING[description])
            bill_descs.append(description)

    if bill_codes:
        answer_dict["BILL_TYPE_CD"] = "|".join(bill_codes)
        answer_dict["BILL_TYPE_CD_DESC"] = "|".join(bill_descs)

    return answer_dict


def get_implicit_runs(answer_dict):
    """
    Determines which code runs should be executed implicitly based on the Claim Type and Bill Type

    Args:
        answer_dict (dict): Dictionary containing, at minimum, the AARETE_DERIVED_CLAIM_TYPE_CD and BILL_TYPE_CD_DESC fields.
    Returns:
        dict: Dictionary containing boolean values for REVENUE_CD and PROCEDURE_CD indicating whether to run those codes implicitly.
    """

    run_dict = {}

    claim_type = answer_dict.get("AARETE_DERIVED_CLAIM_TYPE_CD")
    bill_type = answer_dict.get("BILL_TYPE_CD_DESC")

    if claim_type == "H":
        run_dict["REVENUE_CD"] = True
    elif claim_type == "M":
        if not string_utils.is_empty(bill_type):
            run_dict["REVENUE_CD"] = True
        else:
            run_dict["REVENUE_CD"] = False
    else:
        run_dict["REVENUE_CD"] = False

    # Run proc code implicit for all bill types except 2
    if bill_type in ["Inpatient Hospital", "Skilled Nursing Facility"]:
        run_dict["PROCEDURE_CD"] = False
    else:
        run_dict["PROCEDURE_CD"] = True

    return run_dict


def extract_codes_from_service(answer_dict, constants: Constants):
    """
    Extracts codes from the service term in the answer dictionary.

    Args:
        answer_dict (dict): Dictionary containing, at minimum, the SERVICE_TERM field.
    Returns:
        dict: Dictionary containing extracted codes and their descriptions.

    Codes Extracted:
        Explicit Codes:
            - PROCEDURE_CD: CPT, HCPCS
            - CPT4_PROC_MOD: CPT4 Modifiers
            - REVENUE_CD: Revenue codes
            - DIAG_CD: Diagnosis codes (ICD-10)
            - GROUPER_CD: MS and APR DRG Grouper Codes
            - NDC_CD: National Drug Codes (NDC)
            - CLAIM_ADMIT_TYPE_CD: Claim Admit Type Code
            - CLAIM_STATUS_CD: Claim Status Code
        Implicit Codes (Up to Level 2):
            - PROCEDURE_CD: CPT, HCPCS
            - REVENUE_CD: Revenue codes
    """

    service, bill_type = answer_dict.get("SERVICE_TERM"), answer_dict.get(
        "BILL_TYPE_CD_DESC"
    )

    if string_utils.is_empty(service):
        return {}

    # Fill Bill Type if not filled
    if string_utils.is_empty(bill_type):
        answer_dict = fill_bill_type(
            service,
            answer_dict,
            constants.BILL_TYPE_MAPPING,
            constants.BILL_TYPE_REVERSE_MAPPING,
        )

    # Get list of codes we want implicit - based on claim type
    implicit_run_dict = get_implicit_runs(answer_dict)

    # Preprocessing
    service_clean = clean_service(service, constants)

    # Exit point for unlisted
    if service_clean == "UNLISTED":
        answer_dict["PROCEDURE_CD_DESC"] = "UNLISTED"
        answer_dict["CODE_METHODOLOGY"] = "UNLISTED"
        return answer_dict

    # Exit point if N/A or in DO_NOT_RUN list
    if string_utils.is_empty(service_clean) or any(
        [v in service_clean for v in constants.DO_NOT_RUN]
    ):
        answer_dict["CODE_METHODOLOGY"] = "Generic - Before Prompts"
        return answer_dict

    # Explicit Codes (run always)
    code_answer_dict = code_explicit(service_clean, "")
    if any(
        not string_utils.is_empty(value) for value in code_answer_dict.values()
    ) and all(
        ["Category" not in v for v in code_answer_dict.get("PROCEDURE_CD", "")]
    ):  # if any code value is not empty, return
        if "NOT_ESTABLISHED" in code_answer_dict.get("PROCEDURE_CD", ""):
            code_answer_dict["CODE_METHODOLOGY"] = "Implicit - Not Established"
            code_answer_dict["PROCEDURE_CD"] = "[]"
        else:
            code_answer_dict["CODE_METHODOLOGY"] = "Explicit"
        answer_dict.update(code_answer_dict)
        return answer_dict

    # Implicit Codes: Code Categories
    if any(["Category:" in v for v in code_answer_dict.get("PROCEDURE_CD", "")]):
        code_answer_dict = code_category(
            service_clean,
            code_answer_dict.get("PROCEDURE_CD", ""),
            constants.HCPCS_LEVEL2_MAPPING,
            "",
        )
        if any(
            not string_utils.is_empty(value) for value in code_answer_dict.values()
        ):  # if any code value is not empty, return
            code_answer_dict["CODE_METHODOLOGY"] = "Implicit - Letter Category"
            answer_dict.update(code_answer_dict)
            return answer_dict

    # Implicit Codes: Special Categories
    code_answer_dict = code_implicit_special(service_clean, "")
    if code_answer_dict:
        code_answer_dict["CODE_METHODOLOGY"] = "Implicit - Special Case"
        answer_dict.update(code_answer_dict)
        return answer_dict

    # Implicit Codes: Levels 1 and 2
    code_answer_dict = code_implicit_rag(
        service_clean, implicit_run_dict, "", constants
    )
    if any(
        not string_utils.is_empty(value) for value in code_answer_dict.values()
    ):  # if any code value is not empty, return
        answer_dict.update(code_answer_dict)
        return answer_dict

    # No Match - Why? Generic or Specific
    last_check_answer = code_last_check(service_clean, "")
    if last_check_answer == "Generic":
        answer_dict["CODE_METHODOLOGY"] = "Generic - No Match"
    elif last_check_answer == "Specific":
        answer_dict["CODE_METHODOLOGY"] = "Specific - No Match"

    return answer_dict


def fill_claim_type(answer_dicts):
    """
    Fills in the AARETE_DERIVED_CLAIM_TYPE_CD field in each answer_dict with the mode of the existing values.
    If the field is already populated, it remains unchanged.
    Args:
        answer_dicts (list[dict]):  List of dictionaries containing 1:1 and 1:N fields (after merge process)
    Returns:
        list[dict]: List of dictionaries with AARETE_DERIVED_CLAIM_TYPE
    """
    claim_types = [
        d.get("AARETE_DERIVED_CLAIM_TYPE_CD")
        for d in answer_dicts
        if d.get("AARETE_DERIVED_CLAIM_TYPE_CD")
    ]
    if not claim_types:
        return answer_dicts

    mode_claim_type = max(set(claim_types), key=claim_types.count)
    for answer_dict in answer_dicts:
        if not answer_dict.get("AARETE_DERIVED_CLAIM_TYPE_CD"):
            answer_dict["AARETE_DERIVED_CLAIM_TYPE_CD"] = mode_claim_type
    return answer_dicts


def code_breakout(merged_results: pd.DataFrame, constants: Constants):
    """
    Processes a list of answer dictionaries to extract and fill in code-related information.
    Wrapper function to be run from main Doczy

    Args:
        merged_results (pd.DataFrame): DataFrame containing merged results from the processing pipeline
    Returns:
        pd.DataFrame: DataFrame with updated code fields.
    """

    # Fill empty AARETE_DERIVED_CLAIM_TYPE_CD with mode
    answer_dicts = fill_claim_type(merged_results.to_dict(orient="records"))

    answer_dicts_with_code = []
    for answer_dict in answer_dicts:
        code_answer_dict = extract_codes_from_service(answer_dict, constants)
        answer_dict.update(code_answer_dict)
        answer_dicts_with_code.append(answer_dict)

    return pd.DataFrame(answer_dicts_with_code)


def grouper_breakout(results_with_code: pd.DataFrame):
    """
    Processes a DataFrame to extract and fill in grouper-related information in cases where we got a grouper_cd from code_breakout
    but didn't get grouper fields from methodology + grouper breakout as in special case grouper was not identified and extracted.
    Args:
        results_with_code (pd.DataFrame): DataFrame containing the answer dictionaries with code fields
    Returns:
        pd.DataFrame: DataFrame with updated grouper fields.
    """

    # try grouper fields again if they are empty and grouper_cd is not empty
    answer_dicts_with_code = results_with_code.to_dict(orient="records")
    final_answer_dicts = []
    GROUPER_QUESTIONS = FieldSet(
        file_path=config.FIELD_JSON_PATH, field_type="grouper"
    ).print_prompt_dict()
    for answer_dict in answer_dicts_with_code:
        if (
            string_utils.is_empty(answer_dict.get("GROUPER_TYPE"))
            and not string_utils.is_empty(answer_dict.get("GROUPER_CD"))
        ):
            # run groper breakout prompt
            grouper_breakout_prompt = prompt_templates.GROUPER_BREAKOUT(
                answer_dict.get("SERVICE_TERM", ""),
                answer_dict.get("REIMB_TERM", ""),
                GROUPER_QUESTIONS
            )
            claude_answer_raw = llm_utils.invoke_claude(
                grouper_breakout_prompt, "sonnet_latest", ""
            )
            try:
                grouper_answer = string_utils.universal_json_load(claude_answer_raw)
            except Exception as e:
                grouper_answer = {"GROUPER_TYPE": e}

            answer_dict.update(grouper_answer)
        final_answer_dicts.append(answer_dict)

    return pd.DataFrame(final_answer_dicts)