Merged in feature/doczy-codes (pull request #619)

Feature/doczy codes * Ignore service and procedure generic terms * UPdate PT/OT/ST codes * Standardize stop-loss * UPdate for multiple identical values * Modify code funcs * Run highest match Implicit Level first * Fix implicit code bug * Add PT individually to special case * change revenue code conditions * Merge branch 'main' into feature/doczy-codes * Merged main into feature/doczy-codes * Update main * Relocate create_code_embeddings back to scripts * Remove investment code_funcs (use codes.code_funcs now) * Add docstrings * stop-loss * Update unit tests * Remove test * Remove reimb_dates * fix dtypes * Remove prints * update prompt * Update code breakout unit test * Merge branch 'main' into feature/doczy-codes * Merge branch 'main' into feature/doczy-codes Approved-by: Alex Galarce
2025-08-01 22:36:01 +00:00
parent 84a09002fa
commit c2d5db19d4
47 changed files with 1843 additions and 1362 deletions
@@ -80,21 +80,12 @@ myenv/
 !nltk_data/tokenizers/punkt_tab/english/
 !nltk_data/tokenizers/punkt_tab/english/*.txt
 *.tab
 !nltk_data/
 !nltk_data/tokenizers/
 !nltk_data/tokenizers/punkt_tab/
 !nltk_data/tokenizers/punkt_tab/english/
 !nltk_data/tokenizers/punkt_tab/english/*.tab
 *.pdf
 !requirements.txt
 subset/
 output/
 docs/
 allfiles/
 CNC_testruns_AC/
 cnc-abc-output/
 humana-input/
 humana-output/
 output_consolidated/
 *.zip
 *.xlsx
@@ -163,12 +163,12 @@ class CrosswalkBuilder:
        # First collect all sources for each target
        for source, target in self.mapping.items():
            if target in reverse_dict:
-                reverse_dict[target].append(source)
+                reverse_dict[target].append(str(source))
            else:
-                reverse_dict[target] = [source]
+                reverse_dict[target] = [str(source)]
        # Then join the lists with commas
-        return {target: ", ".join(sources) for target, sources in reverse_dict.items()}
+        return {target: "|".join(sources) for target, sources in reverse_dict.items()}
 def apply_crosswalk(
    val: str, mapping: dict[str, str], default: str = ""
@@ -87,7 +87,7 @@
 96360-96549,"Hydration, Therapeutic, Prophylactic, Diagnostic Injections and Infusions, and Chemotherapy and Other Highly Complex Drug or Highly Complex Biologic Agent Administration"
 96567-96574,Photodynamic Therapy Procedures
 96900-96999,Special Dermatological Procedures
-97010-97799,Physical Medicine and Rehabilitation Evaluations
+97010-97799,Physical Medicine and Rehabilitation
 97151-97158,Adaptive Behavior Services
 97802-97804,Medical Nutrition Therapy Procedures
 97810-97814,Acupuncture Procedures
@@ -0,0 +1,82 @@
 Code,Description
 0810|0811|0812|0813|0814|0819,Acquisition of Body Components
 0390|0391|0392|0399,"Administration, Processing and Storage for Blood and Blood Components"
 3101|3102|3103|3104|3105|3109,Adult Care
 0240|0241|0242|0243|0249,All Inclusive Ancillary
 0100|0101,All inclusive Rate
 2100|2101|2102|2103|2104|2105|2106|2109,Alternative Therapy Services
 0540|0541|0542|0543|0544|0545|0546|0547|0548|0549,Ambulance
 0490|0499,Ambulatory Surgical Care
 0370|0371|0372|0374|0379,Anesthesia
 0470|0471|0472|0479,Audiology
 1000|1001|1002|1003|1004|1005,Behavior Health Accommodations
 0900|0901|0902|0903|0904|0905|0906|0907|0911|0912|0913|0914|0915|0916|0917|0918|0919,Behavioral Health Treatment/Services
 0380|0381|0382|0383|0384|0385|0386|0387|0389,Blood and Blood Components
 0350|0351|0352|0359,CT Scan
 0480|0481|0482|0483|0489,Cardiology
 0700,Cast Room
 0510|0511|0512|0513|0514|0515|0516|0517|0519,Clinic
 0840|0841|0842|0843|0844|0845|0849,Continuous Ambulatory Peritoneal Dialysis (CAPD)
 0855,Continuous Cycling Peritoneal Dialysis (CAPD)
 0850|0851|0852|0853|0854|0859,Continuous Cycling Peritoneal Dialysis (CCPD)
 0210|0211|0212|0213|0214|0219,Coronary Care Unit
 0290|0291|0292|0293|0294|0299,Durable Medical Equipment (other than renal)
 0740,EEG (Electroencephalogram)
 0730|0731|0732|0739,EKG/ECG (Electrocardiogram)
 0450|0451|0452|0456|0459,Emergency Room
 0790,Extra
 0520|0521|0522|0523|0524|0525|0526|0527|0528|0529,Free
 0750,Gastro
 0022|0023|0024,Health Insurance
 0820|0821|0822|0823|0824|0825|0829,Hemodialysis
 0560|0561|0562|0569|0580|0581|0582|0583|0589|0600|0601|0602|0603|0604|0609,Home Health (HH)
 0570|0571|0572|0579,Home Health (HH) Aide
 0590,Home Health (HH) Units of Service
 0640|0641|0642|0643|0644|0645|0646|0647|0648|0649,Home IV Therapy Services
 0650|0651|0652|0655|0656|0657|0658|0659,Hospice Service
 0260|0261|0262|0263|0264|0269,IV Therapy
 0230|0231|0232|0233|0234|0235|0239,Incremental Nursing Charge
 0800|0801|0802|0803|0804|0809,Inpatient Renal Dialysis
 0200|0201|0202|0203|0204|0206|0207|0208|0209,Intensive Care Unit
 0720|0721|0722|0723|0724|0729,Labor Room/Delivery
 0300|0301|0302|0303|0304|0305|0306|0307|0309,Laboratory
 0310|0311|0312|0314|0319,Laboratory Pathology
 0180|0182|0183|0185|0189,Leave of Absence
 0610|0611|0612|0614|0615|0616|0618|0619,Magnetic Resonance Technology (MRT)
 0860|0861,Magnetoencephalography (MEG)
 0931|0932,Medical Rehabilitation Day Program
 0621|0622|0623|0624,Medical Surgical Supplies
 0270|0271|0272|0273|0274|0275|0276|0277|0278|0279,Medical/Surgical Supplies and Devices
 0880|0881|0882|0889,Miscellaneous Dialysis
 9999,Not Applicable
 0340|0341|0342|0343|0344|0349,Nuclear Medicine
 0170|0171|0172|0173|0174|0179,Nursery
 0430|0431|0432|0433|0434|0439,Occupational Therapy
 0280|0289,Oncology
 0360|0361|0362|0367|0369,Operating Room Services
 0530|0531|0539,Osteopathic Services
 0920|0921|0922|0923|0924|0925|0929,Other Diagnostic Services
 0400|0401|0402|0403|0404|0409,Other Imaging Services
 0940|0941|0942|0943|0944|0945|0946|0947|0948|0949|0951|0952,Other Therapeutic Services
 0500|0509,Outpatient Services
 0670|0671|0672|0679,Outpatient Special Residence Charges
 0990|0991|0992|0993|0994|0995|0996|0997|0998|0999,Patient Convenience Items
 0830|0831|0832|0833|0834|0835|0839,Peritoneal Dialysis
 0250|0251|0252|0253|0254|0255|0256|0257|0258|0259|0631|0632|0633|0634|0635|0636|0637,Pharmacy
 0420|0421|0422|0423|0424|0429,Physical Therapy
 0770|0771,Preventive Care Services
 0960|0961|0962|0963|0964|0969|0971|0972|0973|0974|0975|0976|0977|0978|0979|0981|0982|0983|0984|0985|0986|0987|0988|0989,Professional Fees
 0460|0469,Pulmonary Function
 0320|0321|0322|0323|0324|0329|0330|0331|0332|0333|0335|0339,Radiology
 0710,Recovery Room
 0410|0412|0413|0419,Respiratory Services
 0660|0661|0662|0663|0669,Respite Care
 0110|0111|0112|0113|0114|0115|0116|0117|0118|0119|0120|0121|0122|0123|0124|0125|0126|0127|0128|0129|0130|0131|0132|0133|0134|0135|0136|0137|0138|0139|0140|0141|0142|0143|0144|0145|0146|0147|0148|0149|0150|0151|0152|0153|0154|0155|0156|0157|0158|0159|0160|0164|0167|0169,Room & Board
 0550|0551|0552|0559,Skilled Nursing
 0220|0221|0222|0223|0224|0229,Special Charges
 0760|0761|0762|0769,Specialty Services
 0440|0441|0442|0443|0444|0449,Speech Therapy
 0190|0191|0192|0193|0194|0199,Subacute Care
 0780,Telemedicine
 0001,Total Charge
 0681|0682|0683|0684|0689,Trauma Response
@@ -18,7 +18,7 @@
        "033X" : "Home",
        "034X" : "Home",
        "071X" : "Rural Health Clinic",
-        "072X" : "End-Stage Renal Disease Treatment Facility",
+        "072X" : "Dialysis",
        "073X" : "Independent Clinic",
        "074X" : "Independent Clinic",
        "079X" : "Independent Clinic",
@@ -12,7 +12,6 @@
        "Community First Choice (CFC)" : "CFC",
        "Foster Care (FC)" : "FC",
        "Home and Community Based Services Waiver (HCBS)" : "HCBS",
        "Hospice Benefits (HOSPICE)" : "HOSPICE",
        "Intellectual and Developmental Disability (IDD)" : "IDD",
        "Institution for Medical Disease (IMD)" : "IMD",
        "Managed Long Term Care (MLTC)" : "MLTC",
@@ -63,7 +62,7 @@
            "ARKids First" : "CHIP"
        },
        "CA" : {
-            "Medi-Cal" : "MEDICAL",
+            "Medi-Cal" : "MEDI_CAL",
            "Covered California" : "COVEREDCAL",
            "Medi-Cal Access Program (MCAP)" : "CHIP",
            "Healthy Families Benefit Program" : "CHIP"
@@ -57,7 +57,7 @@
        },
        "CA" : {
-            "MEDICAL" : "Medicaid",
+            "MEDI_CAL" : "Medicaid",
            "COVEREDCAL" : "Marketplace"
        },
        "CO" : {
@@ -0,0 +1,30 @@
 import src.utils.embedding_utils as embedding_utils
 from crosswalk.crosswalk_utils import CrosswalkBuilder
 import src.codes.code_funcs as code_funcs
 import os
 import pandas as pd
 from sentence_transformers import SentenceTransformer
 # Load model
 roberta_model = SentenceTransformer("all-roberta-large-v1")
 mapping_dir = "crosswalk/mapping_csvs/proc_cd"
 pkl_dir = "embeddings"
 for filename in os.listdir(mapping_dir):
    if filename.endswith(".csv") and "level" in filename:
        print(f"Creating embeddings for {filename}")
        stripped_filename = filename.replace(".csv", "")
        mapping_df = pd.read_csv(os.path.join(mapping_dir, filename))
        proc_crosswalk = CrosswalkBuilder().from_df(mapping_df, from_col="Code", to_col="Description")
        proc_choices = [x for x in proc_crosswalk.mapping.values()]
        os.makedirs(os.path.join(pkl_dir, stripped_filename), exist_ok=True)
        proc_index = embedding_utils.create_faiss_index(choices=proc_choices, 
                                                        model=roberta_model, 
                                                        save_path=os.path.join(pkl_dir, stripped_filename, "faiss_index.bin"),
                                                        embedding_path=os.path.join(pkl_dir, stripped_filename, "embeddings.npy"),
                                                        choices_path=os.path.join(pkl_dir, stripped_filename, "choices.pkl")
                                                        )        
@@ -1,54 +0,0 @@
 import src.utils.embedding_utils as embedding_utils
 from src.investment.code_funcs import get_proc_crosswalk, get_clean_value
 from crosswalk.crosswalk_utils import CrosswalkBuilder
 import faiss
 import pickle
 import os
 import pandas as pd
 import numpy as np
 from sentence_transformers import SentenceTransformer
 # Load model
 roberta_model = SentenceTransformer("all-roberta-large-v1")
 # If running proc_code for the first time
 """
 mapping_dir = "crosswalk/mapping_csvs/proc_cd"
 pkl_dir = "embeddings"
 for filename in os.listdir(mapping_dir):
    if filename.endswith(".csv"):
        print(f"Creating embeddings for {filename}")
        stripped_filename = filename.strip(".csv")
        mapping_df = pd.read_csv(os.path.join(mapping_dir, filename))
        proc_crosswalk = CrosswalkBuilder().from_df(mapping_df, from_col="Code", to_col="Description")
        proc_choices = [get_clean_value(x) for x in proc_crosswalk.mapping.values()]
        os.makedirs(os.path.join(pkl_dir, stripped_filename), exist_ok=True)
        proc_index = embedding_utils.create_faiss_index(choices=proc_choices, 
                                                        model=roberta_model, 
                                                        save_path=os.path.join(pkl_dir, stripped_filename, "faiss_index.bin"),
                                                        embedding_path=os.path.join(pkl_dir, stripped_filename, "embeddings.npy"),
                                                        choices_path=os.path.join(pkl_dir, stripped_filename, "choices.pkl")
                                                        )        
 """
 # If running diag_code for the first time
 mapping_dir = "crosswalk/mapping_csvs/diag_cd"
 pkl_dir = "embeddings"
 for filename in os.listdir(mapping_dir):
    if filename.endswith(".csv"):
        print(f"Creating embeddings for {filename}")
        stripped_filename = filename.strip(".csv")
        mapping_df = pd.read_csv(os.path.join(mapping_dir, filename))
        mapping_df = mapping_df.astype(str)
        diag_crosswalk = CrosswalkBuilder().from_df(mapping_df, from_col="Code", to_col="Description")
        diag_choices = [get_clean_value(x) for x in diag_crosswalk.mapping.values()]
        os.makedirs(os.path.join(pkl_dir, stripped_filename), exist_ok=True)
        proc_index = embedding_utils.create_faiss_index(choices=diag_choices,
                                                        model=roberta_model,
                                                        save_path=os.path.join(pkl_dir, stripped_filename, "faiss_index.bin"),
                                                        embedding_path=os.path.join(pkl_dir, stripped_filename, "embeddings.npy"),
                                                        choices_path=os.path.join(pkl_dir, stripped_filename, "choices.pkl")
                                                        )
@@ -0,0 +1,58 @@
 from crosswalk.crosswalk_utils import CrosswalkBuilder
 import src.codes.code_funcs as code_funcs
 INPUT_FILE_PATH = "Codes-Test-Doc.xlsx"
 OUTPUT_FILE_PATH = "Doczy-Codes-Test-9.csv"
 SERVICE_COLUMN = "SERVICE_TERM"
 SYNONYM_MAP = {"DME" : "DURABLE MEDICAL EQUIPMENT",
               "SNF" : "SKILLED NURSING FACILITY",
               "ASC" : "AMBULATORY SURGERY CENTER",
               "LABORATORY" : "LABORATORY AND PATHOLOGY",
               "LAB" : "LABORATORY AND PATHOLOGY",
               "STOP LOSS" : "STOP-LOSS",
               "OB": "OBSTETRIC"
 }
 REMOVAL_LIST = [v.upper() for v in 
                list(set(CrosswalkBuilder().from_json(path="crosswalk/mappings/crosswalk_lob.json").mapping.keys())) +          # LOB
                list(set(CrosswalkBuilder().from_json(path="crosswalk/mappings/crosswalk_product_lob.json").mapping.keys())) +  # Product
                list(set(CrosswalkBuilder().from_json(path="crosswalk/mappings/crosswalk_network.json").mapping.values())) +    # Network
                list(set(CrosswalkBuilder().from_json(path="crosswalk/mappings/crosswalk_claim_type.json").mapping.keys())) +   # Claim Type
                list(set(CrosswalkBuilder().from_json(path="crosswalk/mappings/crosswalk_program.json").mapping.values())) +    # Program
                ["INPATIENT", "OUTPATIENT", "IN-PATIENT", "OUT-PATIENT", "IP", "OP"] + 
                ["COVERED SERVICES", "SERVICES", "PROCEDURES", "CONTRACTED", "HEALTHCARE", "HEALTH CARE", "ALL OTHER", 
                 "PER DIEM", "ADMISSION", "REIMBURSEMENT", "ADDITIONAL PAYMENT", "PAYMENT", "RATE", "FEE SCHEDULE", "METHODOLOGY", "OUTLIER", 
                 "INTERMEDIATE", "READMISSION", "ACUTE", "OPPS", "FEE-FOR-SERVICE", "PROGRAM", "PRODUCT", 
                ] # Generic Terms not found in any of the proc code mappings
 ]
 STOP_WORD_LIST = [
    "THE", "AND", "OF", "TO", "A", "IN", "FOR", "ON", "WITH", "AT", "BY", 
    "AN", "BE", "THIS", "THAT", "FROM", "OR", "AS", "IS", "ARE", "WAS", 
    "WERE", "WHICH", "WHO", "ALL", "OTHER", "/", ",", "NOT"
 ]
 DO_NOT_RUN = ["BY REPORT", "STOP-LOSS"]
 from sentence_transformers import SentenceTransformer
 model = SentenceTransformer("all-roberta-large-v1")
 # Code Crosswalks
 cpt_level1_mapping = CrosswalkBuilder().from_excel("crosswalk/mapping_csvs/proc_cd/cpt_level1.csv", "Code", "Description").mapping
 cpt_level2_mapping = CrosswalkBuilder().from_excel("crosswalk/mapping_csvs/proc_cd/cpt_level2.csv", "Code", "Description").mapping
 cpt_level3_mapping = CrosswalkBuilder().from_excel("crosswalk/mapping_csvs/proc_cd/cpt_level3.csv", "Code", "Description").mapping
 hcpcs_level1_mapping = CrosswalkBuilder().from_excel("crosswalk/mapping_csvs/proc_cd/hcpcs_level1.csv", "Code", "Description").mapping
 hcpcs_level2_mapping = CrosswalkBuilder().from_excel("crosswalk/mapping_csvs/proc_cd/hcpcs_level2.csv", "Code", "Description").mapping
 hcpcs_mapping = CrosswalkBuilder().from_excel("crosswalk/mapping_csvs/proc_cd/hcpcs.csv", "Code", "Description").mapping
 rev_level1_mapping = CrosswalkBuilder().from_excel("crosswalk/mapping_csvs/rev_cd/rev_level1.csv", "Code", "Description").mapping
 rev_mapping = CrosswalkBuilder().from_excel("crosswalk/mapping_csvs/rev_cd/rev.csv", "Code", "Description").mapping
 bill_type_mapping = CrosswalkBuilder().from_json("crosswalk/mappings/crosswalk_bill_type.json").mapping
 bill_type_reverse_mapping = CrosswalkBuilder().from_json("crosswalk/mappings/crosswalk_bill_type.json").create_reverse_mapping()
@@ -0,0 +1,552 @@
 from src.prompts.investment_prompts import FieldSet
 import src.config as config
 import src.utils.string_utils as string_utils
 import src.utils.llm_utils as llm_utils
 import src.prompts.investment_prompts as investment_prompts
 import src.utils.embedding_utils as embedding_utils
 import src.codes.code_constants as code_constants
 from src.enums.delimiters import Delimiter
 import pandas as pd
 import os
 import re
 def clean_service(service) -> 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 code_constants.SYNONYM_MAP.items():
        service = re.sub(rf'\b{re.escape(acronym)}\b', full_form, service)
    # Remove other values
    for term in code_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 code_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").get_prompt_dict()
    claude_answer_raw = llm_utils.invoke_claude(
                            investment_prompts.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, 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 code_constants.hcpcs_level2_mapping.items() if key.startswith(category)]
    claude_answer_raw = llm_utils.invoke_claude(
                            investment_prompts.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 code_constants.hcpcs_level2_mapping.values():
            code = list(code_constants.hcpcs_level2_mapping.keys())[list(code_constants.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(
                            investment_prompts.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 not string_utils.is_empty(claude_answer_final):
        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):
    """
    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 = code_constants.cpt_level1_mapping
        hcpcs_mapping = code_constants.hcpcs_level1_mapping
        rev_mapping = code_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 = code_constants.cpt_level2_mapping
        hcpcs_mapping = code_constants.hcpcs_level2_mapping
        rev_mapping = code_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, top_k):
    """
    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 = code_constants.model.encode([service], normalize_embeddings=True).astype("float32")
    match_list = []
    highest_similarity = 0
    for level in embedding_levels:
        index, _, choices = embedding_utils.load_faiss_index(index_path=os.path.join("embeddings", level, "faiss_index.bin"),
                                                            embedding_path=os.path.join("embeddings", level, "embeddings.npy"),
                                                            choices_path=os.path.join("embeddings", level, "choices.pkl"),)
        similarity_scores, match_indices = 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(choices[idx])
    return match_list, highest_similarity
 def code_implicit_rag(service, implicit_run_dict, filename):
    """
    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)        
        if not embedding_levels:
            return {}
        # Get Best Matches
        match_list, highest_similarity = get_match_list(service, embedding_levels, 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)
        # Run Prompt
        claude_answer_raw = llm_utils.invoke_claude(
                                investment_prompts.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 = [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 = [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 = [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(
                            investment_prompts.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):
    """
    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(code_constants.bill_type_mapping.values())))
    claude_answer_raw = llm_utils.invoke_claude(
                            investment_prompts.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 code_constants.bill_type_reverse_mapping:
            bill_codes.append(code_constants.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):
    """
    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")
    # Fill Bill Type if not filled
    if string_utils.is_empty(bill_type):
        answer_dict = fill_bill_type(service, answer_dict)
    # 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)
    # 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 code_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", ""), "")
        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, "")
    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):
    """
    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'))
    final_answer_dicts = []
    for answer_dict in answer_dicts:
        code_answer_dict = extract_codes_from_service(answer_dict)
        answer_dict.update(code_answer_dict)
        final_answer_dicts.append(answer_dict)
    # Convert back to DataFrame
    final_answer_df = pd.DataFrame(final_answer_dicts)
    return final_answer_df
@@ -0,0 +1,87 @@
 import src.utils.io_utils as io_utils
 import src.codes.code_funcs as code_funcs
 import src.codes.code_constants as code_constants
 import src.config as config
 import concurrent.futures
 def main():
    df = io_utils.read_local(code_constants.INPUT_FILE_PATH)
    # Process entire rows instead of just services
    row_indices = list(range(len(df)))
    # Define a function to process a single row
    def process_row(idx):
        try:
            # Create answer_dict from the row
            row = df.iloc[idx]
            answer_dict = {
                "SERVICE_TERM": str(row[code_constants.SERVICE_COLUMN]),
                "AARETE_DERIVED_CLAIM_TYPE_CD": str(row.get("AARETE_DERIVED_CLAIM_TYPE_CD", "")),
                "BILL_TYPE_CD_DESC": str(row.get("BILL_TYPE_CD_DESC", ""))
            }
            # Process the row
            result_dict = code_funcs.extract_codes_from_service(answer_dict)
            return idx, result_dict
        except Exception as e:
            print(f"Error processing row {idx}: {e}")
            return idx, {"SERVICE_TERM": str(row[code_constants.SERVICE_COLUMN]), "ERROR": str(e)}
    # Use ThreadPoolExecutor to process rows in parallel
    results = {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=config.MAX_WORKERS) as executor:
        futures = {executor.submit(process_row, idx): idx for idx in row_indices}
        for future in concurrent.futures.as_completed(futures):
            try:
                idx, result_dict = future.result()
                results[idx] = result_dict
            except Exception as e:
                idx = futures[future]
                print(f"Error retrieving result for row {idx}: {e}")
                row = df.iloc[idx]
                results[idx] = {"SERVICE_TERM": str(row[code_constants.SERVICE_COLUMN]), "ERROR": str(e)}
    # Update the DataFrame with the results
    for idx, result_dict in results.items():
        for key, value in result_dict.items():
            if key not in df.columns:
                df[key] = None
            # Ensure value is scalar
            if isinstance(value, (list, dict)):
                value = str(value)  # Convert lists/dicts to strings
            try:
                # Update the specific cell
                df.at[idx, key] = value
            except Exception as e:
                print(f"Error updating DataFrame at row {idx}, key '{key}': {e}")
    # Save the updated DataFrame
    try:
        if code_constants.OUTPUT_FILE_PATH:
            df.to_csv(code_constants.OUTPUT_FILE_PATH, index=False)
            print(f"Results saved to {code_constants.OUTPUT_FILE_PATH}")
        else:
            print("No output path specified. Results not saved.")
    except Exception as e:
        print(f"Error saving results: {e}")
        # Try to save to a backup location
        backup_path = "backup_results.csv"
        try:
            df.to_csv(backup_path, index=False)
            print(f"Saved backup to {backup_path}")
        except:
            print("Could not save backup file")
    return df
 if __name__ == "__main__":
    main()
@@ -1,547 +0,0 @@
 from src.prompts.investment_prompts import Field, FieldSet
 import src.config as config
 import src.utils.string_utils as string_utils
 import src.utils.llm_utils as llm_utils
 import src.prompts.investment_prompts as investment_prompts
 import src.utils.embedding_utils as embedding_utils
 import pandas as pd
 import os
 import re
 import ast
 import src.constants.investment_values as investment_values
 from crosswalk.crosswalk_utils import CrosswalkBuilder
 import faiss
 import numpy as np
 from sentence_transformers import SentenceTransformer
 # Load model
 model = SentenceTransformer("all-roberta-large-v1")
 def create_faiss_index(choices):
    """
    Creates a FAISS index for a list of choices using cosine similarity (inner product).
    Args:
        choices (list): A list of text items to index.
    Returns:
        faiss.IndexFlatIP: A FAISS index for the given embeddings.
    """
    embeddings = model.encode(choices, normalize_embeddings=True).astype("float32")
    index = faiss.IndexFlatIP(embeddings.shape[1])  # Inner product for cosine similarity
    index.add(embeddings)
    return index
 def get_matches_above_threshold(target, choices, index, threshold=0.7, top_k=5):
    """
    Retrieves the top-k matches from a FAISS index with similarity scores above a given threshold.
    Args:
        target (str): The target text to compare against the choices.
        choices (list): A list of choices to match the target with.
        index (faiss.IndexFlatIP): The FAISS index for fast nearest-neighbor search.
        threshold (float, optional): The minimum similarity score to consider a match. Defaults to 0.7.
        top_k (int, optional): The number of top matches to retrieve. Defaults to 5.
    Returns:
        list: A list of tuples containing matched choice and score for each match above the threshold.
    """
    target_vec = model.encode([target], normalize_embeddings=True).astype("float32")
    similarity_scores, match_indices = index.search(target_vec, top_k)
    return [
        (choices[idx], score)
        for idx, score in zip(match_indices[0], similarity_scores[0])
        if score >= threshold
    ]
 def crosswalk_levels(code_type: str, levels: list[str]) -> dict[str, tuple]:
    """
    Loads and processes crosswalk mappings for different levels (e.g., CPT, HCPCS) from CSV files, 
    and returns a dictionary with clean-to-description and description-to-code mappings for each level.
    Returns:
        dict: A dictionary where each key is a level, and the value is a tuple containing:
              - clean-to-description mapping (str -> str)
              - description-to-code mapping (str -> str)
    """
    # Get necessary mappings
    crosswalk_dict = {}
    for level in levels:
        mapping_df = pd.read_csv(f"crosswalk/mapping_csvs/{code_type}/{level}.csv", dtype=str)
        mapping_df = mapping_df.astype(str)
        level_crosswalk = CrosswalkBuilder().from_df(mapping_df, from_col="Code", to_col="Description")
        clean_to_description = {get_clean_value(x) : x for x in level_crosswalk.mapping.values()}
        description_to_code = level_crosswalk.create_reverse_mapping()
        crosswalk_dict[level] = (clean_to_description, description_to_code)
    return crosswalk_dict  
 def find_containing_ranges(input_range: str, code_mapping: dict) -> list[dict]:
    """
    Find which code ranges in the dictionary contain or overlap with the input range.
    Args:
        input_range (str): Code range in format 'START-END' (e.g., '70010-79999')
        code_mapping (dict): Dictionary containing CPT codes with code ranges as keys and descriptions as values
    Returns:
        list: List of dictionaries containing matching code ranges and descriptions
    """
    # print("finding ranges for", input_range) # comment for debugging
    # Parse the input range
    match = re.match(r'(\d+)-(\d+)', input_range)
    if not match:
        raise ValueError("Invalid input format. Expected format: 'START-END'")
    input_start = int(match.group(1))
    input_end = int(match.group(2))
    # Initialize results list
    results = []
    # Process each code range in the dictionary
    for code_range, description in code_mapping.items():
        # Parse the code range from the dictionary
        range_match = re.match(r'(\d+)-(\d+)', str(code_range))
        if range_match:
            csv_start = int(range_match.group(1))
            csv_end = int(range_match.group(2))
            # Check if there's any overlap between ranges
            if max(input_start, csv_start) <= min(input_end, csv_end):
                # The ranges overlap in some way
                results.append({
                    'code_range': code_range,
                    'description': description,
                    'overlap_start': max(input_start, csv_start),
                    'overlap_end': min(input_end, csv_end)
                })
    # If no single range contains all others, return the highest level ranges
    highest_level_ranges = []
    for result in results:
        is_highest = True
        for other in results:
            if result != other and int(result['overlap_start']) >= int(other['overlap_start']) and int(result['overlap_end']) <= int(other['overlap_end']):
                is_highest = False
                break
        if is_highest:
            highest_level_ranges.append(result['code_range'])
    return highest_level_ranges
 def code_primary(service, code_primary_questions, filename):
    """
    Breaks down the direct code information from SERVICE_TERM and REIMB_TERM fields
    using an LLM prompt and adds the results to the provided answer dictionaries.
    Args:
        answer_dicts (list of dict): A list of dictionaries containing service and reimbursement data.
        filename (str): The name of the file being processed, used for LLM processing.
    Returns:
        list of dict: The input answer dictionaries with additional primary code breakout information. These are only the codes that are explicitly written in the Service/Methodology
    """
    if not string_utils.is_empty(service):
        claude_answer_raw = llm_utils.invoke_claude(
                            investment_prompts.CODE_PRIMARY_BREAKOUT(str(service), code_primary_questions),
                            "sonnet_latest",
                            filename
                        )
        code_answer_dict = string_utils.universal_json_load(claude_answer_raw)
    else:
        code_answer_dict = {}
    return code_answer_dict
 def get_proc_crosswalk():
    """
    Loads and concatenates all CSV files in the 'proc_cd' directory, and creates a crosswalk mapping from the 'Code' to 'Description' columns.
    Returns:
        CrosswalkBuilder: A CrosswalkBuilder object containing the mapping from procedure codes to descriptions.
    """
    all_mappings = []
    mapping_dir = "crosswalk/mapping_csvs/proc_cd"
    for filename in os.listdir(mapping_dir):
        if filename.endswith(".csv"):
            mapping_df = pd.read_csv(os.path.join(mapping_dir, filename))
            all_mappings.append(mapping_df)
    proc_df = pd.concat(all_mappings, ignore_index=True)
    return CrosswalkBuilder().from_df(proc_df, from_col="Code", to_col="Description")
 def get_clean_value(service_description):
    """
    Cleans a service description by removing stop words, replacing certain terms, and standardizing acronyms.
    Args:
        service_description (str): The service description string to clean.
    Returns:
        str: The cleaned service description with stop words removed, specific terms replaced, and acronyms expanded.
    """
    stop_words = {"the", "is", "in", "and", "to", "of", "a", "an", "that", "this", "it", "for", "on", "with", "as", "was", "were", "at", "by", "but", "be"}  # Experiment with not removing 'and'
    def remove_stop_words(text):
        words = text.split()
        return " ".join(word for word in words if word.lower() not in stop_words)
    service_clean = str(service_description).lower()
    # Replace "Services" and "Procedures"
    service_clean = service_clean.replace("covered services", "").replace("services", "").replace(
        "procedures", "").replace("covered service", "")
    service_clean = service_clean.replace("DME", "Durable Medical Equipment") # Explore adding other acronyms
    service_clean = remove_stop_words(service_clean)
    # Explore lemmatization here
    return service_clean
 def get_best_match(
        code: str,
        all_dataset: dict,
        target: str,
        model: SentenceTransformer,
        thresholds: dict[str, float]={"upper": 0.9, "lower": 0.9},
        grace: float = 0.1
        ): #TODO: fix return type; sometimes it's tuple[str, str, float]; but if no valid match is found it's tuple[None, None, 0]
    """
    Finds the best matching code from a crosswalk dictionary by comparing a target string to precomputed embeddings
    using a model and similarity threshold values for different levels.
    Args:
        code (str): code for which nearest description is to be found. e.g. proc_cd
        all_dataset (dict): Dictionary containing the crosswalk data for different levels for different codes.
        target (str): The target string to match.
        model: The model used to encode the target string into an embedding.
        thresholds (dict, optional): A dictionary containing upper threshold and lower threshold. If any level has score above 
        upper threshold, the match corresponding to that level will be picked. If no level has value above upper threshold, level with 
        higher weighted score (and above lower threshold score) will be picked
    Returns:
        tuple: A tuple containing the best matching code, its description, and the similarity score.
               If no match is found above the threshold, returns (None, None, 0).
    """
    target = get_clean_value(target)
    target_vec = model.encode([target], normalize_embeddings=True).astype("float32")
    #create list of tuples for level, best match text and score
    match_list = []
    max_score = 0
    weighted_max_score = (None, 0)
    if code == 'proc_cd':
        levels = investment_values.proc_levels
        crosswalk_dict = all_dataset['proc_codes_crosswalk']
    else:
        levels = investment_values.diag_levels
        crosswalk_dict = all_dataset['diag_codes_crosswalk']
    for level in levels:
        index, _, choices = embedding_utils.load_faiss_index(index_path=os.path.join("embeddings", level, "faiss_index.bin"),
                                                             embedding_path=os.path.join("embeddings", level, "embeddings.npy"),
                                                             choices_path=os.path.join("embeddings", level, "choices.pkl"),)
        similarity_scores, match_indices = index.search(target_vec, 1)
        best_match, best_score = choices[match_indices[0][0]], similarity_scores[0][0]
        match_list.append((level, best_score, best_match))
        # calculate and save best_score and weighted_best_score (in weighted_best_score, higher levels have some grace scores added)
        if best_score > max_score:
            max_score = best_score
        if best_score > weighted_max_score[1] and level in ["cpt", "hcpcs", "diag"]:
            weighted_max_score = (level, best_score)
        elif best_score + grace > weighted_max_score[1] and level in ["cpt_level3", "hcpcs_level2"]:
            weighted_max_score = (level, best_score + grace)
        elif best_score + 2*grace > weighted_max_score[1] and level in ["cpt_level2", "hcpcs_level1"]:
            weighted_max_score = (level, best_score + 2*grace)
        elif best_score + 3*grace > weighted_max_score[1] and level in ["cpt_level1"]:
            weighted_max_score = (level, best_score + 3*grace)
    for match in match_list:
        level, score, match_text = match
        if max_score >= thresholds["upper"] and score == max_score:
            best_description = crosswalk_dict[level][0][match_text]
            best_code = crosswalk_dict[level][1][best_description]
            return best_code, best_description, best_score
        elif max_score < thresholds["upper"] and weighted_max_score[1] >= thresholds["lower"] and level == weighted_max_score[0]:
            best_description = crosswalk_dict[level][0][match_text]
            best_code = crosswalk_dict[level][1][best_description]
            return best_code, best_description, best_score
    return None, None, 0 # No valid match found
 def code_implicit(service, filename):
    """
    Maps free-form service descriptions to the best-matching code and description. 
    Caches results for efficiency and skips entries with existing codes or irrelevant text.
    ONLY FOR PROC CODES RN
    Args:
        answer_dicts (list): List of dictionaries containing 1:N answers
        all_dataset (dict): dict containing all crosswalk and mapping.
    Returns:
        list: Updated list of dictionaries with mapped CPT4 procedure codes and descriptions.
    """
    if not string_utils.is_empty(service):
        claude_answer_raw = llm_utils.invoke_claude(investment_prompts.CODE_IMPLICIT(service), "sonnet_latest", filename)
        code_answer_final = string_utils.universal_json_load(claude_answer_raw)
    else:
        code_answer_final = ""
    return code_answer_final
 def get_code_description(code, code_mapping):
    """
    Retrieves descriptions for a given code or list of codes using the provided code mapping.
    If no description is found, attempts to resolve by modifying the code (e.g., prepending "0" or handling code patterns).
    Args:
        code (str, list): A single code or a list of codes to look up.
        code_mapping (Crosswalk): The mapping object containing code-to-description mappings.
    Returns:
        str: A comma-separated string of descriptions for the provided codes.
    """
    if isinstance(code, str):
        if string_utils.is_empty(code):
            return "N/A"
        elif "[" in code and "]" in code:
            code_list = ast.literal_eval(code)
        elif "," in code:
            code_list = code.split(",")      
        else:
            code_list = [code] # added this line to handle range strings like 'J0120-J7175' or "70010-79999"
    elif isinstance(code, list):
        code_list = code
    else:
        code_list = [code]
    code_description = []
    for code in code_list:
        code_str = str(code).strip()
        description = code_mapping.mapping.get(code_str, "mismap")
        if string_utils.is_empty(description) or description == "mismap":
            if code_str.isdigit():
                description = code_mapping.mapping.get(int(code_str), "mismap")
                if description == "mismap":
                    description = code_mapping.mapping.get(str(int(code_str)), "mismap")
            else:
                description = "mismap"
        # If mismap, try to handle ranges
        if string_utils.is_empty(description) or description == "mismap":
            if "X" in code_str:
                base_code = code_str.replace("X", "")
                expanded_code_range = f"{base_code}0-{base_code}9"
                description = code_mapping.mapping.get(expanded_code_range, "mismap")
                # If mismap, try to handle numeric codes
                if string_utils.is_empty(description) or description == "mismap":
                    try:
                        x_range_code_list = [int(f"{base_code}{i}") for i in range(10)]
                        description = get_code_description(x_range_code_list, code_mapping)
                    except Exception as e:
                        print(f"Error message: {e}")
                        description = "N/A"
            elif "-" in code_str and all(part.isdigit() for part in code_str.split("-")):
                try:
                    range_code_list = find_containing_ranges(code_str, code_mapping.mapping)
                    description = get_code_description(range_code_list, code_mapping)
                except ValueError as e:
                    print(f"Failed to parse code range: {code_str}; error {e}")
                    description = "mismap"
                # # If mismap, try to handle just starting range value and ending range value separately
                # if string_utils.is_empty(description) or description == "mismap":
                #     description = get_code_description(code_str.split("-"), code_mapping)
            else:
                # If mismap, prepend "0" and try again (useful for rev codes)
                code_str = "0" + code_str
                description = code_mapping.mapping.get(code_str, "mismap")
        if not string_utils.is_empty(description):
            if isinstance(description, list):
                code_description.extend(description)
            else:
                code_description.append(description)
    # Remove repeated values and retain unique values, excluding "mismap" if other values are present
    unique_descriptions = list(set(code_description))
    if "mismap" in unique_descriptions and len(unique_descriptions) > 1:
        unique_descriptions.remove("mismap")
    code_description = unique_descriptions
    return ",".join(code_description)
 def get_mappings():
    """
    Loads and processes multiple CSV and Excel files containing code mappings for various categories
    (e.g., PROC, DIAG, Grouper, Rev, Admit Type), and returns a dictionary of crosswalks for each code type.
    Returns:
        dict: A dictionary containing mappings for PROC, DIAG, Grouper, Rev, and Admit Type codes.
    """
    all_mappings = {}
    # PROC codes
    mapping_path = "crosswalk/mapping_csvs/"
    all_proc_mappings = pd.concat([pd.read_csv(os.path.join(mapping_path, "proc_cd", file), dtype=str) for file in os.listdir(os.path.join(mapping_path, "proc_cd")) if file.endswith(".csv")])
    proc_mapping = CrosswalkBuilder().from_df(all_proc_mappings, from_col="Code", to_col="Description")
    all_mappings['proc_mapping'] = proc_mapping
    # Diag codes
    all_diag_mappings = pd.concat([pd.read_csv(os.path.join(mapping_path, "diag_cd", file), dtype=str) for file in os.listdir(os.path.join(mapping_path, "diag_cd")) if file.endswith(".csv")])
    diag_mapping = CrosswalkBuilder().from_df(all_diag_mappings, from_col="Code", to_col="Description")
    all_mappings['diag_mapping'] = diag_mapping
    # Grouper codes
    all_grouper_mappings = pd.concat([pd.read_csv(os.path.join(mapping_path, "grouper_cd", file), dtype=str) for file in os.listdir(os.path.join(mapping_path, "grouper_cd")) if file.endswith(".csv")])
    grouper_mapping_ms = CrosswalkBuilder().from_df(all_grouper_mappings[all_grouper_mappings['Version']=='MS-DRG'], from_col="Code", to_col="Description")
    grouper_mapping_apr = CrosswalkBuilder().from_df(all_grouper_mappings[all_grouper_mappings['Version']=='APR-DRG'], from_col="Code", to_col="Description")
    all_mappings['grouper_mapping_ms'] = grouper_mapping_ms
    all_mappings['grouper_mapping_apr'] = grouper_mapping_apr
    # Rev codes
    rev_mapping = CrosswalkBuilder().from_excel(path=os.path.join(mapping_path, "rev_cd", "rev_mapping.csv"), from_col="Code", to_col="Description")
    all_mappings['rev_mapping'] = rev_mapping
    # Admit Type codes
    admit_mapping = CrosswalkBuilder().from_excel(path=os.path.join(mapping_path, "admit_type_cd", "admit_mapping.csv"), from_col="Code", to_col="Description")
    all_mappings['admit_mapping'] = admit_mapping
    return all_mappings
 def code_description_mapping(answer_dict, all_mappings):
    """
    Maps code descriptions to relevant fields in the answer dictionaries using provided mappings.
    Args:
        answer_dicts (list): List of dictionaries containing codes to be mapped.
        all_mappings (dict): Dictionary containing the mappings for different code types.
    Returns:
        list: Updated list of dictionaries with code descriptions added.
    """
    proc_code = answer_dict.get("PROCEDURE_CD")
    proc_mod = answer_dict.get("CPT4_PROC_MOD")
    diag_code = answer_dict.get("DIAG_CD")
    rev_code = answer_dict.get("REVENUE_CD")
    grouper_code = answer_dict.get("GROUPER_CD")
    grouper_type = answer_dict.get("GROUPER_TYPE")
    admit_type_code = answer_dict.get("CLAIM_ADMIT_TYPE_CD")
    answer_dict["PROCEDURE_CD_DESC"] = get_code_description(proc_code, all_mappings['proc_mapping'])
    answer_dict["CPT4_PROC_MOD_DESC"] = get_code_description(proc_mod, all_mappings['proc_mapping'])
    answer_dict["DIAG_CD_DESC"] = get_code_description(diag_code, all_mappings['diag_mapping'])
    answer_dict["REVENUE_CD_DESC"] = get_code_description(rev_code, all_mappings['rev_mapping'])
    if grouper_type == 'MS-DRG':
        answer_dict["GROUPER_CD_DESC"] = get_code_description(grouper_code, all_mappings['grouper_mapping_ms'])
    elif grouper_type == 'APR-DRG':
        answer_dict["GROUPER_CD_DESC"] = get_code_description(grouper_code, all_mappings['grouper_mapping_apr'])
    elif string_utils.is_empty(grouper_code):
        answer_dict["GROUPER_CD_DESC"] = "N/A"  # No grouper code exists
    else:
        answer_dict["GROUPER_CD_DESC"] = "UNKNOWN"  # Grouper code exists but type isn't recognized
    answer_dict["AUTH_ADMIT_TYPE_DESC"] = get_code_description(admit_type_code, all_mappings['admit_mapping'])
    return answer_dict
 def clean_service_term(service_term: str,
                       LOBS_TO_REMOVE,
                       PROGRAMS_TO_REMOVE,
                       NETWORKS_TO_REMOVE) -> str:
    if string_utils.is_empty(service_term):
        return str(service_term)
    service_term = service_term.upper()
    # replace whole words from LOBS_TO_REMOVE with ""
    for lob in LOBS_TO_REMOVE:
        service_term = re.sub(r'\b' + re.escape(lob.upper()) + r'\b', '', service_term)
    # replace whole words from PROGRAMS_TO_REMOVE with ""
    for program in PROGRAMS_TO_REMOVE:
        if program not in ["HOSPICE"]:
            service_term = re.sub(r'\b' + re.escape(program.upper()) + r'\b', '', service_term)
    # replace whole words from NETWORKS_TO_REMOVE with ""
    for network in NETWORKS_TO_REMOVE:
        service_term = re.sub(r'\b' + re.escape(network.upper()) + r'\b', '', service_term)
    # replace specific phrases
    for val in ["INPATIENT", "OUTPATIENT", 
                "PROFESSIONAL", "INSITUTIONAL", "FACILITY", "ANCILLARY", "HOSPITAL", "HEALTH"]:
        service_term = re.sub(r'\b' + re.escape(val) + r'\b', '', service_term)
    # remove extra spaces
    service_term = re.sub(r'\s+', ' ', service_term).strip()
    return service_term
 def get_code_breakout(answer_dicts: list[dict], filename: str, all_dataset: dict):
    """
    Extracts and maps codes from contract data, including primary, indirect, and description mappings.
    Args:
        answer_dicts (list[dict]): List of dictionaries containing contract-related answers.
        filename (str): The name of the file being processed.
        all_dataset (dict): Dictionary containing crosswalk and mappings for code translation.
    Returns:
        dict: A dictionary with primary, indirect/derived, and mapped code descriptions.
    """
    LOBS_TO_REMOVE = list(set(CrosswalkBuilder().from_json(path="crosswalk/mappings/crosswalk_lob.json").mapping.values()))
    PROGRAMS_TO_REMOVE = list(set(CrosswalkBuilder().from_json(path="crosswalk/mappings/crosswalk_program.json").mapping.values()))
    NETWORKS_TO_REMOVE = list(set(CrosswalkBuilder().from_json(path="crosswalk/mappings/crosswalk_network.json").mapping.values()))
    code_primary_questions = FieldSet(file_path=config.FIELD_JSON_PATH, field_type="code_primary_breakout").get_prompt_dict()
    final_answer_dicts = []
    for answer_dict in answer_dicts:
        service_term = answer_dict.get("SERVICE_TERM", "")
        service_term_clean = clean_service_term(service_term,
                                                LOBS_TO_REMOVE,
                                                PROGRAMS_TO_REMOVE,
                                                NETWORKS_TO_REMOVE)
        if service_term_clean in ['COVERED SERVICES', 'ALL COVERED SERVICES']:
            final_answer_dicts.append(answer_dict)
            continue
        # Get Explicit Codes
        code_primary_answers = code_primary(service_term, code_primary_questions, filename)
        answer_dict.update(code_primary_answers)
        if all(string_utils.is_empty(value) for value in code_primary_answers.values()):
            implicit_proc_cd = code_implicit(service_term_clean, filename)
            answer_dict["PROCEDURE_CD"] = implicit_proc_cd
        # Code Descriptions
        try:
            answer_dict = code_description_mapping(answer_dict, all_dataset['codes_mappings'])
        except Exception as e:
            print(f"Error in code_description_mapping: {e}")
            answer_dict["PROCEDURE_CD_DESC"] = "N/A"
            answer_dict["CPT4_PROC_MOD_DESC"] = "N/A"
            answer_dict["DIAG_CD_DESC"] = "N/A"
            answer_dict["REVENUE_CD_DESC"] = "N/A"
            answer_dict["GROUPER_CD_DESC"] = "N/A"
            answer_dict["AUTH_ADMIT_TYPE_DESC"] = "N/A"
        final_answer_dicts.append(answer_dict)
    return final_answer_dicts
@@ -13,20 +13,18 @@ import src.investment.tin_npi_funcs as tin_npi_funcs
 import src.investment.postprocess as postprocess
 import src.investment.preprocess as preprocess
 import src.investment.row_funcs as row_funcs
-import src.investment.smart_chunking_funcs as smart_chunking_funcs
+import src.codes.code_funcs as code_funcs
 import src.utils.io_utils as io_utils
 import src.utils.string_utils as string_utils
 from src import config
 from src.config import FIELD_JSON_PATH
 from src.investment.one_to_n_funcs import (combine_one_to_n_answers,
                                           reimbursement_level)
-from src.investment.tin_npi_funcs import (merge_provider_info,
+from src.prompts.investment_prompts import FieldSet
                                          reimbursement_tin_npi)
 from src.prompts.investment_prompts import Field, FieldSet
 from src.utils.string_utils import datetime_str
-def process_file(file_object, all_dataset, run_timestamp):
+def process_file(file_object, run_timestamp):
    filename, contract_text = file_object
    print(f"{datetime_str()} Processing {filename}...")
@@ -39,7 +37,7 @@ def process_file(file_object, all_dataset, run_timestamp):
    ################## ONE TO N ##################
    if string_utils.contains_reimbursement(contract_text):
-        one_to_n_results, dynamic_one_to_one_fields = run_one_to_n_prompts(filename, exhibit_dict, all_exhibit_headers, all_dataset) # Return df
+        one_to_n_results, dynamic_one_to_one_fields = run_one_to_n_prompts(filename, exhibit_dict, all_exhibit_headers) # Return df
        one_to_n_results['FILE_NAME'] = filename
        one_to_n_results = investment_postprocessing_funcs.generate_reimb_ids(one_to_n_results) # Add reimb_id
        print(f"{datetime_str()} One to N Complete - {filename}")
@@ -54,7 +52,10 @@ def process_file(file_object, all_dataset, run_timestamp):
    print(f"{datetime_str()} One to One Complete - {filename}")
    ################## MERGE ##################
-    final_results = row_funcs.merge_one_to_one_into_one_to_n(one_to_n_results, one_to_one_results)
+    merged_results = row_funcs.merge_one_to_one_into_one_to_n(one_to_n_results, one_to_one_results)
    ################## CODES ##################
    final_results = code_funcs.code_breakout(merged_results)
    ################## POSTPROCESS ##################
    final_df = postprocess.postprocess(final_results)
@@ -101,7 +102,7 @@ def run_one_to_one_prompts(filename, contract_text, text_dict, top_sheet_dict, d
    return one_to_one_results[0]
-def run_one_to_n_prompts(filename, exhibit_dict, all_exhibit_headers, all_dataset):
+def run_one_to_n_prompts(filename, exhibit_dict, all_exhibit_headers):
    ################## RUN PROMPTS ##################
    one_to_n_results = []
@@ -154,7 +155,7 @@ def run_one_to_n_prompts(filename, exhibit_dict, all_exhibit_headers, all_datase
        reimbursement_level_fields.combine(code_to_reimbursement_level_fields, inplace=True)
        ################## GET REIMBURSEMENT-LEVEL ANSWERS (INCLUDING DYNAMIC) ##################
-        reimbursement_level_answers = reimbursement_level(exhibit_text, filename, reimbursement_level_fields, all_dataset, exhibit_page, seen_pairs, exhibit_lesser_of=exhibit_lesser_of) # Return list of dictionaries
+        reimbursement_level_answers = reimbursement_level(exhibit_text, filename, reimbursement_level_fields, exhibit_page, seen_pairs, exhibit_lesser_of=exhibit_lesser_of) # Return list of dictionaries
        ################# COMBINE ANSWERS ##################
        full_answer_dict = combine_one_to_n_answers(exhibit_level_answers, reimbursement_level_answers, tin_npi_answers={})
@@ -9,17 +9,14 @@ import pandas as pd
 random.seed(42)
 # from src.tracking.batch_tracking import BatchTracker
 import traceback
 import src.investment.file_processing as file_processing
 import src.tracking.tracking_utils as tracking_utils
 import src.utils.io_utils as io_utils
 import src.utils.string_utils as string_utils
 from src import config
-def safe_process_file(item, all_dataset, run_timestamp):
+def safe_process_file(item, run_timestamp):
    """call process_file and catch any exceptions that occur
    Args:
@@ -34,33 +31,11 @@ def safe_process_file(item, all_dataset, run_timestamp):
        item[0] if isinstance(item, (list, tuple)) and len(item) > 0 else item
    )  # unpack file_id from item (passed in as a tuple below)
    try:
-        return file_processing.process_file(item, all_dataset, run_timestamp)
+        return file_processing.process_file(item, run_timestamp)
-    except (
+    except Exception as e: # When there's an issue with the processing inside the future
-        Exception
+        print(f"Error processing file {item[0]}: {str(e)}")
-    ) as e:  # When there's an issue with the processing inside the future
+        print(traceback.format_exc())
-        error_type = type(e).__name__
+        return pd.DataFrame([{"error": str(e), "file_id": item[0]}]) # Return a single-row dataframe so we can still concat it later
        error_message = str(e)
        full_traceback = traceback.format_exc()
        logging.error(f"Error processing file {file_id}: {error_message}")
        logging.error(f"Error type: {error_type}")
        logging.error(f"Error Message: {error_message}")
        logging.error(f"Full traceback:\n{full_traceback}")
        return pd.DataFrame(
            [
                {
                    "file_id": file_id,
                    "error": error_message.replace(",", ";").replace(
                        "\n", " "
                    ),  # Replace problematic characters for CSV compatibility
                    "error_type": error_type,
                    "traceback": full_traceback.replace(",", ";").replace(
                        "\n", " | "
                    ),  # keep line breaks as separators
                }
            ]
        )  # Return a single-row dataframe so we can still concat it later
 def main(testing=False, test_params={}):
@@ -80,21 +55,15 @@ def main(testing=False, test_params={}):
    print(f"Total Input Files : {len(input_dict)}")
-    # load all crosswalk files, mappings and embeddings
+    
    all_dataset = io_utils.load_all_dataset()
    io_utils.load_embeddings()
    # Process files with run_timestamp
    successful_results = []
    error_results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        # use executor to process files concurrently; use submit() instead of map()
        # so that we can catch exceptions and continue processing other files
-        futures = [
+        futures = [executor.submit(safe_process_file, item, run_timestamp) for item in input_dict.items()]
-            executor.submit(safe_process_file, item, all_dataset, run_timestamp)
+        
            for item in input_dict.items()
        ]
        # collect results as they complete
        for future in concurrent.futures.as_completed(futures):
            try:
@@ -4,7 +4,7 @@ import re
 from typing import Callable
 import src.constants.investment_values as investment_values
-import src.investment.code_funcs as code_funcs
+import src.codes.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
@@ -475,7 +475,6 @@ def reimbursement_level(
    exhibit_text,
    filename,
    reimbursement_level_fields,
    all_dataset,
    exhibit_page: str,
    seen_pairs: set,
    exhibit_lesser_of: str
@@ -528,10 +527,7 @@ def reimbursement_level(
        if deduplicated_answers:  # Only process if we have valid reimbursement pairs
            # This will also run the methodology breakout prompt for non-carveout rows
            carveout_answers = get_special_cases(deduplicated_answers, filename)
-            code_breakout_answers = code_funcs.get_code_breakout(
+            return carveout_answers
                carveout_answers, filename, all_dataset
            )
            return code_breakout_answers
        else:
            return []  # All pairs were duplicates, return empty list
@@ -538,7 +538,7 @@
        "field_name": "NDC_CD",
        "relationship": "one_to_n",
        "field_type": "code_primary_breakout",
-        "prompt": "Return a list of any National Drug Code numbers present. Look for an 11-digit number of the format XXXXX-XXXX-XX."
+        "prompt": "Identify National Drug Code (NDC) numbers present. Look for an 11-digit number of the format XXXXX-XXXX-XX."
    },
    {
        "field_name": "LINE_NDC_CD_DESC",
@@ -550,7 +550,7 @@
        "field_name": "CLAIM_ADMIT_TYPE_CD",
        "relationship": "one_to_n",
        "field_type": "code_primary_breakout",
-        "prompt": "Identify admit type codes in the text: - Single digit numeric codes (1-9) - Used to indicate hospital admission type"
+        "prompt": "Identify Admit Type codes in the text: - Single digit numeric codes (1-9) - Used to indicate hospital admission type."
    },
    {
        "field_name": "AUTH_ADMIT_TYPE_DESC",
@@ -561,8 +561,8 @@
    {
        "field_name": "CLAIM_STATUS_CD",
        "relationship": "one_to_n",
-        "field_type": "TBD",
+        "field_type": "code_primary_breakout",
-        "prompt": "TBD"
+        "prompt": "Identify Claim Status codes in the text. Look for numbers between 1-999 that refer specifically to the status of a claim being submitted."
    },
    {
        "field_name": "CLAIM_STATUS_CD_DESC",
@@ -657,17 +657,19 @@ Methodology: {methodology.replace('"', "'")}
 Write your JSON dictionary below:
 """
-def CODE_PRIMARY_BREAKOUT(service, questions):
+def CODE_EXPLICIT(service, questions):
    return f"""Analyze a given medical service:
 Use the text in the Service to populate a JSON dictionary with the following fields:
 {questions}.
-For each field, return a list of all codes EXPLICITLY written for that field. The code itself must be written, not language that describes the code.
+- For each field, return a list of all codes EXPLICITLY written for that field. The code itself must be written, not language that describes the code.
-Note that codes may be present, but not explicitly labeled as codes. For example, you may simply see "J1098", which is a PROCEDURE_CD. You may see "155" which is a REVENUE_CD, etc.
+- Note that codes may be present, but not explicitly labeled as codes. For example, you may simply see "J1098", which is a PROCEDURE_CD. You may see "155" which is a REVENUE_CD, etc.
-If the contract gives a range of codes, without listing each code in the range individually, return the range in the format 'LowestCode-HighestCode'.
+- If the text gives a range of codes, without listing each code in the range individually, return the range in the format 'LowestCode-HighestCode'.
-If any of the codes are not found, do not write a list for that field. Simply populate the field with 'N/A'.
+- If the text describes an entire Code Category (not one specific code) based on a start letter, return the Category in the form "Category: X". e.g. "A-Codes" --> "Category: A", "K-Codes" --> "Category K", etc.
 - If the text explicitly says that there is no published or established rate, write "NOT_ESTABLISHED" for the PROCEDURE_CD value.
 - If any of the codes are not found, do not write a list for that field. Simply populate the field with an empty list [].
 Here is the Service to analyze:
 Service: {service.replace('"', "'")}
@@ -675,21 +677,114 @@ Service: {service.replace('"', "'")}
 Briefly explain your answer before putting the final answer in JSON dictionary format.
 """
 def CODE_IMPLICIT(service):
    return f"""Analyze a given medical Service.
-What Procedure code or code range the Service does the Service describe?
+def CODE_CATEGORY(service, choices):
    return f"""Analyze a given medical Service Term.
-If the Service is broad enough to describe a range of codes, return the range in the format 'LowestCode-HighestCode'. The code descriptions for ranges may be very broad (for general services like 'Surgery', 'Radiology', etc.) or more specific (for services like 'Surgical Procedures on the Shoulder', 'Diagnostic Radiology Procedures of the Heart', etc.).
+What Procedure Code description does the Service Term describe?
 If the Service is specific, and therefore references a single code, return just that code.
 If the Service is not specific enough to identify a code, return 'N/A'.
-Ensure that there is no meaningful difference between the description of the code you provide and the service you are analyzing. These codes will be used to price healthcare claims, so they must be accurate. Do not attempt to force a code to fit the service if it will result in different pricing.
+[VALID DESCRIPTIONS]
 Here are the descriptions to choose from:
 {choices}
-Here is the Service to analyze:
+[INSTRUCTIONS]
 1. Determine if the descriptions are needed. If the Service Term mentions the Code Category (a single-character e.g. A, B, C, ...), without any additional subcategories or descriptors, then there is no need to reference the descriptions. For example, if the Service Term is "A-Codes", the answer is "A0000-A9999". This represents the full range of A-Codes.
 2. If there are additional subcategories or descriptors accompanying the code category, match them to the closest item or items from the VALID DESCRIPTIONS.
    - There can be multiple correct answers. When this is the case, return them all in a list. This may be phrased like "Service1/Service2" or "Service1 and Service2", in which case both Services must be considered.
    - There may be no correct answers. When this is the case, return the full range of codes in that category (e.g. A0000-A9999, etc.)
 [CONTEXT]
 Here is the Service Term to analyze:
 {service.replace('"', "'")}
-Explain your answer in 1-2 sentences. Then return your final answer in JSON list-of-string format.
+[OUTPUT FORMAT]
 Explain your answer in 1-2 sentences. Write your final answer in JSON list format.
 """
 def CODE_IMPLICIT_SPECIAL(service):
    return f"""Analyze a given medical Service Term.
 [INSTRUCTIONS]
 - If the Service refers to the broad categories of Drugs, Medications, Injectible Medications, Pharmaceuticals or similar, without referring to a more specific drug category, respond "Drugs". Do NOT use this category if vaccines are included.
 - If the Service refers to Drugs and/or Vaccines, Immunizations, or similar, without referring to a more specific drug or vaccine category, respond "Vaccines".
 - If the Service refers to Physical, Occupational, and Speech Therapy (aka PT/OT/ST), respond "PT/OT/ST". All three must be present to be considered in this group.
    - If the Service refers to just one of these 3 (PT, OT, or ST), respond with "PT", "OT", or "ST"
 - If the Service refers to Multiple procedures, bilateral procedures, or similar, without referring to a single specific procedure or type of procedure, respond "Surgery"
 - If none of the above cases apply, respond "N/A". If any of the above are mentioned but only in the context of being excluded from the service, respond "N/A".
 Here is the Service to analyze:
 {service}
 Briefly explain your answer, but put your final answer in |pipes|
 """
 def CODE_IMPLICIT(service, choices):
    return f"""Analyze a given medical Service Term.
 What Procedure Code description does the Service Term describe?
 [VALID DESCRIPTIONS]
 Here are the descriptions to choose from:
 {choices}
 [INSTRUCTIONS]
 1. Determine which part of the Service Term is the Service itself. The term may include additional confusing information that is irrelevant. Here are some things to look out for:
    - The Provider Type is NOT the Service being provided. For example, in "Medical Services provided by an Anesthesiologist", the Service part is simply "Medical Services". The fact that they are provided by an Anesthesiologist is irrelevant.
    - The Place of Service is NOT the Service being provided. For example, in "Anesthesia Services provided in an Outpatient Laboratory", the Service part is simply "Anesthesia Services". The fact that they are provided in an Outpatient Laboratory is irrelevant.
    - Any other Line of Business (e.g. Medicare, Medicaid) or Program (e.g. CHIP, STAR) information is not relevant to the Service.
    - There may be multiple Services within the Service term. This may be phrased like "Service1/Service2" or "Service1 and Service2", in which case both Services must be considered.
 2. Once you have isolated the Service part of the text and determined that it is not generic, match it to the corresponding item or items from the VALID DESCRIPTIONS. Here's how to do it:
    - IGNORE distracting terms like "Services" or "Procedures" in both the Service and the Description. They are meaningless for the purpose of this exercise.
    - For each Description in VALID DESCRIPTIONS, ask yourself the following question: Is the Description SYNONYMOUS with the Service?
    - **Note**: Do NOT simply categorize the Service into the Description that it falls under. You are looking specifically for synonymous Service-Description matches. For example, "Nutritional Evaluations" is not synonymous with "Evaluations", because not ALL "Evaluations" are "Nutritional Evaluations".
    - There can be multiple correct answers. When this is the case, return them all in a list. 
    - There may be no correct answers, especially when the Service is much more specific than any of the descriptions. When this is the case, simply return an empty list "[]".
 [CONTEXT]
 Here is the Service Term to analyze:
 {service.replace('"', "'")}
 [OUTPUT FORMAT]
 Explain your answer, ensuring each point in the instructions is addressed. Then return your final answer in JSON list format.
 """
 def CODE_LAST_CHECK(service):
    return f"""Analyze the given healthcare-related term.
 Determine if the term is a specific medical service or category of services. If so, return "Specific".
 If the term is a generic medical service, not a service at all, the name of a hospital or other provider, or other nonsensical or non-service phrase, return "Generic".
 Here is the service to analyze: {service}
 Briefly explain your answer, then put your final answer in |pipes|.
 """
 def FILL_BILL_TYPE(service, choices):
    return f"""Analyze a given medical Service Term.
 What Bill Type Code Description does the Service Term fall under?
 [VALID DESCRIPTIONS]
 Here are the descriptions to choose from:
 {choices}
 [INSTRUCTIONS]
 - Answer with the Description or Descriptions from [VALID DESCRIPTIONS] only if it is clearly and unambiguously referred to by the Service. Do not attempt much inference, the answer will be very clear if it is present.
 - It is highly likely that none of the Descriptions will be a clear and unambiguous match. When this is the case, return an empty list: []
 - Here are some examples:
  - "Ambulatory Surgery Procedures" --> ["Ambulatory Surgery Center"]
  - "Hospital Services" --> ["Inpatient Hospital", "Outpatient Hospital"]
 [CONTEXT]
 Here is the Service Term to analyze:
 {service.replace('"', "'")}
 [OUTPUT FORMAT]
 Briefly explain your answer, then return your final answer in JSON list format.
 """
@@ -1,8 +1,11 @@
 import faiss
 import pickle
 import numpy as np
 import os
 import src.config as config
 from src.utils.string_utils import datetime_str
 def create_faiss_index(choices, model, save_path="faiss_index.bin", embedding_path="embeddings.npy", choices_path="choices.pkl"):
    embeddings = model.encode(choices, normalize_embeddings=True).astype("float32")
@@ -33,3 +36,40 @@ def load_faiss_index(index_path="faiss_index.bin", embedding_path="embeddings.np
 def load_embeddings():
    """loads all embedding files required for processing indirect codes (e.g. proc codes)
    """
    # list all embeddings present in s3
    s3_client = config.S3_CLIENT
    response = s3_client.list_objects_v2(Bucket="doczy-investment", Prefix="embeddings")
    file_list = [
        obj["Key"]
        for obj in response.get("Contents", [])
        if not obj["Key"].endswith("/")
    ]
    # list embeddings already present locally
    if not os.path.exists("embeddings"):
        os.makedirs("embeddings")
    local_files = []
    for root, dirs, files in os.walk("embeddings"):
        for name in files:
            local_files.append(os.path.join(root, name))
    # download embeddings not present locally
    for file in file_list:
        local_path = file
        # create directory structure for this file
        directory = os.path.dirname(local_path)
        if not os.path.exists(directory):
            os.makedirs(directory)
        # check if file already exists
        if not os.path.exists(local_path):
            try:
                print(f"{datetime_str()} Downloading {file} to {local_path}...")
                s3_client.download_file("doczy-investment", file, local_path)
            except Exception as e:
                print(f"{datetime_str()} Error downloading {file}: {e}")
@@ -7,9 +7,7 @@ from pyxlsb import open_workbook
 import src.tracking.tracking_utils as tracking_utils
 from src import config
-from src.investment.code_funcs import crosswalk_levels, get_mappings
+
 import src.constants.investment_values as investment_values
 from src.utils.string_utils import datetime_str
 import tempfile
@@ -296,51 +294,3 @@ def write_s3(df, filename, run_timestamp, output_type):
        print(f"Saved individual output file: {output_path}")
 def load_all_dataset():
    """_summary_: loads all csv files required for processing contracts including crosswalk files
    and mappings (embeddings, past results may be added later)
    """
    all_dataset = {}
    all_dataset['proc_codes_crosswalk'] = crosswalk_levels("proc_cd", investment_values.proc_levels)
    all_dataset['diag_codes_crosswalk'] = crosswalk_levels("diag_cd", investment_values.diag_levels)
    all_dataset['codes_mappings'] = get_mappings()
    return all_dataset
 def load_embeddings():
    """loads all embedding files required for processing indirect codes (e.g. proc codes)
    """
    # list all embeddings present in s3
    s3_client = config.S3_CLIENT
    response = s3_client.list_objects_v2(Bucket="doczy-investment", Prefix="embeddings")
    file_list = [
        obj["Key"]
        for obj in response.get("Contents", [])
        if not obj["Key"].endswith("/")
    ]
    # list embeddings already present locally
    if not os.path.exists("embeddings"):
        os.makedirs("embeddings")
    local_files = []
    for root, dirs, files in os.walk("embeddings"):
        for name in files:
            local_files.append(os.path.join(root, name))
    # download embeddings not present locally
    for file in file_list:
        local_path = file
        # create directory structure for this file
        directory = os.path.dirname(local_path)
        if not os.path.exists(directory):
            os.makedirs(directory)
        # check if file already exists
        if not os.path.exists(local_path):
            try:
                print(f"{datetime_str()} Downloading {file} to {local_path}...")
                s3_client.download_file("doczy-investment", file, local_path)
            except Exception as e:
                print(f"{datetime_str()} Error downloading {file}: {e}")
@@ -131,32 +131,26 @@ class TestOneToN(unittest.TestCase):
    @patch('src.investment.one_to_n_funcs.get_reimbursement_primary')
    @patch('src.investment.one_to_n_funcs.get_special_cases')
    @patch('src.investment.code_funcs.get_code_breakout')
    @patch('src.investment.one_to_n_funcs.validate_reimbursements_for_llm')
-    def test_reimbursement_level(self, mock_validate_llm, mock_code_breakout,
+    def test_reimbursement_level(self, mock_validate_llm, mock_special_cases, mock_reimbursement_primary):
                                mock_special_cases, mock_reimbursement_primary):
        # Setup
        mock_reimbursement_primary.return_value = [{"SERVICE_TERM": "Test Service", "REIMB_TERM": "term1$"}]
        mock_special_cases.return_value = [{"SERVICE_TERM": "Test Service", "REIMB_TERM": "term1$", "CARVEOUT_IND": "N"}]
        mock_code_breakout.return_value = [{"SERVICE_TERM": "Test Service", "REIMB_TERM": "term1$", "CODE": "code1"}]
        mock_validate_llm.return_value = True  # Mock all entries pass LLM validation
        mock_fields = MagicMock()
        mock_dataset = {}
        exhibit_page = "1"
        seen_pairs = set()
        # Execute
        result = one_to_n_funcs.reimbursement_level(
-            self.exhibit_text, self.filename, mock_fields, mock_dataset, exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
+            self.exhibit_text, self.filename, mock_fields, exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
        )
        # Assert
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["CODE"], "code1")
        mock_reimbursement_primary.assert_called_once()
        mock_special_cases.assert_called_once()
        mock_code_breakout.assert_called_once()
 if __name__ == '__main__':
    unittest.main()
@@ -1,2 +0,0 @@
 def test_always_passes():
    assert True
@@ -1,59 +0,0 @@
 import pytest
 from unittest.mock import patch, MagicMock
 from src.investment.code_funcs import find_containing_ranges
 def test_find_containing_ranges_valid_input():
    input_range = '70010-79999'
    code_mapping = {
        '70000-70050': 'Description 1',
        '70040-70060': 'Description 2',
        '70010-70020': 'Description 3',
        '80000-80050': 'Description 4'
    }
    expected_output = ['70000-70050', '70040-70060']
    assert find_containing_ranges(input_range, code_mapping) == expected_output
 def test_find_containing_ranges_no_overlap():
    input_range = '70010-70020'
    code_mapping = {
        '70030-70040': 'Description 1',
        '70050-70060': 'Description 2',
        '80000-80050': 'Description 3'
    }
    expected_output = []
    assert find_containing_ranges(input_range, code_mapping) == expected_output
 def test_find_containing_ranges_partial_overlap():
    input_range = '70010-70050'
    code_mapping = {
        '70000-70020': 'Description 1',
        '70040-70060': 'Description 2',
        '70050-70070': 'Description 3'
    }
    expected_output = ['70000-70020', '70040-70060']
    assert find_containing_ranges(input_range, code_mapping) == expected_output
 def test_find_containing_ranges_exact_match():
    input_range = '70000-70050'
    code_mapping = {
        '70000-70050': 'Description 1',
        '70040-70060': 'Description 2',
        '70010-70020': 'Description 3'
    }
    expected_output = ['70000-70050']
    assert find_containing_ranges(input_range, code_mapping) == expected_output
 def test_find_containing_ranges_empty_mapping():
    input_range = '70000-70050'
    code_mapping = {}
    expected_output = []
    assert find_containing_ranges(input_range, code_mapping) == expected_output
 def test_find_containing_ranges_invalid_input():
    input_range = 'invalid-range'
    code_mapping = {
        '70000-70050': 'Description 1',
        '70040-70060': 'Description 2'
    }
    with pytest.raises(ValueError):
        find_containing_ranges(input_range, code_mapping)
@@ -3,9 +3,7 @@ import unittest
 from unittest.mock import patch, MagicMock
 import src.investment.one_to_n_funcs as one_to_n_funcs
-import src.preprocessing_funcs as preprocessing_funcs
+
 import src.prompts.investment_prompts as investment_prompts
 import src.config as config
 class TestOneToN(unittest.TestCase):
    def setUp(self):
@@ -93,18 +91,14 @@ class TestOneToN(unittest.TestCase):
    @patch('src.investment.one_to_n_funcs.get_reimbursement_primary')
    @patch('src.investment.one_to_n_funcs.get_special_cases')
    @patch('src.investment.code_funcs.get_code_breakout')
    @patch('src.investment.one_to_n_funcs.validate_reimbursements_for_llm')
-    def test_reimbursement_level(self, mock_validate_llm, mock_code_breakout,
+    def test_reimbursement_level(self, mock_validate_llm, mock_special_cases, mock_reimbursement_primary):
                                mock_special_cases, mock_reimbursement_primary):
        # Setup
        mock_reimbursement_primary.return_value = [{"SERVICE_TERM": "Test Service", "REIMB_TERM": "term1 paid at rate"}]
        mock_validate_llm.return_value = True
        mock_special_cases.return_value = [{"SERVICE_TERM": "Test Service", "REIMB_TERM": "term1 paid at rate", "CARVEOUT_IND": "N"}]
        mock_code_breakout.return_value = [{"SERVICE_TERM": "Test Service", "REIMB_TERM": "term1 paid at rate", "CODE": "code1"}]
        mock_fields = MagicMock()
        mock_dataset = {}
        exhibit_page = "1"
        seen_pairs = set()
        exhibit_lesser_of = "N/A"  # Simulate no lesser of statement
@@ -112,15 +106,13 @@ class TestOneToN(unittest.TestCase):
        # Execute
        result = one_to_n_funcs.reimbursement_level(
-            self.exhibit_text, self.filename, mock_fields, mock_dataset, exhibit_page, seen_pairs, exhibit_lesser_of
+            self.exhibit_text, self.filename, mock_fields, exhibit_page, seen_pairs, exhibit_lesser_of
        )
        # Assert
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["CODE"], "code1")
        mock_reimbursement_primary.assert_called_once()
        mock_special_cases.assert_called_once()
        mock_code_breakout.assert_called_once()
    @patch('src.investment.one_to_n_funcs.get_reimbursement_primary')
    @patch('src.investment.one_to_n_funcs.validate_reimbursements_for_llm')
@@ -135,20 +127,17 @@ class TestOneToN(unittest.TestCase):
        mock_validate_llm.return_value = True
        mock_fields = MagicMock()
        mock_dataset = {}
        exhibit_page = "1"
        seen_pairs = set()
-        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases, \
+        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases:
            patch('src.investment.code_funcs.get_code_breakout') as mock_code_breakout:
            # Mock the downstream functions
            mock_special_cases.side_effect = lambda x, _: x  # Pass through
            mock_code_breakout.side_effect = lambda x, _, __: x  # Pass through
            # Execute
            result = one_to_n_funcs.reimbursement_level(
-                self.exhibit_text, self.filename, mock_fields, mock_dataset, 
+                self.exhibit_text, self.filename, mock_fields,
                exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
            )
@@ -174,18 +163,15 @@ class TestOneToN(unittest.TestCase):
        ]
        mock_fields = MagicMock()
        mock_dataset = {}
        exhibit_page = "2"
-        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases, \
+        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases:
            patch('src.investment.code_funcs.get_code_breakout') as mock_code_breakout:
            mock_special_cases.side_effect = lambda x, _: x
            mock_code_breakout.side_effect = lambda x, _, __: x
            # Execute
            result = one_to_n_funcs.reimbursement_level(
-                self.exhibit_text, self.filename, mock_fields, mock_dataset, 
+                self.exhibit_text, self.filename, mock_fields,
                exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
            )
@@ -213,15 +199,13 @@ class TestOneToN(unittest.TestCase):
        mock_dataset = {}
        exhibit_page = "2.1"  # Same base page (2) as seen_pairs
-        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases, \
+        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases:
            patch('src.investment.code_funcs.get_code_breakout') as mock_code_breakout:
            mock_special_cases.side_effect = lambda x, _: x
            mock_code_breakout.side_effect = lambda x, _, __: x
            # Execute
            result = one_to_n_funcs.reimbursement_level(
-                self.exhibit_text, self.filename, mock_fields, mock_dataset, 
+                self.exhibit_text, self.filename, mock_fields, 
                exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
            )
@@ -250,15 +234,13 @@ class TestOneToN(unittest.TestCase):
        mock_dataset = {}
        exhibit_page = "24.0"  # Different base page (24) from seen_pairs (23)
-        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases, \
+        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases:
            patch('src.investment.code_funcs.get_code_breakout') as mock_code_breakout:
            mock_special_cases.side_effect = lambda x, _: x
            mock_code_breakout.side_effect = lambda x, _, __: x
            # Execute
            result = one_to_n_funcs.reimbursement_level(
-                self.exhibit_text, self.filename, mock_fields, mock_dataset, 
+                self.exhibit_text, self.filename, mock_fields, 
                exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
            )
@@ -352,7 +334,7 @@ class TestOneToN(unittest.TestCase):
            # Execute
            result = one_to_n_funcs.reimbursement_level(
-                self.exhibit_text, self.filename, mock_fields, mock_dataset, 
+                self.exhibit_text, self.filename, mock_fields,
                exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
            )
@@ -0,0 +1,317 @@
 import unittest
 from unittest.mock import patch, MagicMock
 import json
 import os
 import pandas as pd
 import src.codes.code_funcs as code_funcs
 import src.codes.code_constants as code_constants
 from src.enums.delimiters import Delimiter
 class TestCodeFuncs(unittest.TestCase):
    def test_clean_service(self):
        """Test the clean_service function with various inputs"""
        # Test empty input
        self.assertEqual(code_funcs.clean_service(""), "")
        self.assertEqual(code_funcs.clean_service(None), "")
        # Test uppercase conversion
        self.assertEqual(code_funcs.clean_service("test service"), "TEST SERVICE")
        # Test synonym mapping
        with patch.dict(code_constants.SYNONYM_MAP, {"CT": "COMPUTED TOMOGRAPHY"}):
            self.assertEqual(code_funcs.clean_service("CT SCAN"), "COMPUTED TOMOGRAPHY SCAN")
        # Test removal list
        with patch.object(code_constants, "REMOVAL_LIST", ["SCAN"]):
            self.assertEqual(code_funcs.clean_service("CT SCAN"), "CT")
        # Test unlisted service
        self.assertEqual(code_funcs.clean_service("UNLISTED PROCEDURE"), "UNLISTED")
        # Test stop words
        with patch.object(code_constants, "STOP_WORD_LIST", ["THE", "AND"]):
            self.assertEqual(code_funcs.clean_service("THE AND"), "")
    @patch('src.utils.llm_utils.invoke_claude')
    @patch('src.utils.string_utils.universal_json_load')
    def test_code_explicit(self, mock_json_load, mock_invoke_claude):
        """Test the code_explicit function for invoking Claude"""
        # Setup mock responses
        mock_invoke_claude.return_value = '{"CPT4_PROC_CD": "12345"}'
        mock_json_load.return_value = {"CPT4_PROC_CD": "12345"}
        # Test successful extraction
        result = code_funcs.code_explicit("TEST SERVICE", "test.pdf")
        mock_invoke_claude.assert_called_once()
        self.assertEqual(result, {"CPT4_PROC_CD": "12345"})
    @patch('src.utils.llm_utils.invoke_claude')
    @patch('src.utils.string_utils.extract_text_from_delimiters')
    def test_code_implicit_special(self, mock_extract, mock_invoke_claude):
        """Test the code_implicit_special function for special case handling"""
        # Setup mock responses
        mock_invoke_claude.return_value = "mock_response"
        mock_extract.return_value = "Drugs"
        result = code_funcs.code_implicit_special("DRUG SERVICE", "test.pdf")
        self.assertEqual(result["PROCEDURE_CD"], "J0000-J9999")
        self.assertEqual(result["PROCEDURE_CD_DESC"], "Drugs")
        # Test no match
        mock_extract.return_value = ""
        result = code_funcs.code_implicit_special("OTHER SERVICE", "test.pdf")
        self.assertEqual(result, {})
    def test_get_embedding_levels(self):
        """Test the get_embedding_levels function for determining embedding levels"""
        # Test level 1 with both codes enabled
        implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True}
        embedding_levels, _, _, _ = code_funcs.get_embedding_levels(1, implicit_run_dict)
        self.assertIn("cpt_level1", embedding_levels)
        self.assertIn("hcpcs_level1", embedding_levels)
        self.assertIn("rev_level1", embedding_levels)
        # Test level 2 with only procedure code enabled
        implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": False}
        embedding_levels, _, _, _ = code_funcs.get_embedding_levels(2, implicit_run_dict)
        self.assertIn("cpt_level2", embedding_levels)
        self.assertIn("hcpcs_level2", embedding_levels)
        self.assertNotIn("rev", embedding_levels)
    @patch('src.utils.embedding_utils.load_faiss_index')
    @patch('src.codes.code_constants.model.encode')
    def test_get_match_list(self, mock_encode, mock_load_faiss):
        """Test the get_match_list function for retrieving matches from embeddings"""
        # Setup mock responses
        mock_encoded_vector = MagicMock()
        mock_encoded_vector.astype.return_value = "mock_vector"
        mock_encode.return_value = mock_encoded_vector
        # Create a mock FAISS index
        mock_index = MagicMock()
        # Configure the search method to return appropriate values
        mock_index.search.return_value = (
            # Scores array - 2D array with shape [1, top_k]
            [[0.9, 0.8]],
            # Indices array - 2D array with shape [1, top_k]
            [[0, 1]]
        )
        # Setup mock choices for the returned index
        choices = ["Description A", "Description B"]
        # Configure load_faiss_index to return our mock objects
        mock_load_faiss.return_value = (mock_index, None, choices)
        # Call the function
        match_list, highest_similarity = code_funcs.get_match_list("TEST SERVICE", ["test_level"], 2)
        # Verify results
        self.assertEqual(match_list, ["Description A", "Description B"])
        self.assertEqual(highest_similarity, 0.9)
        # Verify calls were made correctly
        mock_encode.assert_called_once_with(["TEST SERVICE"], normalize_embeddings=True)
        # Make sure load_faiss_index was called for the test_level
        mock_load_faiss.assert_called_once()
        # Check that search was called with the encoded vector
        mock_index.search.assert_called_once_with("mock_vector", 2)
    @patch('src.codes.code_funcs.get_embedding_levels')
    @patch('src.codes.code_funcs.get_match_list')
    @patch('src.utils.llm_utils.invoke_claude')
    @patch('src.utils.string_utils.universal_json_load')
    def test_code_implicit_rag(self, mock_json_load, mock_invoke_claude, mock_get_match_list, mock_get_embedding_levels):
        """Test the code_implicit_rag function for RAG-based code extraction"""
        # Setup mock responses
        mock_get_embedding_levels.return_value = (["test_level"], 
                                                 {"12345": "Test Procedure"}, 
                                                 {}, 
                                                 {"678": "Test Revenue"})
        mock_get_match_list.return_value = (["Test Procedure"], 0.9)
        mock_invoke_claude.return_value = '["Test Procedure"]'
        mock_json_load.return_value = ["Test Procedure"]
        result = code_funcs.code_implicit_rag("TEST SERVICE", {"PROCEDURE_CD": True, "REVENUE_CD": True}, "test.pdf")
        self.assertEqual(result["PROCEDURE_CD"], "12345")
        self.assertEqual(result["PROCEDURE_CD_DESC"], "Test Procedure")
        self.assertIn("CODE_METHODOLOGY", result)
        # Test exception handling
        mock_json_load.side_effect = Exception("Test error")
        result = code_funcs.code_implicit_rag("TEST SERVICE", {"PROCEDURE_CD": True, "REVENUE_CD": True}, "test.pdf")
        self.assertIn("CODE_METHODOLOGY", result)
    @patch('src.utils.llm_utils.invoke_claude')
    @patch('src.utils.string_utils.extract_text_from_delimiters')
    def test_code_last_check(self, mock_extract, mock_invoke_claude):
        """Test the code_last_check function for categorizing non-matched services"""
        # Setup mock responses
        mock_invoke_claude.return_value = "mock_response"
        mock_extract.return_value = "Specific"
        result = code_funcs.code_last_check("SPECIAL SERVICE", "test.pdf")
        self.assertEqual(result, "Specific")
        # Test generic case
        mock_extract.return_value = "Generic"
        result = code_funcs.code_last_check("GENERIC SERVICE", "test.pdf")
        self.assertEqual(result, "Generic")
    @patch('src.utils.llm_utils.invoke_claude')
    @patch('src.utils.string_utils.universal_json_load')
    def test_fill_bill_type(self, mock_json_load, mock_invoke_claude):
        """Test the fill_bill_type function for populating bill type information"""
        # Setup mock responses
        mock_invoke_claude.return_value = '["Inpatient Hospital"]'
        mock_json_load.return_value = ["Inpatient Hospital"]
        # Mock bill type mappings
        with patch.dict(code_constants.bill_type_reverse_mapping, {"Inpatient Hospital": "11X"}):
            result = code_funcs.fill_bill_type("INPATIENT SERVICE", {})
            self.assertEqual(result["BILL_TYPE_CD"], "11X")
            self.assertEqual(result["BILL_TYPE_CD_DESC"], "Inpatient Hospital")
    def test_get_implicit_runs(self):
        """Test the get_implicit_runs function for determining which code runs to execute"""
        # Test hospital claim type
        result = code_funcs.get_implicit_runs({"AARETE_DERIVED_CLAIM_TYPE_CD": "H"})
        self.assertTrue(result["REVENUE_CD"])
        # Test medical claim type with bill type
        result = code_funcs.get_implicit_runs({
            "AARETE_DERIVED_CLAIM_TYPE_CD": "M",
            "BILL_TYPE_CD_DESC": "Outpatient"
        })
        self.assertTrue(result["REVENUE_CD"])
        self.assertTrue(result["PROCEDURE_CD"])
        # Test inpatient bill type
        result = code_funcs.get_implicit_runs({
            "AARETE_DERIVED_CLAIM_TYPE_CD": "M",
            "BILL_TYPE_CD_DESC": "Inpatient Hospital"
        })
        self.assertTrue(result["REVENUE_CD"])
        self.assertFalse(result["PROCEDURE_CD"])
    @patch('src.codes.code_funcs.clean_service')
    @patch('src.codes.code_funcs.fill_bill_type')
    @patch('src.codes.code_funcs.get_implicit_runs')
    @patch('src.codes.code_funcs.code_explicit')
    @patch('src.codes.code_funcs.code_category')
    @patch('src.codes.code_funcs.code_implicit_special')
    @patch('src.codes.code_funcs.code_implicit_rag')
    @patch('src.codes.code_funcs.code_last_check')
    def test_extract_codes_from_service(self, mock_last_check, mock_implicit_rag, mock_implicit_special,
                                        mock_category, mock_explicit, mock_implicit_runs, mock_fill_bill_type,
                                        mock_clean_service):
        """Test the extract_codes_from_service function end-to-end"""
        # Setup mock responses
        mock_clean_service.return_value = "CLEAN SERVICE"
        mock_fill_bill_type.return_value = {"SERVICE_TERM": "TEST SERVICE", "BILL_TYPE_CD_DESC": "Outpatient"}
        mock_implicit_runs.return_value = {"PROCEDURE_CD": True, "REVENUE_CD": True}
        # Test explicit code match
        mock_explicit.return_value = {"PROCEDURE_CD": "12345"}
        result = code_funcs.extract_codes_from_service({"SERVICE_TERM": "TEST SERVICE"})
        self.assertEqual(result["PROCEDURE_CD"], "12345")
        self.assertEqual(result["CODE_METHODOLOGY"], "Explicit")
        # Test unlisted service
        mock_clean_service.return_value = "UNLISTED"
        result = code_funcs.extract_codes_from_service({"SERVICE_TERM": "UNLISTED SERVICE"})
        self.assertEqual(result["PROCEDURE_CD_DESC"], "UNLISTED")
        # Test category match
        mock_clean_service.return_value = "CLEAN SERVICE"
        mock_explicit.return_value = {"PROCEDURE_CD": ["Category: J"]}
        mock_category.return_value = {"PROCEDURE_CD": ["J0001"], "PROCEDURE_CD_DESC": ["Drug A"]}
        result = code_funcs.extract_codes_from_service({"SERVICE_TERM": "J CODES"})
        self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Letter Category")
        # Test special case match
        mock_explicit.return_value = {}
        mock_implicit_special.return_value = {"PROCEDURE_CD": "J0000-J9999", "PROCEDURE_CD_DESC": "Drugs"}
        result = code_funcs.extract_codes_from_service({"SERVICE_TERM": "DRUG SERVICE"})
        self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Special Case")
        # Test RAG match
        mock_implicit_special.return_value = {}
        mock_implicit_rag.return_value = {"PROCEDURE_CD": "12345", "CODE_METHODOLOGY": "Implicit - Level 1"}
        result = code_funcs.extract_codes_from_service({"SERVICE_TERM": "TEST SERVICE"})
        self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Level 1")
        # Test no match
        mock_implicit_rag.return_value = {}
        mock_last_check.return_value = "Specific"
        result = code_funcs.extract_codes_from_service({"SERVICE_TERM": "TEST SERVICE"})
        self.assertEqual(result["CODE_METHODOLOGY"], "Specific - No Match")
    def test_fill_claim_type(self):
        """Test the fill_claim_type function for determining claim types"""
        # Test with existing claim types
        answer_dicts = [
            {"AARETE_DERIVED_CLAIM_TYPE_CD": "H"},
            {"AARETE_DERIVED_CLAIM_TYPE_CD": "H"},
            {"AARETE_DERIVED_CLAIM_TYPE_CD": "M"},
            {}
        ]
        result = code_funcs.fill_claim_type(answer_dicts)
        self.assertEqual(result[3]["AARETE_DERIVED_CLAIM_TYPE_CD"], "H")
        # Test with no claim types
        answer_dicts = [{}, {}]
        result = code_funcs.fill_claim_type(answer_dicts)
        self.assertEqual(result, answer_dicts)
    @patch('src.codes.code_funcs.fill_claim_type')
    @patch('src.codes.code_funcs.extract_codes_from_service')
    def test_code_breakout(self, mock_extract_codes, mock_fill_claim_type):
        """Test the code_breakout function for processing a DataFrame of service records"""
        # Create a test DataFrame
        input_df = pd.DataFrame([
            {"SERVICE_TERM": "SERVICE A"},
            {"SERVICE_TERM": "SERVICE B"}
        ])
        # Setup mock responses for the records after conversion to list of dicts
        mock_fill_claim_type.return_value = [
            {"SERVICE_TERM": "SERVICE A"},
            {"SERVICE_TERM": "SERVICE B"}
        ]
        # Setup extract_codes_from_service mock to return augmented records
        mock_extract_codes.side_effect = [
            {"SERVICE_TERM": "SERVICE A", "PROCEDURE_CD": "12345"},
            {"SERVICE_TERM": "SERVICE B", "PROCEDURE_CD": "67890"}
        ]
        # Call the function
        result_df = code_funcs.code_breakout(input_df)
        # Verify the result is a DataFrame
        self.assertIsInstance(result_df, pd.DataFrame)
        # Verify the correct procedure codes were added
        self.assertEqual(result_df.iloc[0]["PROCEDURE_CD"], "12345")
        self.assertEqual(result_df.iloc[1]["PROCEDURE_CD"], "67890")
        # Verify the DataFrame has the expected number of rows
        self.assertEqual(len(result_df), 2)
        # Verify fill_claim_type was called with the correct data
        mock_fill_claim_type.assert_called_once()
        args, _ = mock_fill_claim_type.call_args
        self.assertIsInstance(args[0], list)
        # Verify extract_codes_from_service was called for each record
        self.assertEqual(mock_extract_codes.call_count, 2)
 if __name__ == '__main__':
    unittest.main()