From 863c9db9a1649cd8bb6fc42eb19c0edcf9216030 Mon Sep 17 00:00:00 2001 From: Praneel Panchigar Date: Wed, 11 Mar 2026 17:39:31 +0000 Subject: [PATCH] Merged in feature/code_optimization (pull request #903) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Feature/code optimization * Ran Black * made a small change in code_last_check, fixed so it returns string and not single char * Made changes to make sure that default_ind postprocess only happens to the cc output and not dashboard * Merge bugfix/default_ind_postprocess into feature/code_optimization - Parent-child: child_rank column init and cols_to_keep filter - Postprocess: default_ind only on cc output - code_last_check: parser str/list return handling + exception logging - Prompt: FIELD ASSIGNMENT for code extraction (no CRITICAL prefix) * More prompt changes and debugging print statements * Added system to group same service term + bill type cd + claim type cd and then parallelize the code extraction (will improve consistency) * Revert one-time PROV_INFO_JSON ad hoc logic; retain json_utils and output format - Remove temporary postprocess step that filled empty TIN in PROV_INFO_JSON from FILENAME_TIN (fill_prov_info_tin_from_filename_tin call in postprocess.py). This behavior is intended to move upstream per JIRA (tin_npi_funcs / extraction). - Revert postprocess_existing_output.py to generic config: empty INPUT_DIR, INPUT_FILENAME, OUTPUT_CSV. Remove one-time hyphen-strip for TIN/NPI in PROV_INFO_JSON and remove _strip_hyphens_from_prov_info_json_cell. - Keep json_utils PROV_INFO_JSON helpers (parse_prov_info_json_cell, serialize_prov_info_json, format_prov_info_json) and (str,str) serialization in contract_config_postprocess. Keep fill_prov_info_tin_from_filename_tin and related helpers in postprocessing_funcs for potential upstream reuse. - Tests in test_json_parsers.py updated for (str,str) PROV_INFO_JSON output. * Add arbitration-based implicit code flow and prefer Level 2 RAG (Fix 5d, 4c+7) Implicit flow no longer uses first-match-wins. We build a case from all applicable strategies and run a single arbitration step to choose the best candidate. - build_implicit_candidates: Run Category (when explicit had Category), Special, and RAG; return list of {source, code_answer_dict}. RAG result uses preferred level when both return (Fix 5d). - code_implicit_rag (Fix 5d): Run both Level 1 and Level 2, collect results; when both return non-empty, prefer Level 2 over Level 1. - code_implicit_arbitration: 0 candidates -> None; 1 candidate -> return it with CODE_METHODOLOGY = 'Implicit - Arbitration (Source)' (no LLM); 2+ candidates -> one CODE_IMPLICIT_ARBITRATION LLM call (chosen_index or no_match). If no candidates or no_match, code_last_check runs as before. - prompt_templates: Add CODE_IMPLICIT_ARBITRATION and CODE_IMPLICIT_ARBITRATION_INSTRUCTION (cached). Use _create_json_dict_parser(field_names=['chose… * fix: test_code_breakout avoid LLM call and flaky order on CI - Add BILL_TYPE_CD_DESC to test data so pre-pass skips fill_bill_type (no Bedrock on CI) - Use side_effect function keyed by SERVICE_TERM for deterministic result under ThreadPoolExecutor * ran black * feat: hybrid code validation with unmatched tagging and code quality fixes Validation and unmatched tagging: - Add format-only filtering for explicit codes (drop wrong-format, preserve valid) - Add unified _has_unmatched_codes() used by both explicit and implicit paths - Tag CODE_METHODOLOGY with "- Unmatched" when codes are format-valid but not explicit keys in our mappings (codes and range keys count as mapped) - Remove retry logic and "Explicit - Validation Failed" in favor of single LLM call with Unmatched tagging - Include all revenue level mappings in valid revenue set Code quality and bug fixes: - Fix fill_bill_type += string bug (was splitting "11X" into chars) - Fix bare except clauses in code_explicit, code_category, code_implicit_special, fill_bill_type (now except Exception as e) - Replace eval() with json.loads() in fill_grouper_cd_desc - Implement INVALID_SERVICE short-circuit in RAG (Fix 5a) - Cache valid procedure/revenue code sets per Constants instance - Store mappings in level_dicts… * chore: remove debug print statements from code extraction pipeline Remove all [DEBUG_CODE_EXTRACTION] print statements added during development of the hybrid validation and implicit arbitration flow. * refactor: robust PROV_INFO_JSON sanitization with multi-format parsing Replace the simple format_prov_info_json with a layered sanitize_prov_info_json that handles all known malformed variants: single-quoted dicts, empty-value-after-colon patterns, list-typed field values, and hyphenated TINs. Extract _normalize_prov_entries and _prov_value_to_str helpers for uniform str-valued output. format_prov_info_json now delegates to sanitize_prov_info_json. * merge: sync feature/code_optimization with DEV Resolve conflict in prompt_templates.py: keep detailed CRITICAL field assignment instructions from feature branch; accept new AARETE_DERIVED_PROVIDER_NAME prompt function from DEV. * Minor change to explict prompt * fix: resolve mypy errors and remove stale pipe-delim splits in RAG - Narrow grouper_cd type to str|None with explicit None check to satisfy mypy in fill_grouper_cd_desc. - Remove dead pipe-delim split logic from CPT and HCPCS matching blocks (no pipes in those mappings). - Retain pipe split for revenue block only (rev_level1.csv keys are pipe-delimited, e.g. "0810|0811|0812"). * Code extraction: format-only explicit, keep unmapped tag, RAG Level 1/2 as separate candidates - Explicit: accept all format-valid codes (single + range); do not drop for mapping. Validity check only appends "- Unmatched" to methodology when code not in mappings. - Add _procedure_code_or_range_format_valid; use in validate_explicit_codes and _has_unmatched_codes. - RAG: return list of Level 1 and Level 2 candidates (distinct source keys); remove Level 2 preference. - build_implicit_candidates: consume RAG list and append each as separate candidate. - CODE_EXPLICIT: clarify range extraction when service term describes a range. - CODE_IMPLICIT_ARBITRATION: choose most appropriate candidate(s), avoid overly broad. - Tests updated for new RAG return shape and validate_explicit_codes behavior. * Resolved Merge Conflicts * Fixed failing tests * Remove unused functions * Remove unused functions * Revert "Remove unused functions" This reverts commit 39954db6a421e55501282d4e6a270c88624049fe. * remove unused * Merged DEV into feature/code_optimization * Merged DEV into feature/code_optimization Approved-by: Katon Minhas --- src/codes/code_funcs.py | 1013 ++++++++++++++++++++++-------- src/constants/constants.py | 21 +- src/prompts/prompt_templates.py | 83 ++- src/tests/test_code_funcs.py | 629 ++++++++++++++++++- src/tests/test_prompt_caching.py | 33 + 5 files changed, 1500 insertions(+), 279 deletions(-) diff --git a/src/codes/code_funcs.py b/src/codes/code_funcs.py index e8918b5..87de6b6 100644 --- a/src/codes/code_funcs.py +++ b/src/codes/code_funcs.py @@ -1,18 +1,16 @@ -import ast +import concurrent.futures import json import logging -import concurrent.futures -import os import re +from collections import defaultdict import pandas as pd + import src.config as config import src.prompts.prompt_templates as prompt_templates import src.utils.llm_utils as llm_utils import src.utils.string_utils as string_utils from src.constants.constants import Constants -from src.constants.delimiters import Delimiter -from src.prompts.fieldset import FieldSet def clean_service(service, constants: Constants) -> str: @@ -57,55 +55,19 @@ def clean_service(service, constants: Constants) -> str: return service -def get_regex_answers(service, code_answer_dict): +def code_explicit(service: str, methodology: str, filename: str) -> 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. + Invoke the LLM to extract explicit codes from a service term. Args: - service (str): The service name or identifier. - code_answer_dict (dict): Dictionary to store the extracted codes. + service: The service name or identifier. + methodology: The reimbursement methodology for additional context. + filename: The name of the file, used for logging or tracking. + Returns: - dict: A dictionary containing the extracted codes. If no codes are found, returns an empty dictionary. + A dictionary of code fields parsed from the LLM response, or empty + dict on failure. """ - - # 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"] = rev_matches - - return code_answer_dict - - -def code_explicit(service, methodology, 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 - field definitions are now included in CODE_EXPLICIT_INSTRUCTION() for caching prompt, _parser = prompt_templates.CODE_EXPLICIT(service, methodology) llm_answer_raw = llm_utils.invoke_claude( @@ -119,25 +81,35 @@ def code_explicit(service, methodology, filename): try: code_answer_dict = _parser(llm_answer_raw) return code_answer_dict - except: + except Exception as e: logging.error( - f"Failed to parse LLM response in code_explicit: {llm_answer_raw}" + "Failed to parse LLM response in code_explicit: %s | error: %s", + llm_answer_raw, + e, ) return {} -def code_category(service, proc_category, hcpcs_level2_mapping, filename): - """ - Determines if the service explicitly names a code category, such as "A-Codes" or "J-Codes", and returns the corresponding codes. - Args: - service (str): The service name or identifier. - proc_category (list): List of procedure categories to check against. - filename (str): The name of the file associated with the operation, used for logging or tracking purposes. - Returns: - dict: A dictionary containing the procedure codes and their descriptions. If no codes are found, returns an empty dictionary. - If the service does not explicitly name a code category, it returns an empty dictionary. +def code_category( + service: str, + proc_category: list, + hcpcs_level2_mapping: dict, + filename: str, +) -> dict: """ + Determines if the service explicitly names a code category, such as + "A-Codes" or "J-Codes", and returns the corresponding codes. + Args: + service: The service name or identifier. + proc_category: List of procedure categories to check against. + hcpcs_level2_mapping: HCPCS level 2 code-to-description mapping. + filename: The file name for logging/tracking. + + Returns: + A dictionary with PROCEDURE_CD and PROCEDURE_CD_DESC lists, or + empty dict on failure. + """ code_category_list = [v.split(":")[1].strip() for v in proc_category if ":" in v] matching_values = [] @@ -158,40 +130,44 @@ def code_category(service, proc_category, hcpcs_level2_mapping, filename): instruction=prompt_templates.CODE_CATEGORY_INSTRUCTION(), ) try: - llm_answer_final = _parser(llm_answer_raw) # Returns list - except: + llm_answer_final = _parser(llm_answer_raw) + except Exception as e: logging.error( - f"Failed to parse LLM response in code_category: {llm_answer_raw}" + "Failed to parse LLM response in code_category: %s | error: %s", + llm_answer_raw, + e, ) return {} - code_answer_dict = {"PROCEDURE_CD": [], "PROCEDURE_CD_DESC": []} + # Pre-build reverse dict for O(1) description-to-code lookup. + desc_to_code: dict[str, str] = {v: k for k, v in hcpcs_level2_mapping.items()} + + code_answer_dict: dict = {"PROCEDURE_CD": [], "PROCEDURE_CD_DESC": []} for answer in llm_answer_final: if "0000" in answer: code_answer_dict["PROCEDURE_CD"].append(answer) code_answer_dict["PROCEDURE_CD_DESC"].append("") - elif answer in hcpcs_level2_mapping.values(): - code = list(hcpcs_level2_mapping.keys())[ - list(hcpcs_level2_mapping.values()).index(answer) - ] + elif answer in desc_to_code: + code = desc_to_code[answer] code_answer_dict["PROCEDURE_CD"].append(str(code)) code_answer_dict["PROCEDURE_CD_DESC"].append(str(answer)) return code_answer_dict -def code_implicit_special(service, filename): +def code_implicit_special(service: str, filename: str) -> dict: """ - 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 + Allow certain Service classes to be handled as special cases, such as + Drugs, Vaccines, Surgery, and PT/OT/ST. 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. - """ + service: The service name or identifier. + filename: The file name for logging/tracking. + Returns: + A dictionary with PROCEDURE_CD and PROCEDURE_CD_DESC, or empty + dict when no special case matched. + """ prompt, _parser = prompt_templates.CODE_IMPLICIT_SPECIAL(service) try: llm_answer_raw = llm_utils.invoke_claude( @@ -202,10 +178,8 @@ def code_implicit_special(service, filename): instruction=prompt_templates.CODE_IMPLICIT_SPECIAL_INSTRUCTION(), ) llm_answer_final = _parser(llm_answer_raw) - except: - logging.error( - f"Failed to parse LLM response in code_implicit_special: {llm_answer_raw}" - ) + except Exception as e: + logging.error("Failed to parse LLM response in code_implicit_special: %s", e) return {} special_case_mapping = { @@ -259,18 +233,23 @@ def code_implicit_special(service, filename): return code_answer_dict -def get_embedding_levels(level_suffix, implicit_run_dict, constants: Constants): +def get_embedding_levels( + level_suffix: int, implicit_run_dict: dict, constants: Constants +) -> tuple[list, dict, dict, dict]: """ - Retrieves the embedding levels and their corresponding mappings for a given level suffix. + 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. + level_suffix: 1 or 2; determines which mappings to use. + implicit_run_dict: Dict indicating which codes to run implicitly. + constants: Application constants with mapping data. + 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. + (embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping). + + Raises: + ValueError: If level_suffix is not 1 or 2. """ if level_suffix == 1: cpt_mapping = constants.CPT_LEVEL1_MAPPING @@ -299,20 +278,29 @@ def get_embedding_levels(level_suffix, implicit_run_dict, constants: Constants): if implicit_run_dict["REVENUE_CD"]: embedding_levels += [f"rev"] + else: + raise ValueError(f"Unsupported level_suffix: {level_suffix}") + return embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping -def get_match_list(service, embedding_levels, constants: Constants, top_k: int): +def get_match_list( + service: str, + embedding_levels: list, + constants: Constants, + top_k: int, +) -> tuple[list, float]: """ - Retrieves the best matches for a given service from multiple embedding levels using FAISS index search. + Retrieves the best matches for a service from embedding levels via FAISS. + 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. + service: The service name or identifier. + embedding_levels: List of embedding levels to search in. + constants: Application constants with embedding data. + top_k: 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. + (match_list, highest_similarity). """ target_vec = constants.EMBEDDING_MODEL.encode( [service], normalize_embeddings=True @@ -332,28 +320,38 @@ def get_match_list(service, embedding_levels, constants: Constants, top_k: int): return match_list, highest_similarity -def code_implicit_rag(service, implicit_run_dict, filename, constants): - """ - Processes a service string to find implicit codes using RAG (Retrieval-Augmented Generation) methodology. - Args: - service (str): The service name or identifier. - implicit_run_dict (dict): Dictionary indicating which codes to run implicitly. - filename (str): The name of the file associated with the operation, used for logging or tracking purposes. - Returns: - dict: A dictionary containing the implicit codes found for the service. If no codes are found, returns an empty dictionary. +def code_implicit_rag( + service: str, + implicit_run_dict: dict, + filename: str, + constants: Constants, +) -> list[dict]: """ + Find implicit codes using RAG (Retrieval-Augmented Generation). - # Get Embedding Embeddings and Mappings for Levels 1 and 2 - level_dicts = [] + Runs Level 1 and Level 2; returns each non-empty result as a separate + candidate so arbitration can choose or combine. Source keys "Level 1" and + "Level 2" identify each candidate (same pattern as Category, Special). + + Args: + service: The service name or identifier. + implicit_run_dict: Dict indicating which codes to run implicitly. + filename: The file name for logging/tracking. + constants: Application constants with mapping and embedding data. + + Returns: + List of {"source": "Level 1"|"Level 2", "code_answer_dict": dict}. + Only includes entries that have PROCEDURE_CD or REVENUE_CD (no placeholder-only). + """ + # Get embeddings, mappings, and best matches for Levels 1 and 2. + level_dicts: list[dict] = [] for level_suffix in [1, 2]: - # Get Mappings and Embeddings embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping = ( get_embedding_levels(level_suffix, implicit_run_dict, constants) ) if not embedding_levels: - return {} + return [] - # Get Best Matches match_list, highest_similarity = get_match_list( service, embedding_levels, constants, top_k=5 ) @@ -362,25 +360,24 @@ def code_implicit_rag(service, implicit_run_dict, filename, constants): "level_suffix": level_suffix, "match_list": match_list, "highest_similarity": highest_similarity, + "cpt_mapping": cpt_mapping, + "hcpcs_mapping": hcpcs_mapping, + "rev_mapping": rev_mapping, } ) - # 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 + candidates: list[dict] = [] for level_dict in level_dicts: - embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping = ( - get_embedding_levels( - level_dict["level_suffix"], implicit_run_dict, constants - ) - ) + cpt_mapping = level_dict["cpt_mapping"] + hcpcs_mapping = level_dict["hcpcs_mapping"] + rev_mapping = level_dict["rev_mapping"] prompt, _parser = prompt_templates.CODE_IMPLICIT( service, level_dict["match_list"] ) - # Run Prompt llm_answer_raw = llm_utils.invoke_claude( prompt, "sonnet_latest", @@ -393,61 +390,48 @@ def code_implicit_rag(service, implicit_run_dict, filename, constants): llm_answer_final = _parser(llm_answer_raw) except Exception as e: logging.error(str(e)) - return {} + continue if not llm_answer_final: continue - # Populate answers, if any - code_answer_dict = {} - proc_codes, rev_codes = [], [] - proc_descs, rev_descs = [], [] - for description in llm_answer_final: - if description == "INVALID_SERVICE": - code_answer_dict["CODE_METHODOLOGY"] = "Generic - Prompt" - continue + # INVALID_SERVICE is a hard stop for this level; do not add a candidate. + if "INVALID_SERVICE" in llm_answer_final: + continue + code_answer_dict: dict = {} + proc_codes: list[str] = [] + rev_codes: list[str] = [] + proc_descs: list[str] = [] + rev_descs: list[str] = [] + for description in llm_answer_final: if description in cpt_mapping.values(): matching_codes = [ str(key) for key, val in cpt_mapping.items() if val == description ] if matching_codes: - # Split pipe-delimited codes into individual codes - for code in matching_codes: - # Handle pipe-delimited codes (e.g., "99170-99199|99200-99210") - if "|" in code: - proc_codes.extend(code.split("|")) - else: - proc_codes.append(code) + proc_codes.extend(matching_codes) proc_descs.append(description) if description in hcpcs_mapping.values(): matching_codes = [ str(key) for key, val in hcpcs_mapping.items() if val == description ] if matching_codes: - # Split pipe-delimited codes into individual codes - for code in matching_codes: - # Handle pipe-delimited codes (e.g., "99170-99199|99200-99210") - if "|" in code: - proc_codes.extend(code.split("|")) - else: - proc_codes.append(code) + proc_codes.extend(matching_codes) proc_descs.append(description) if description in rev_mapping.values(): matching_codes = [ str(key) for key, val in rev_mapping.items() if val == description ] if matching_codes: - # Split pipe-delimited codes into individual codes + # REV_LEVEL1_MAPPING keys are pipe-delimited (e.g. "0810|0811|0812"). for code in matching_codes: - # Handle pipe-delimited codes (e.g., "99170-99199|99200-99210") if "|" in code: rev_codes.extend(code.split("|")) else: rev_codes.append(code) rev_descs.append(description) - # Combine answers if proc_codes: code_answer_dict["PROCEDURE_CD"] = proc_codes code_answer_dict["PROCEDURE_CD_DESC"] = proc_descs @@ -462,21 +446,153 @@ def code_implicit_rag(service, implicit_run_dict, filename, constants): ) if code_answer_dict: - return code_answer_dict + source_label = f"Level {level_dict['level_suffix']}" + candidates.append( + {"source": source_label, "code_answer_dict": code_answer_dict} + ) - return {} + return candidates -def code_last_check(service, filename): +def _format_implicit_candidate_for_prompt(candidate: dict, index: int) -> str: + """Format a single implicit candidate for the arbitration prompt.""" + source = candidate.get("source", "Unknown") + code_dict = candidate.get("code_answer_dict", {}) + parts = [f"Candidate {index} (Source: {source})"] + for key in ["PROCEDURE_CD", "PROCEDURE_CD_DESC", "REVENUE_CD", "REVENUE_CD_DESC"]: + if key in code_dict and code_dict[key]: + val = code_dict[key] + if isinstance(val, list): + val = json.dumps(val) + parts.append(f" {key}: {val}") + return "\n".join(parts) + + +def build_implicit_candidates( + service_clean: str, + code_answer_dict_after_explicit: dict, + implicit_run_dict: dict, + filename: str, + constants: Constants, +) -> list: """ - For Services that did not match any codes, this function categorizes the reason as "Specific" or "Generic". + Build the case for implicit arbitration: run Category (if applicable), Special, + and RAG; return a list of non-empty candidates with source labels. + + Returns: + List of dicts: [{"source": str, "code_answer_dict": dict}, ...] + """ + candidates = [] + proc_cd = code_answer_dict_after_explicit.get("PROCEDURE_CD", []) + + if any("Category:" in str(v) for v in proc_cd): + cat_result = code_category( + service_clean, + proc_cd, + constants.HCPCS_LEVEL2_MAPPING, + filename, + ) + if cat_result and any( + not string_utils.is_empty(v) for v in cat_result.values() + ): + candidates.append({"source": "Category", "code_answer_dict": cat_result}) + + special_result = code_implicit_special(service_clean, filename) + if special_result and any( + not string_utils.is_empty(v) for v in special_result.values() + ): + candidates.append({"source": "Special", "code_answer_dict": special_result}) + + rag_candidates = code_implicit_rag( + service_clean, implicit_run_dict, filename, constants + ) + for item in rag_candidates: + if not isinstance(item, dict): + continue + source = item.get("source", "RAG") + code_answer_dict = item.get("code_answer_dict", {}) + if code_answer_dict and any( + not string_utils.is_empty(v) for v in code_answer_dict.values() + ): + candidates.append({"source": source, "code_answer_dict": code_answer_dict}) + + return candidates + + +def code_implicit_arbitration( + service: str, candidates: list, filename: str +) -> dict | None: + """ + Run one LLM arbitration call to choose the best candidate from the case. Args: - service (str): The service term. - filename (str): The name of the file associated with the operation. + service: The service term. + candidates: List of {"source": str, "code_answer_dict": dict}. + filename: For logging/cache. + 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. + The chosen code_answer_dict with CODE_METHODOLOGY set to reflect arbitration + and the winning source, or None if arbitrator returned no_match. + """ + if not candidates: + return None + if len(candidates) == 1: + chosen = candidates[0]["code_answer_dict"].copy() + source = candidates[0]["source"] + chosen["CODE_METHODOLOGY"] = f"Implicit - Arbitration ({source})" + return chosen + + candidates_text = "\n\n".join( + _format_implicit_candidate_for_prompt(c, i) for i, c in enumerate(candidates) + ) + prompt, _parser = prompt_templates.CODE_IMPLICIT_ARBITRATION( + service, candidates_text + ) + try: + llm_answer_raw = llm_utils.invoke_claude( + prompt, + "sonnet_latest", + filename, + cache=True, + instruction=prompt_templates.CODE_IMPLICIT_ARBITRATION_INSTRUCTION(), + ) + parsed = _parser(llm_answer_raw) + except Exception as e: + logging.error("code_implicit_arbitration LLM failed: %s", e) + return None + + if not isinstance(parsed, dict): + return None + no_match_val = parsed.get("no_match") + if no_match_val in (True, "True", "true", "1"): + return None + idx_val = parsed.get("chosen_index") + if idx_val is None: + return None + try: + result = int(idx_val) if not isinstance(idx_val, int) else idx_val + except (TypeError, ValueError): + logging.warning("code_implicit_arbitration invalid chosen_index %s", idx_val) + return None + if result < 0 or result >= len(candidates): + logging.warning("code_implicit_arbitration invalid index %s, using 0", result) + result = 0 + chosen = candidates[result]["code_answer_dict"].copy() + source = candidates[result]["source"] + chosen["CODE_METHODOLOGY"] = f"Implicit - Arbitration ({source})" + return chosen + + +def code_last_check(service: str, filename: str) -> str: + """ + Categorize why a service did not match any codes. + + Args: + service: The service term. + filename: The file name for logging/tracking. + + Returns: + "Specific" or "Generic". """ prompt, _parser = prompt_templates.CODE_LAST_CHECK(service) @@ -498,35 +614,33 @@ def code_last_check(service, filename): def fill_bill_type( - service, - answer_dict, - BILL_TYPE_MAPPING, - BILL_TYPE_REVERSE_MAPPING, - reimb_term=None, - exhibit_text=None, -): + service: str, + answer_dict: dict, + BILL_TYPE_MAPPING: dict, + BILL_TYPE_REVERSE_MAPPING: dict, + reimb_term: str | None = None, + exhibit_text: str | None = None, +) -> 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. - BILL_TYPE_MAPPING (dict): Mapping from codes to descriptions. - BILL_TYPE_REVERSE_MAPPING (dict): Mapping from descriptions to codes. - reimb_term (str, optional): The reimbursement term for additional context. - exhibit_text (str, optional): Full exhibit text for contextual analysis. + service: The service term. + answer_dict: The answer dictionary to update. + BILL_TYPE_MAPPING: Mapping from codes to descriptions. + BILL_TYPE_REVERSE_MAPPING: Mapping from descriptions to codes. + reimb_term: The reimbursement term for additional context. + exhibit_text: Full exhibit text for contextual analysis. Returns: - dict: The updated answer dictionary. + The updated answer dictionary. """ - valid_bill_type = sorted(list(set(BILL_TYPE_MAPPING.values()))) prompt, _parser = prompt_templates.FILL_BILL_TYPE( service, valid_bill_type, reimb_term=reimb_term, exhibit_text=None ) - # First attempt: Service term + reimbursement term llm_answer_raw = llm_utils.invoke_claude( prompt, "sonnet_latest", @@ -536,10 +650,11 @@ def fill_bill_type( ) try: llm_answer_final = _parser(llm_answer_raw) - except: + except Exception as e: + logging.error("Failed to parse fill_bill_type response: %s", e) return answer_dict - # Second attempt: Service term + reimbursement term + exhibit context (if first attempt failed and context available) + # Second attempt with exhibit context if first attempt returned empty. if string_utils.is_empty(llm_answer_final) and exhibit_text: prompt, _parser = prompt_templates.FILL_BILL_TYPE( service, @@ -557,22 +672,22 @@ def fill_bill_type( ) try: llm_answer_final = _parser(llm_answer_raw) - except: + except Exception as e: + logging.error("Failed to parse fill_bill_type response: %s", e) return answer_dict - # Keep legacy behavior: return answer_dict unchanged if no result found if not llm_answer_final: return answer_dict - bill_codes, bill_descs = [], [] + bill_codes: list[str] = [] + bill_descs: list[str] = [] for description in llm_answer_final: if description in BILL_TYPE_REVERSE_MAPPING: - codes = BILL_TYPE_REVERSE_MAPPING[description] - # Reverse mapping returns a list of codes for each description - if isinstance(codes, list): - bill_codes.extend(codes) + value = BILL_TYPE_REVERSE_MAPPING[description] + if isinstance(value, list): + bill_codes.extend(value) else: - bill_codes.append(codes) + bill_codes.append(str(value)) bill_descs.append(description) if bill_codes: @@ -582,26 +697,27 @@ def fill_bill_type( return answer_dict -def fill_grouper_cd_desc(answer_dict, constants: Constants): +def fill_grouper_cd_desc(answer_dict: dict, constants: Constants) -> dict: """ - Fills the GROUPER_CD_DESC field in the answer dictionary based on the GROUPER_CD. + Fills the GROUPER_CD_DESC field in the answer dictionary based on + the GROUPER_CD. Args: - answer_dict (dict): The answer dictionary to update. + answer_dict: The answer dictionary to update. + constants: Application constants with grouper mappings. Returns: - dict: The updated answer dictionary. + The updated answer dictionary. """ - - grouper_cd = answer_dict.get("GROUPER_CD") - if string_utils.is_empty(grouper_cd): + grouper_cd_raw: str | None = answer_dict.get("GROUPER_CD") + if string_utils.is_empty(grouper_cd_raw) or grouper_cd_raw is None: return answer_dict grouper_type = answer_dict.get("GROUPER_TYPE") if string_utils.is_empty(grouper_type): return answer_dict grouper_descs = [] - grouper_cd = eval(grouper_cd) if "[" in grouper_cd else grouper_cd + grouper_cd = json.loads(grouper_cd_raw) if "[" in grouper_cd_raw else grouper_cd_raw for code in grouper_cd: try: @@ -634,14 +750,17 @@ def fill_grouper_cd_desc(answer_dict, constants: Constants): return answer_dict -def get_implicit_runs(answer_dict): +def get_implicit_runs(answer_dict: dict) -> dict: """ - Determines which code runs should be executed implicitly based on the Claim Type and Bill Type + 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. + answer_dict: Dict with 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. + Dict with boolean values for REVENUE_CD and PROCEDURE_CD. """ run_dict = {} @@ -668,14 +787,15 @@ def get_implicit_runs(answer_dict): return run_dict -def normalize_answer_dict_codes(answer_dict): +def normalize_answer_dict_codes(answer_dict: dict) -> dict: """ Normalizes code fields in an answer dictionary to JSON list format. Args: - answer_dict (dict): Dictionary containing code fields to normalize. + answer_dict: Dictionary containing code fields to normalize. + Returns: - dict: Dictionary with normalized code fields. + Dictionary with normalized code fields. """ fields_to_normalize = ["REVENUE_CD", "PROCEDURE_CD", "CPT4_PROC_CD"] @@ -686,7 +806,331 @@ def normalize_answer_dict_codes(answer_dict): return answer_dict -def extract_codes_from_service(answer_dict, constants: Constants): +def _code_in_range(range_key: str, code: str) -> bool: + """ + Return True if code is within the range given by range_key. + + range_key format: "low-high" (e.g. "00100-01999", "A0021-A0999"). + Leaves _DESC unchanged when no mapping (e.g. NOT_ESTABLISHED). + """ + if not range_key or "-" not in range_key or not code: + return False + parts = range_key.split("-", 1) + if len(parts) != 2: + return False + low, high = parts[0].strip(), parts[1].strip() + if not low or not high: + return False + # Numeric range (e.g. 00100-01999): compare as zero-padded strings + # so 5-digit CPT and 5-digit ranges align. + if low.isdigit() and high.isdigit() and code.isdigit(): + pad = max(len(low), len(high), len(code)) + low_p = low.zfill(pad) + high_p = high.zfill(pad) + code_p = code.zfill(pad) + return low_p <= code_p <= high_p + # Alphanumeric (e.g. A0021-A0999): lexicographic comparison. + return low <= code <= high + + +def _lookup_procedure_description(code: str, constants: Constants) -> str: + """ + Look up description for a procedure code (or range string). + + Tries exact match in full CPT/HCPCS, then exact key in level mappings, + then code-in-range in level mappings. Returns empty string if no mapping. + """ + if not code or not isinstance(code, str): + return "" + code = code.strip() + if code in ("NOT_ESTABLISHED", ""): + return "" + # Exact match in full mappings first. + if hasattr(constants, "CPT_FULL_MAPPING") and code in constants.CPT_FULL_MAPPING: + return constants.CPT_FULL_MAPPING[code] + if ( + hasattr(constants, "HCPCS_FULL_MAPPING") + and code in constants.HCPCS_FULL_MAPPING + ): + return constants.HCPCS_FULL_MAPPING[code] + # Exact range key in level mappings (e.g. "00100-01999"). + for mapping in ( + getattr(constants, "CPT_LEVEL2_MAPPING", {}), + getattr(constants, "CPT_LEVEL1_MAPPING", {}), + getattr(constants, "HCPCS_LEVEL2_MAPPING", {}), + getattr(constants, "HCPCS_LEVEL1_MAPPING", {}), + ): + if code in mapping: + return mapping[code] + # Code-in-range in level mappings. + for mapping in ( + getattr(constants, "CPT_LEVEL2_MAPPING", {}), + getattr(constants, "CPT_LEVEL1_MAPPING", {}), + getattr(constants, "HCPCS_LEVEL2_MAPPING", {}), + getattr(constants, "HCPCS_LEVEL1_MAPPING", {}), + ): + for range_key, desc in mapping.items(): + if _code_in_range(range_key, code): + return desc + return "" + + +def _lookup_revenue_description(code: str, constants: Constants) -> str: + """Look up description for a revenue code. Returns empty string if no mapping.""" + if not code or not isinstance(code, str): + return "" + code = str(code).strip() + rev = getattr(constants, "REV_MAPPING", None) + return rev.get(code, "") if rev else "" + + +# Module-level cache for valid code sets, keyed by id(constants). +# Constants don't change during a run, so we build each set once. +_VALID_PROC_CACHE: dict[int, set] = {} +_VALID_REV_CACHE: dict[int, set] = {} + + +def _valid_procedure_codes_set(constants: Constants) -> set: + """Build (or return cached) set of valid procedure codes: all mapping keys.""" + key = id(constants) + if key in _VALID_PROC_CACHE: + return _VALID_PROC_CACHE[key] + valid: set[str] = set() + for attr in ( + "CPT_FULL_MAPPING", + "HCPCS_FULL_MAPPING", + "CPT_LEVEL1_MAPPING", + "CPT_LEVEL2_MAPPING", + "HCPCS_LEVEL1_MAPPING", + "HCPCS_LEVEL2_MAPPING", + ): + m = getattr(constants, attr, None) + if m and isinstance(m, dict): + valid.update(m.keys()) + _VALID_PROC_CACHE[key] = valid + return valid + + +def _matches_cpt_or_hcpcs_format(s: str) -> bool: + """ + Return True only if the code matches CPT (5 alphanumeric) or HCPCS (1 letter + 4 alphanumeric). + + Rejects ICD-10-PCS (e.g. 7-char 0E0V0CZ) that can otherwise pass code-in-range + via lexicographic inclusion in ranges like 0001F-9007F (Category II). + """ + if not s or len(s) != 5: + return False + # HCPCS: 1 letter + 4 alphanumeric. + if s[0].isalpha() and s[1:].isalnum(): + return True + # CPT: 5 digits or 5 alphanumeric (e.g. 00100, 0124A, 0001F). + if s.isalnum(): + return True + return False + + +def _procedure_code_or_range_format_valid(s: str) -> bool: + """ + Return True if s is a valid procedure code or range format for explicit acceptance. + + Accepts single 5-char CPT/HCPCS codes and range strings (e.g. 10021-69990, A0001-A9999). + Used so we keep format-valid codes in explicit output; mapping check only tags Unmatched. + """ + if not s or not isinstance(s, str): + return False + s = s.strip() + if _matches_cpt_or_hcpcs_format(s): + return True + if "-" in s and len(s) >= 11: + parts = s.split("-", 1) + if len(parts) == 2: + low, high = parts[0].strip(), parts[1].strip() + if len(low) == 5 and len(high) == 5 and low.isalnum() and high.isalnum(): + return True + return False + + +def _revenue_code_format_valid(s: str) -> bool: + """True if s looks like a revenue code (4 digits).""" + if not s or not isinstance(s, str): + return False + s = s.strip() + return len(s) == 4 and s.isdigit() + + +def _valid_revenue_codes_set(constants: Constants) -> set: + """Build (or return cached) set of valid revenue codes from rev mappings.""" + key = id(constants) + if key in _VALID_REV_CACHE: + return _VALID_REV_CACHE[key] + valid: set[str] = set() + for attr in ("REV_MAPPING", "REV_LEVEL1_MAPPING", "REV_LEVEL2_MAPPING"): + rev = getattr(constants, attr, None) + if rev and isinstance(rev, dict): + valid.update(rev.keys()) + _VALID_REV_CACHE[key] = valid + return valid + + +def _has_unmatched_codes(code_answer_dict: dict, constants: Constants) -> bool: + """ + Return True if any procedure or revenue code is format-valid but not an + explicit key in our mappings (codes and range keys count as mapped). + Single source of truth for unmatched check; used by both explicit and + implicit paths. + """ + valid_proc = _valid_procedure_codes_set(constants) + valid_rev = _valid_revenue_codes_set(constants) + has_unmatched = False + + proc = code_answer_dict.get("PROCEDURE_CD") + if proc is not None: + if isinstance(proc, list): + proc_list = proc + elif isinstance(proc, str) and proc.startswith("["): + try: + proc_list = json.loads(proc) + except (json.JSONDecodeError, TypeError): + proc_list = [proc] + else: + proc_list = [proc] if proc else [] + for c in proc_list: + s = str(c).strip() if c is not None else "" + if _procedure_code_or_range_format_valid(s) and s not in valid_proc: + has_unmatched = True + break + + if not has_unmatched: + rev = code_answer_dict.get("REVENUE_CD") + if rev is not None: + if isinstance(rev, list): + rev_list = rev + elif isinstance(rev, str) and rev.startswith("["): + try: + rev_list = json.loads(rev) + except (json.JSONDecodeError, TypeError): + rev_list = [rev] + else: + rev_list = [rev] if rev else [] + for c in rev_list: + s = str(c).strip() if c is not None else "" + if _revenue_code_format_valid(s) and s not in valid_rev: + has_unmatched = True + break + + return has_unmatched + + +def validate_explicit_codes( + code_answer_dict: dict, constants: Constants +) -> tuple[dict, bool]: + """ + Filter code fields by format only; keep all format-valid codes. Do not drop + codes for failing mapping check; use mapping only to set has_unmatched for + the "-Unmatched" methodology tag. Returns (filtered_dict, has_unmatched). + """ + filtered = dict(code_answer_dict) + + proc = code_answer_dict.get("PROCEDURE_CD") + if proc is not None: + if isinstance(proc, list): + proc_list = proc + elif isinstance(proc, str) and proc.startswith("["): + try: + proc_list = json.loads(proc) + except (json.JSONDecodeError, TypeError): + proc_list = [proc] + else: + proc_list = [proc] if proc else [] + kept = [] + for c in proc_list: + s = str(c).strip() if c is not None else "" + if ( + s in ("", "NOT_ESTABLISHED") + or (isinstance(s, str) and "Category:" in s) + or _procedure_code_or_range_format_valid(s) + ): + kept.append(c) + filtered["PROCEDURE_CD"] = kept + + rev = code_answer_dict.get("REVENUE_CD") + if rev is not None: + if isinstance(rev, list): + rev_list = rev + elif isinstance(rev, str) and rev.startswith("["): + try: + rev_list = json.loads(rev) + except (json.JSONDecodeError, TypeError): + rev_list = [rev] + else: + rev_list = [rev] if rev else [] + kept = [] + for c in rev_list: + s = str(c).strip() if c is not None else "" + if _revenue_code_format_valid(s) or s == "": + kept.append(c) + filtered["REVENUE_CD"] = kept + + has_unmatched = _has_unmatched_codes(filtered, constants) + return filtered, has_unmatched + + +def fill_code_descriptions_from_mappings( + answer_dict: dict, constants: Constants +) -> None: + """ + Fill missing code descriptions from mapping CSVs (in-place). + + For each code field (PROCEDURE_CD, REVENUE_CD) that has a corresponding _DESC + field and for which we have a mapping, set the _DESC entry. One description + per code in the code list. Leaves _DESC unchanged where we have no mapping + (e.g. NOT_ESTABLISHED, unmapped codes). + """ + # PROCEDURE_CD -> PROCEDURE_CD_DESC (one description per code; leave unchanged if no mapping) + proc_codes = answer_dict.get("PROCEDURE_CD") + if proc_codes is not None: + proc_list = ( + proc_codes + if isinstance(proc_codes, list) + else string_utils.normalize_to_json_list(proc_codes) + ) + existing_descs = answer_dict.get("PROCEDURE_CD_DESC") + if isinstance(existing_descs, list) and len(existing_descs) == len(proc_list): + descs = list(existing_descs) + else: + descs = [""] * len(proc_list) if proc_list else [] + for i, code in enumerate(proc_list): + if not isinstance(code, str): + continue + code = code.strip() + if code in ("NOT_ESTABLISHED", ""): + continue + if i < len(descs) and string_utils.is_empty(descs[i]): + descs[i] = _lookup_procedure_description(code, constants) + answer_dict["PROCEDURE_CD_DESC"] = descs + + # REVENUE_CD -> REVENUE_CD_DESC + rev_codes = answer_dict.get("REVENUE_CD") + if rev_codes is not None: + rev_list = ( + rev_codes + if isinstance(rev_codes, list) + else string_utils.normalize_to_json_list(rev_codes) + ) + existing_rev_descs = answer_dict.get("REVENUE_CD_DESC") + if isinstance(existing_rev_descs, list) and len(existing_rev_descs) == len( + rev_list + ): + rev_descs = list(existing_rev_descs) + else: + rev_descs = [""] * len(rev_list) if rev_list else [] + for i, code in enumerate(rev_list): + if i < len(rev_descs) and string_utils.is_empty(rev_descs[i]): + rev_descs[i] = _lookup_revenue_description(str(code).strip(), constants) + answer_dict["REVENUE_CD_DESC"] = rev_descs + + +def extract_codes_from_service(answer_dict: dict, constants: Constants) -> dict: """ Extracts codes from the service term in the answer dictionary. @@ -751,8 +1195,11 @@ def extract_codes_from_service(answer_dict, constants: Constants): answer_dict["CODE_METHODOLOGY"] = "Generic - Before Prompts" return normalize_answer_dict_codes(answer_dict) - # Explicit Codes (run always) + # Explicit Codes (run always). Format-only validation; tag Unmatched if not in mapping. code_answer_dict = code_explicit(service_clean, methodology, filename) + code_answer_dict, has_unmatched = validate_explicit_codes( + code_answer_dict, constants + ) if any( not string_utils.is_empty(value) for value in code_answer_dict.values() ) and all( @@ -762,75 +1209,147 @@ def extract_codes_from_service(answer_dict, constants: Constants): code_answer_dict["CODE_METHODOLOGY"] = "Implicit - Not Established" code_answer_dict["PROCEDURE_CD"] = "[]" else: - code_answer_dict["CODE_METHODOLOGY"] = "Explicit" + code_answer_dict["CODE_METHODOLOGY"] = ( + "Explicit - Unmatched" if has_unmatched else "Explicit" + ) answer_dict.update(code_answer_dict) + fill_code_descriptions_from_mappings(answer_dict, constants) return normalize_answer_dict_codes(answer_dict) - # Implicit Codes: Code Categories - if any(["Category:" in v for v in code_answer_dict.get("PROCEDURE_CD", "")]): - code_answer_dict = code_category( - service_clean, - code_answer_dict.get("PROCEDURE_CD", ""), - constants.HCPCS_LEVEL2_MAPPING, - "", - ) - if any( - not string_utils.is_empty(value) for value in code_answer_dict.values() - ): # if any code value is not empty, return - code_answer_dict["CODE_METHODOLOGY"] = "Implicit - Letter Category" - answer_dict.update(code_answer_dict) - return normalize_answer_dict_codes(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 normalize_answer_dict_codes(answer_dict) - - # Implicit Codes: Levels 1 and 2 - code_answer_dict = code_implicit_rag( - service_clean, implicit_run_dict, "", constants + # Implicit Codes: build case (Category + Special + RAG) then arbitrate (Fix 4c + 7) + candidates = build_implicit_candidates( + service_clean, + code_answer_dict, + implicit_run_dict, + "", + constants, ) - if any( - not string_utils.is_empty(value) for value in code_answer_dict.values() - ): # if any code value is not empty, return - answer_dict.update(code_answer_dict) + chosen = code_implicit_arbitration(service_clean, candidates, "") + if chosen: + if _has_unmatched_codes(chosen, constants): + chosen["CODE_METHODOLOGY"] = ( + chosen.get("CODE_METHODOLOGY", "") + " - Unmatched" + ) + answer_dict.update(chosen) + fill_code_descriptions_from_mappings(answer_dict, constants) return normalize_answer_dict_codes(answer_dict) - # No Match - Why? Generic or Specific + # No Match - Why? Generic or Specific (code_last_check) 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" + fill_code_descriptions_from_mappings(answer_dict, constants) return normalize_answer_dict_codes(answer_dict) -def code_breakout(merged_results: pd.DataFrame, constants: Constants): +# Code fields set by extract_codes_from_service; spread to all rows in same group. +CODE_FIELDS_TO_SPREAD = { + "PROCEDURE_CD", + "PROCEDURE_CD_DESC", + "CPT4_PROC_MOD", + "CPT4_PROC_MOD_DESC", + "REVENUE_CD", + "REVENUE_CD_DESC", + "DIAG_CD", + "DIAG_CD_DESC", + "GROUPER_CD", + "GROUPER_CD_DESC", + "NDC_CD", + "NDC_CD_DESC", + "CLAIM_ADMIT_TYPE_CD", + "AUTH_ADMIT_TYPE_DESC", + "CLAIM_STATUS_CD", + "CLAIM_STATUS_CD_DESC", + "CODE_METHODOLOGY", +} + + +def _grouping_key(answer_dict: dict) -> tuple: """ - Processes a list of answer dictionaries to extract and fill in code-related information. - Wrapper function to be run from main Doczy + Return (service_term, bill_type_desc, claim_type) for group-then-split. + + Same key => same explicit/implicit gating; extract once per key and spread. + """ + service = str(answer_dict.get("SERVICE_TERM") or "").strip() + bt = answer_dict.get("BILL_TYPE_CD_DESC") + if isinstance(bt, list): + bill_desc = (bt[0] if bt else "") or "" + else: + bill_desc = str(bt or "").strip() + claim = str(answer_dict.get("AARETE_DERIVED_CLAIM_TYPE_CD") or "").strip() + return (service, bill_desc, claim) + + +def code_breakout(merged_results: pd.DataFrame, constants: Constants) -> pd.DataFrame: + """ + Extract and fill code-related information using group-then-split. + + After fill_bill_type, groups by SERVICE_TERM + BILL_TYPE_CD_DESC + + AARETE_DERIVED_CLAIM_TYPE_CD, runs extract_codes_from_service once + per group, then assigns the result to every row in that group. Args: - merged_results (pd.DataFrame): DataFrame containing merged results from the processing pipeline + merged_results: DataFrame containing merged results from the pipeline. + constants: Application constants with mappings. + Returns: - pd.DataFrame: DataFrame with updated code fields. + DataFrame with updated code fields (same row order as input). """ - import concurrent.futures + # Register this Constants instance in the prompt-templates cache so that + # when worker threads call CODE_EXPLICIT_INSTRUCTION() -> _get_constants(), + # they get this (main-thread-created) instance instead of creating a new + # Constants() inside a thread, which can trigger PyTorch meta-tensor errors. + prompt_templates.set_constants_for_instruction_cache(constants) answer_dicts = merged_results.to_dict(orient="records") - # Parallelize code extraction - def process_single_code(answer_dict): - code_answer_dict = extract_codes_from_service(answer_dict, constants) - answer_dict.update(code_answer_dict) - return answer_dict + # Pre-pass: run fill_bill_type for any row missing BILL_TYPE_CD_DESC (before grouping). + for answer_dict in answer_dicts: + bill_type = answer_dict.get("BILL_TYPE_CD_DESC") + if string_utils.is_empty(bill_type): + service = answer_dict.get("SERVICE_TERM", "") + exhibit_text = answer_dict.get("EXHIBIT_TEXT", None) + reimb_term = answer_dict.get("REIMB_TERM", None) + fill_bill_type( + service, + answer_dict, + constants.BILL_TYPE_MAPPING, + constants.BILL_TYPE_REVERSE_MAPPING, + reimb_term=reimb_term, + exhibit_text=exhibit_text, + ) - max_workers = min(len(answer_dicts), 10) + # Group by (SERVICE_TERM, BILL_TYPE_CD_DESC, AARETE_DERIVED_CLAIM_TYPE_CD). + groups = defaultdict(list) + for idx, answer_dict in enumerate(answer_dicts): + key = _grouping_key(answer_dict) + groups[key].append((idx, answer_dict)) + + # Run extract_codes_from_service once per group; spread result to all rows in group. + def process_group(items): + indices, rows = zip(*items) + rep = rows[0] + result = extract_codes_from_service(rep, constants) + # Apply return value to rep so single-row groups and mocks get codes. + if result: + for k in CODE_FIELDS_TO_SPREAD: + if k in result: + rep[k] = result[k] + for idx, row in zip(indices, rows): + if row is not rep: + for k in CODE_FIELDS_TO_SPREAD: + if k in rep: + row[k] = rep[k] + return items + + max_workers = min(len(groups), 10) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: - answer_dicts_with_code = list(executor.map(process_single_code, answer_dicts)) + list(executor.map(process_group, groups.values())) + + answer_dicts_with_code = answer_dicts # Fill in GROUPER_CD_DESC based on GROUPER_TYPE and GROUPER_CD answer_dicts_with_code = [ @@ -857,14 +1376,16 @@ def code_breakout(merged_results: pd.DataFrame, constants: Constants): return df -def grouper_breakout(results_with_code: pd.DataFrame): +def grouper_breakout(results_with_code: pd.DataFrame) -> pd.DataFrame: """ - Processes a DataFrame to extract and fill in grouper-related information in cases where we got a grouper_cd from code_breakout - but didn't get grouper fields from methodology + grouper breakout as in special case grouper was not identified and extracted. + Fill grouper-related information for rows that have GROUPER_CD but no + GROUPER_TYPE (special case where grouper was not identified upstream). + Args: - results_with_code (pd.DataFrame): DataFrame containing the answer dictionaries with code fields + results_with_code: DataFrame with code fields from code_breakout. + Returns: - pd.DataFrame: DataFrame with updated grouper fields. + DataFrame with updated grouper fields. """ # try grouper fields again if they are empty and grouper_cd is not empty answer_dicts_with_code = results_with_code.to_dict(orient="records") diff --git a/src/constants/constants.py b/src/constants/constants.py index fad51e3..c37a539 100644 --- a/src/constants/constants.py +++ b/src/constants/constants.py @@ -125,7 +125,26 @@ class Constants: ) .mapping ) - # self.HCPCS_MAPPING = CrosswalkBuilder().from_excel(_path("constants/mapping_csvs/proc_cd/hcpcs.csv", "Code", "Description").mapping + # Full procedure code mappings (exact code to description) for validation and + # description fill. Used by code extraction validation and fill_code_descriptions. + self.CPT_FULL_MAPPING = ( + CrosswalkBuilder() + .from_excel( + _path("constants/mapping_csvs/proc_cd/cpt.csv"), + "Code", + "Description", + ) + .mapping + ) + self.HCPCS_FULL_MAPPING = ( + CrosswalkBuilder() + .from_excel( + _path("constants/mapping_csvs/proc_cd/hcpcs.csv"), + "Code", + "Description", + ) + .mapping + ) self.REV_LEVEL1_MAPPING = ( CrosswalkBuilder() diff --git a/src/prompts/prompt_templates.py b/src/prompts/prompt_templates.py index 2de5c5b..d1b2ff9 100644 --- a/src/prompts/prompt_templates.py +++ b/src/prompts/prompt_templates.py @@ -1,9 +1,9 @@ -from typing import Optional, Callable, Tuple, Any +from typing import Any, Callable, Optional, Tuple import src.config as config -from src.prompts.fieldset import FieldSet from src.constants.constants import Constants -from src.utils import string_utils, json_utils +from src.prompts.fieldset import FieldSet +from src.utils import json_utils, string_utils # JSON format instructions for standardized LLM outputs JSON_DICT_FORMAT_INSTRUCTIONS = """Return your final answer as a valid JSON dictionary with the specified field names as keys. @@ -74,6 +74,15 @@ def _create_json_list_parser( _cached_constants = None +def set_constants_for_instruction_cache(constants: Optional[Constants]) -> None: + """Set the Constants instance used by _get_constants (e.g. for tests or + contract-by-contract runs that want to inject or clear between runs). + Pass None to clear the cache so the next _get_constants() creates a new instance. + """ + global _cached_constants + _cached_constants = constants + + def _get_constants() -> Constants: """Get or create the cached Constants instance for instruction functions.""" global _cached_constants @@ -230,6 +239,9 @@ def get_cacheable_instructions(): CODE_IMPLICIT_SPECIAL_INSTRUCTION() ) cacheable_instructions["CODE_IMPLICIT"] = CODE_IMPLICIT_INSTRUCTION() + cacheable_instructions["CODE_IMPLICIT_ARBITRATION"] = ( + CODE_IMPLICIT_ARBITRATION_INSTRUCTION() + ) cacheable_instructions["CODE_LAST_CHECK"] = CODE_LAST_CHECK_INSTRUCTION() cacheable_instructions["FILL_BILL_TYPE"] = FILL_BILL_TYPE_INSTRUCTION() except Exception: @@ -1292,9 +1304,24 @@ Extract explicit procedure, revenue, diagnosis, and other healthcare codes from [INSTRUCTIONS] - 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. -- FIELD ASSIGNMENT: Assign each code to exactly one field that matches its format. Do not put diagnosis codes (ICD-10: letter followed by digits, e.g. Z00.00, A01.1) in PROCEDURE_CD. Do not put procedure codes (CPT 5 digits or HCPCS 1 letter + 4 digits) in DIAG_CD. Revenue codes (3-4 digits, may end in X) belong only in REVENUE_CD. When in doubt, use format: PROCEDURE_CD = CPT/HCPCS only; DIAG_CD = ICD-10 only. +- FIELD ASSIGNMENT (map by code system, not by the word "procedure" or + "diagnosis" in the text): + - PROCEDURE_CD (CPT4): Only CPT (5 digits or 5 alphanumeric, e.g. 99213, 0124A) + and HCPCS (1 letter + 4 alphanumeric, e.g. J1098, S2900). Put these ONLY in + PROCEDURE_CD. + - DIAG_CD: Includes ICD-10-CM (letter then digits, e.g. Z00.00) + and ICD-10-PCS (7 alphanumeric characters). When the text mentions both + diagnosis/procedure codes and uses "ICD-10" or lists codes that look like + ICD-10 (e.g. 7-character alphanumeric), put those in DIAG_CD. Do NOT put + ICD-10 codes in PROCEDURE_CD. + - When the same sentence or list contains both ICD-10-style codes and + CPT/HCPCS-style codes (5-char), assign each by format: ICD-10 to DIAG_CD, + CPT/HCPCS to PROCEDURE_CD. Do not put all codes in one field. + - REVENUE_CD: Revenue codes only (3-4 digits, may end in X). When in doubt: + PROCEDURE_CD = CPT/HCPCS only; DIAG_CD = ICD-10 only (including ICD-10-PCS). - 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 text gives a range of codes, without listing each code in the range individually, return the range in the format 'LowestCode-HighestCode'. +- If the text explicitly gives a range of codes (e.g. "Surgery codes 10021 to 69990", "codes X to Y", "X–Y"), return that range in the format 'LowestCode-HighestCode' (e.g. 10021-69990). A single range string is acceptable and preferred when the source text describes a single range; do not list every code in the range. +- If the text gives a range of codes without listing each code individually, return the range in the format 'LowestCode-HighestCode'. - 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 []. @@ -1445,6 +1472,52 @@ Here is the Service Term to analyze: return (prompt, _json_list_parser) +def CODE_IMPLICIT_ARBITRATION_INSTRUCTION() -> str: + """Static instruction for CODE_IMPLICIT_ARBITRATION prompt caching.""" + return """[OBJECTIVE] +Choose the most appropriate code-set candidate(s) for the given Service Term from the listed candidates, or indicate that none are appropriate. + +[INSTRUCTIONS] +- Each candidate has a source (Category, Special, Level 1, Level 2) and procedure/revenue codes with descriptions. +- Return the value(s) that most accurately describe the service. Prefer the candidate that is most appropriate and narrowest when one clearly fits; when two candidates are equally relevant and complementary (e.g. one Level 1 set and one Level 2 set that both apply to the service), choose the single candidate that best represents the combined intent, or the narrowest single candidate that covers the service. Do not return all candidates when that would be overly broad. +- Aim for most accurate and not too broad: avoid choosing a candidate that is redundant or overly broad if a better option exists. +- If no candidate is appropriate (e.g. all are placeholders or too broad and none match the service), return no_match. +- Return exactly one of: {"chosen_index": N} where N is the 0-based index of the chosen candidate, or {"no_match": true}. + +[OUTPUT FORMAT] +Briefly explain your answer, then return your final answer as a valid JSON dictionary with keys chosen_index (integer) or no_match (boolean). Only one key should be present. +{JSON_DICT_FORMAT_INSTRUCTIONS}""" + + +def CODE_IMPLICIT_ARBITRATION( + service: str, candidates_text: str +) -> Tuple[str, Callable[[str], dict]]: + """Build prompt and parser for implicit code arbitration. + + Call CODE_IMPLICIT_ARBITRATION_INSTRUCTION() separately for the cached instruction. + Uses the same field-aware dict parser as other code prompts; code_funcs interprets + the parsed dict (chosen_index vs no_match). + + Args: + service: The service term. + candidates_text: Formatted list of candidates (source + code sets). + + Returns: + Tuple of (prompt_string, parser_function). Parser returns normalized dict + with keys chosen_index (str) and/or no_match (str) per FIELD_FORMAT_MAPPING. + """ + prompt = f"""[SERVICE TERM] +{service.replace('"', "'")} + +[CANDIDATES] +{candidates_text} + +Choose the most appropriate candidate by index (or no_match if none are appropriate).""" + + parser = _create_json_dict_parser(field_names=["chosen_index", "no_match"]) + return (prompt, parser) + + def CODE_LAST_CHECK_INSTRUCTION() -> str: """Static instruction for CODE_LAST_CHECK prompt caching.""" return f"""[OBJECTIVE] diff --git a/src/tests/test_code_funcs.py b/src/tests/test_code_funcs.py index bc7f615..561bf3a 100644 --- a/src/tests/test_code_funcs.py +++ b/src/tests/test_code_funcs.py @@ -4,7 +4,9 @@ import unittest from unittest.mock import MagicMock, patch import pandas as pd + import src.codes.code_funcs as code_funcs +import src.prompts.prompt_templates as prompt_templates from src.constants.constants import Constants from src.constants.delimiters import Delimiter @@ -13,7 +15,10 @@ class TestCodeFuncs(unittest.TestCase): def setUp(self): """Set up test fixtures before each test method.""" - # Create a mock Constants object + # Clear the module-level valid code set caches between tests. + code_funcs._VALID_PROC_CACHE.clear() + code_funcs._VALID_REV_CACHE.clear() + self.constants = MagicMock(spec=Constants) # Configure mock constants @@ -32,7 +37,16 @@ class TestCodeFuncs(unittest.TestCase): self.constants.CPT_LEVEL2_MAPPING = {"67890": "Advanced Procedure"} self.constants.HCPCS_LEVEL2_MAPPING = {"J0002": "Advanced Drug"} - self.constants.REV_MAPPING = {"456": "Advanced Revenue"} + self.constants.REV_MAPPING = {"0456": "Advanced Revenue"} + # Full mappings for validation and description fill (Fix 1 / Fix 2) + self.constants.CPT_FULL_MAPPING = { + "12345": "Test Procedure", + "99213": "Office visit, established patient", + } + self.constants.HCPCS_FULL_MAPPING = { + "J0001": "Test Drug", + "G0525": "Care management", + } # Bill type mappings matching the actual JSON structure self.constants.BILL_TYPE_MAPPING = { @@ -56,6 +70,10 @@ class TestCodeFuncs(unittest.TestCase): mock_encoded.astype.return_value = "mock_vector" self.constants.EMBEDDING_MODEL.encode.return_value = mock_encoded + def tearDown(self): + """Clear prompt_templates constants cache so a MagicMock is not reused by other tests.""" + prompt_templates.set_constants_for_instruction_cache(None) + def test_clean_service(self): """Test the clean_service function with various inputs.""" # Test empty input @@ -229,22 +247,76 @@ class TestCodeFuncs(unittest.TestCase): "TEST SERVICE", implicit_run_dict, "test.pdf", self.constants ) - # Verify results - code_implicit_rag returns lists - self.assertEqual(result["PROCEDURE_CD"], ["12345"]) - self.assertEqual(result["PROCEDURE_CD_DESC"], ["Test Procedure"]) - self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Level 1") + # Verify results - code_implicit_rag returns list of candidates with source Level 1/Level 2 + self.assertIsInstance(result, list) + self.assertGreaterEqual(len(result), 1) + codes_found = False + for item in result: + self.assertIn("source", item) + self.assertIn("code_answer_dict", item) + d = item["code_answer_dict"] + if d.get("PROCEDURE_CD") == ["12345"]: + codes_found = True + self.assertEqual(d["PROCEDURE_CD_DESC"], ["Test Procedure"]) + break + self.assertTrue( + codes_found, "Expected PROCEDURE_CD ['12345'] in a candidate" + ) - # Test exception handling + # Test exception handling - returns list (possibly empty) mock_json_load.side_effect = Exception("Test error") result = code_funcs.code_implicit_rag( "TEST SERVICE", implicit_run_dict, "test.pdf", self.constants ) - self.assertIn("CODE_METHODOLOGY", result) + self.assertIsInstance(result, list) + + @patch("src.utils.llm_utils.invoke_claude") + @patch("src.utils.string_utils.universal_json_load") + @patch("src.codes.code_funcs.get_match_list") + def test_code_implicit_rag_invalid_service_short_circuits( + self, mock_get_match_list, mock_json_load, mock_invoke_claude + ): + """INVALID_SERVICE in RAG response is a hard stop: no codes added.""" + mock_get_match_list.return_value = (["Test Procedure"], 0.9) + mock_invoke_claude.return_value = '["INVALID_SERVICE", "Test Procedure"]' + mock_json_load.return_value = ["INVALID_SERVICE", "Test Procedure"] + implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} + + with patch( + "src.codes.code_funcs.get_embedding_levels" + ) as mock_get_embedding_levels: + mock_get_embedding_levels.return_value = ( + ["cpt_level1"], + {"12345": "Test Procedure"}, + {}, + {}, + ) + result = code_funcs.code_implicit_rag( + "TEST SERVICE", implicit_run_dict, "test.pdf", self.constants + ) + + # INVALID_SERVICE yields no candidates with codes + self.assertIsInstance(result, list) + self.assertEqual(len(result), 0) + + @patch("src.utils.llm_utils.invoke_claude") + def test_fill_bill_type_list_reverse_mapping(self, mock_invoke_claude): + """fill_bill_type handles list values in BILL_TYPE_REVERSE_MAPPING.""" + mock_invoke_claude.return_value = '["Inpatient Hospital"]' + reverse_mapping = {"Inpatient Hospital": ["11X", "12X"]} + answer_dict = {} + result = code_funcs.fill_bill_type( + "INPATIENT SERVICE", + answer_dict, + self.constants.BILL_TYPE_MAPPING, + reverse_mapping, + ) + self.assertEqual(result["BILL_TYPE_CD"], ["11X", "12X"]) @patch("src.utils.llm_utils.invoke_claude") def test_code_last_check(self, mock_invoke_claude): """Test the code_last_check function for categorizing non-matched services.""" - # Setup mocks - parser returns a list, function extracts first element + # Parser is format-aware and returns 'Specific' or 'Generic' directly. # Test specific case mock_invoke_claude.return_value = '["Specific"]' result = code_funcs.code_last_check("SPECIAL SERVICE", "test.pdf") @@ -255,10 +327,218 @@ class TestCodeFuncs(unittest.TestCase): result = code_funcs.code_last_check("GENERIC SERVICE", "test.pdf") self.assertEqual(result, "Generic") + def test_format_implicit_candidate_for_prompt(self): + """Test _format_implicit_candidate_for_prompt formats a candidate for arbitration.""" + candidate = { + "source": "Category", + "code_answer_dict": { + "PROCEDURE_CD": ["J0001", "J0002"], + "PROCEDURE_CD_DESC": ["Drug A", "Drug B"], + }, + } + result = code_funcs._format_implicit_candidate_for_prompt(candidate, 0) + self.assertIn("Candidate 0 (Source: Category)", result) + self.assertIn("PROCEDURE_CD", result) + self.assertIn("J0001", result) + + @patch("src.codes.code_funcs.code_implicit_rag") + @patch("src.codes.code_funcs.code_implicit_special") + @patch("src.codes.code_funcs.code_category") + def test_build_implicit_candidates_empty_when_none_return( + self, mock_category, mock_special, mock_rag + ): + """Test build_implicit_candidates returns empty list when no strategy returns.""" + mock_category.return_value = {} + mock_special.return_value = {} + mock_rag.return_value = [] # RAG returns list of candidates + implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} + result = code_funcs.build_implicit_candidates( + "SERVICE", + {}, + implicit_run_dict, + "test.pdf", + self.constants, + ) + self.assertEqual(result, []) + + @patch("src.codes.code_funcs.code_implicit_rag") + @patch("src.codes.code_funcs.code_implicit_special") + @patch("src.codes.code_funcs.code_category") + def test_build_implicit_candidates_category_when_explicit_has_category( + self, mock_category, mock_special, mock_rag + ): + """Test build_implicit_candidates runs category and adds candidate when explicit had Category.""" + mock_category.return_value = { + "PROCEDURE_CD": ["J0001"], + "PROCEDURE_CD_DESC": ["Drug A"], + } + mock_special.return_value = {} + mock_rag.return_value = [] # RAG returns list of candidates + implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} + code_after_explicit = {"PROCEDURE_CD": ["Category: J"]} + result = code_funcs.build_implicit_candidates( + "J CODES", + code_after_explicit, + implicit_run_dict, + "test.pdf", + self.constants, + ) + self.assertEqual(len(result), 1) + self.assertEqual(result[0]["source"], "Category") + self.assertEqual(result[0]["code_answer_dict"]["PROCEDURE_CD"], ["J0001"]) + mock_category.assert_called_once() + + @patch("src.codes.code_funcs.code_implicit_rag") + @patch("src.codes.code_funcs.code_implicit_special") + @patch("src.codes.code_funcs.code_category") + def test_build_implicit_candidates_skips_category_without_category_in_explicit( + self, mock_category, mock_special, mock_rag + ): + """Test build_implicit_candidates does not run category when explicit has no Category.""" + mock_special.return_value = { + "PROCEDURE_CD": ["J0000-J9999"], + "PROCEDURE_CD_DESC": ["Drugs"], + } + mock_rag.return_value = [] # RAG returns list of candidates + implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} + result = code_funcs.build_implicit_candidates( + "DRUG SERVICE", + {}, + implicit_run_dict, + "test.pdf", + self.constants, + ) + mock_category.assert_not_called() + self.assertEqual(len(result), 1) + self.assertEqual(result[0]["source"], "Special") + + @patch("src.codes.code_funcs.code_implicit_rag") + @patch("src.codes.code_funcs.code_implicit_special") + @patch("src.codes.code_funcs.code_category") + def test_build_implicit_candidates_collects_category_special_rag( + self, mock_category, mock_special, mock_rag + ): + """Test build_implicit_candidates collects candidates from all three when all return.""" + mock_category.return_value = { + "PROCEDURE_CD": ["A0000-A9999"], + "PROCEDURE_CD_DESC": ["A range"], + } + mock_special.return_value = { + "PROCEDURE_CD": ["J0000-J9999"], + "PROCEDURE_CD_DESC": ["Drugs"], + } + mock_rag.return_value = [ + { + "source": "Level 2", + "code_answer_dict": { + "PROCEDURE_CD": ["99213"], + "PROCEDURE_CD_DESC": ["Office visit"], + "CODE_METHODOLOGY": "Implicit - Level 2", + }, + }, + ] + implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} + code_after_explicit = {"PROCEDURE_CD": ["Category: A"]} + result = code_funcs.build_implicit_candidates( + "SERVICE", + code_after_explicit, + implicit_run_dict, + "test.pdf", + self.constants, + ) + self.assertEqual(len(result), 3) + sources = [c["source"] for c in result] + self.assertIn("Category", sources) + self.assertIn("Special", sources) + self.assertIn("Level 2", sources) + + def test_code_implicit_arbitration_zero_candidates_returns_none(self): + """Test code_implicit_arbitration returns None when candidates list is empty.""" + result = code_funcs.code_implicit_arbitration("SERVICE", [], "test.pdf") + self.assertIsNone(result) + + def test_code_implicit_arbitration_single_candidate_returns_it_no_llm(self): + """Test code_implicit_arbitration with one candidate returns it without calling LLM.""" + candidates = [ + { + "source": "Category", + "code_answer_dict": { + "PROCEDURE_CD": ["J0001"], + "PROCEDURE_CD_DESC": ["Drug A"], + }, + } + ] + result = code_funcs.code_implicit_arbitration("J CODES", candidates, "test.pdf") + self.assertIsNotNone(result) + self.assertEqual(result["PROCEDURE_CD"], ["J0001"]) + self.assertEqual( + result["CODE_METHODOLOGY"], + "Implicit - Arbitration (Category)", + ) + + @patch("src.utils.llm_utils.invoke_claude") + def test_code_implicit_arbitration_multiple_candidates_uses_chosen_index( + self, mock_invoke_claude + ): + """Test code_implicit_arbitration with 2+ candidates calls LLM and returns chosen candidate.""" + mock_invoke_claude.return_value = '{"chosen_index": 1}' + candidates = [ + { + "source": "Category", + "code_answer_dict": {"PROCEDURE_CD": ["A0000-A9999"]}, + }, + {"source": "Level 2", "code_answer_dict": {"PROCEDURE_CD": ["99213"]}}, + ] + result = code_funcs.code_implicit_arbitration( + "Office visit", candidates, "test.pdf" + ) + self.assertIsNotNone(result) + self.assertEqual(result["PROCEDURE_CD"], ["99213"]) + self.assertEqual( + result["CODE_METHODOLOGY"], + "Implicit - Arbitration (Level 2)", + ) + mock_invoke_claude.assert_called_once() + + @patch("src.utils.llm_utils.invoke_claude") + def test_code_implicit_arbitration_no_match_returns_none(self, mock_invoke_claude): + """Test code_implicit_arbitration returns None when LLM returns no_match.""" + mock_invoke_claude.return_value = '{"no_match": true}' + candidates = [ + { + "source": "Category", + "code_answer_dict": {"PROCEDURE_CD": ["A0000-A9999"]}, + }, + { + "source": "Special", + "code_answer_dict": {"PROCEDURE_CD": ["10004-69990"]}, + }, + ] + result = code_funcs.code_implicit_arbitration( + "Unclear service", candidates, "test.pdf" + ) + self.assertIsNone(result) + + @patch("src.utils.llm_utils.invoke_claude") + def test_code_implicit_arbitration_invalid_index_fallback_to_zero( + self, mock_invoke_claude + ): + """Test code_implicit_arbitration falls back to index 0 when chosen_index out of range.""" + mock_invoke_claude.return_value = '{"chosen_index": 99}' + candidates = [ + {"source": "Category", "code_answer_dict": {"PROCEDURE_CD": ["J0001"]}}, + ] + result = code_funcs.code_implicit_arbitration("SERVICE", candidates, "test.pdf") + self.assertIsNotNone(result) + self.assertEqual(result["PROCEDURE_CD"], ["J0001"]) + self.assertEqual( + result["CODE_METHODOLOGY"], + "Implicit - Arbitration (Category)", + ) + @patch("src.utils.llm_utils.invoke_claude") def test_fill_bill_type(self, mock_invoke_claude): """Test the fill_bill_type function for populating bill type information.""" - # Setup mocks - parser returns a list mock_invoke_claude.return_value = '["Inpatient Hospital"]' # Test with valid bill type @@ -353,7 +633,7 @@ class TestCodeFuncs(unittest.TestCase): ) self.assertEqual(result["PROCEDURE_CD_DESC"], "UNLISTED") - # Test category match + # Test category match (arbitration with single candidate returns it) mock_clean_service.return_value = "CLEAN SERVICE" mock_explicit.return_value = {"PROCEDURE_CD": ["Category: J"]} mock_category.return_value = { @@ -363,9 +643,11 @@ class TestCodeFuncs(unittest.TestCase): result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "J CODES"}, self.constants ) - self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Letter Category") + self.assertEqual( + result["CODE_METHODOLOGY"], "Implicit - Arbitration (Category)" + ) - # Test special case match + # Test special case match (arbitration with single candidate) mock_explicit.return_value = {} mock_implicit_special.return_value = { "PROCEDURE_CD": "J0000-J9999", @@ -374,18 +656,25 @@ class TestCodeFuncs(unittest.TestCase): result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "DRUG SERVICE"}, self.constants ) - self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Special Case") + # Range J0000-J9999 not in mock mapping so methodology gets - Unmatched + self.assertIn("Implicit - Arbitration (Special)", result["CODE_METHODOLOGY"]) + self.assertIn("Unmatched", result["CODE_METHODOLOGY"]) - # Test RAG match + # Test RAG match (arbitration with single candidate; RAG returns list with Level 1) mock_implicit_special.return_value = {} - mock_implicit_rag.return_value = { - "PROCEDURE_CD": "12345", - "CODE_METHODOLOGY": "Implicit - Level 1", - } + mock_implicit_rag.return_value = [ + { + "source": "Level 1", + "code_answer_dict": { + "PROCEDURE_CD": "12345", + "CODE_METHODOLOGY": "Implicit - Level 1", + }, + }, + ] result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "TEST SERVICE"}, self.constants ) - self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Level 1") + self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Arbitration (Level 1)") # Test no match - specific mock_implicit_rag.return_value = {} @@ -402,19 +691,305 @@ class TestCodeFuncs(unittest.TestCase): ) self.assertEqual(result["CODE_METHODOLOGY"], "Generic - No Match") + def test_code_in_range(self): + """Test _code_in_range for numeric and alphanumeric range keys.""" + self.assertTrue(code_funcs._code_in_range("00100-01999", "01234")) + self.assertTrue(code_funcs._code_in_range("00100-01999", "00100")) + self.assertTrue(code_funcs._code_in_range("00100-01999", "01999")) + self.assertFalse(code_funcs._code_in_range("00100-01999", "99999")) + self.assertFalse(code_funcs._code_in_range("00100-01999", "02000")) + self.assertTrue(code_funcs._code_in_range("A0021-A0999", "A0500")) + self.assertFalse(code_funcs._code_in_range("A0021-A0999", "B0001")) + self.assertFalse(code_funcs._code_in_range("", "12345")) + self.assertFalse(code_funcs._code_in_range("00100-01999", "")) + + def test_lookup_procedure_description(self): + """Test _lookup_procedure_description: exact match and range match.""" + desc = code_funcs._lookup_procedure_description("12345", self.constants) + self.assertEqual(desc, "Test Procedure") + desc = code_funcs._lookup_procedure_description("99213", self.constants) + self.assertEqual(desc, "Office visit, established patient") + desc = code_funcs._lookup_procedure_description( + "NOT_ESTABLISHED", self.constants + ) + self.assertEqual(desc, "") + desc = code_funcs._lookup_procedure_description("UNMAPPED99", self.constants) + self.assertEqual(desc, "") + + def test_lookup_revenue_description(self): + """Test _lookup_revenue_description.""" + desc = code_funcs._lookup_revenue_description("0456", self.constants) + self.assertEqual(desc, "Advanced Revenue") + desc = code_funcs._lookup_revenue_description("9999", self.constants) + self.assertEqual(desc, "") + + def test_fill_code_descriptions_from_mappings(self): + """Test fill_code_descriptions_from_mappings fills _DESC from mappings.""" + answer_dict = { + "PROCEDURE_CD": ["12345", "99213"], + "REVENUE_CD": ["0456"], + } + code_funcs.fill_code_descriptions_from_mappings(answer_dict, self.constants) + self.assertEqual( + answer_dict["PROCEDURE_CD_DESC"], + ["Test Procedure", "Office visit, established patient"], + ) + self.assertEqual(answer_dict["REVENUE_CD_DESC"], ["Advanced Revenue"]) + + def test_fill_code_descriptions_leave_unmapped_unchanged(self): + """Test fill leaves _DESC empty for unmapped codes and NOT_ESTABLISHED.""" + answer_dict = { + "PROCEDURE_CD": ["12345", "NOT_ESTABLISHED", "UNMAPPED99"], + "PROCEDURE_CD_DESC": ["", "", ""], + } + code_funcs.fill_code_descriptions_from_mappings(answer_dict, self.constants) + self.assertEqual(answer_dict["PROCEDURE_CD_DESC"][0], "Test Procedure") + self.assertEqual(answer_dict["PROCEDURE_CD_DESC"][1], "") + self.assertEqual(answer_dict["PROCEDURE_CD_DESC"][2], "") + + def test_validate_explicit_codes_keeps_valid_filters_invalid(self): + """Format-invalid codes dropped; format-valid preserved; has_unmatched if not in mapping.""" + code_answer_dict = { + "PROCEDURE_CD": ["12345", "99213", "garbage7"], + "REVENUE_CD": ["0456", "9999"], + } + filtered, has_unmatched = code_funcs.validate_explicit_codes( + code_answer_dict, self.constants + ) + self.assertTrue(has_unmatched) + self.assertEqual(filtered["PROCEDURE_CD"], ["12345", "99213"]) + self.assertEqual(filtered["REVENUE_CD"], ["0456", "9999"]) + + def test_validate_explicit_codes_keeps_category_and_not_established(self): + """Test validate_explicit_codes keeps Category: and NOT_ESTABLISHED; no unmatched.""" + code_answer_dict = {"PROCEDURE_CD": ["Category: J", "NOT_ESTABLISHED"]} + filtered, has_unmatched = code_funcs.validate_explicit_codes( + code_answer_dict, self.constants + ) + self.assertFalse(has_unmatched) + self.assertEqual(filtered["PROCEDURE_CD"], ["Category: J", "NOT_ESTABLISHED"]) + + def test_validate_explicit_codes_all_valid(self): + """Test validate_explicit_codes returns has_unmatched False when all in mapping.""" + code_answer_dict = {"PROCEDURE_CD": ["12345"], "REVENUE_CD": ["0456"]} + filtered, has_unmatched = code_funcs.validate_explicit_codes( + code_answer_dict, self.constants + ) + self.assertFalse(has_unmatched) + self.assertEqual(filtered["PROCEDURE_CD"], ["12345"]) + self.assertEqual(filtered["REVENUE_CD"], ["0456"]) + + def test_validate_explicit_codes_multiple_and_range(self): + """Range key that is an explicit key in mapping is kept and counts as mapped.""" + self.constants.CPT_LEVEL1_MAPPING["00100-01999"] = "Anesthesia" + code_answer_dict = { + "PROCEDURE_CD": ["12345", "99213", "00100-01999"], + "REVENUE_CD": ["0456"], + } + filtered, has_unmatched = code_funcs.validate_explicit_codes( + code_answer_dict, self.constants + ) + self.assertFalse(has_unmatched) + self.assertEqual(filtered["PROCEDURE_CD"], ["12345", "99213", "00100-01999"]) + self.assertEqual(filtered["REVENUE_CD"], ["0456"]) + + def test_validate_explicit_codes_range_not_in_mapping_kept_tagged_unmatched(self): + """Format-valid range not in mapping is kept; has_unmatched True for tag.""" + # Mock has 99213 in CPT_FULL_MAPPING but not 10021-69990; range is kept and tagged + code_answer_dict = { + "PROCEDURE_CD": ["99213", "10021-69990"], + "REVENUE_CD": ["0456"], + } + filtered, has_unmatched = code_funcs.validate_explicit_codes( + code_answer_dict, self.constants + ) + self.assertTrue(has_unmatched) + self.assertEqual(filtered["PROCEDURE_CD"], ["99213", "10021-69990"]) + self.assertEqual(filtered["REVENUE_CD"], ["0456"]) + + def test_validate_explicit_codes_single_code_unmapped(self): + """Format-valid code not in exact mapping is preserved and has_unmatched True.""" + code_answer_dict = {"PROCEDURE_CD": ["01234"], "REVENUE_CD": []} + filtered, has_unmatched = code_funcs.validate_explicit_codes( + code_answer_dict, self.constants + ) + self.assertTrue(has_unmatched) + self.assertEqual(filtered["PROCEDURE_CD"], ["01234"]) + + def test_validate_explicit_codes_rejects_icd10_pcs_in_procedure(self): + """ICD-10-PCS (7-char) filtered from PROCEDURE_CD; CPT/HCPCS format kept; S2900 in mapping.""" + self.constants.HCPCS_FULL_MAPPING["S2900"] = "Robotic surgical system" + code_answer_dict = { + "PROCEDURE_CD": ["0E0V0CZ", "0E0W0CZ", "S2900"], + "REVENUE_CD": [], + } + filtered, has_unmatched = code_funcs.validate_explicit_codes( + code_answer_dict, self.constants + ) + self.assertFalse(has_unmatched) + self.assertEqual(filtered["PROCEDURE_CD"], ["S2900"]) + + def test_has_unmatched_codes_true_for_unmapped_procedure(self): + """_has_unmatched_codes returns True when procedure code is format-valid but not in mapping.""" + d = {"PROCEDURE_CD": ["99999"], "REVENUE_CD": []} + self.assertTrue(code_funcs._has_unmatched_codes(d, self.constants)) + + def test_has_unmatched_codes_true_for_unmapped_revenue(self): + """_has_unmatched_codes returns True when revenue code is format-valid but not in mapping.""" + d = {"PROCEDURE_CD": [], "REVENUE_CD": ["9999"]} + self.assertTrue(code_funcs._has_unmatched_codes(d, self.constants)) + + def test_has_unmatched_codes_false_when_all_mapped(self): + """_has_unmatched_codes returns False when all procedure/revenue codes are in mapping.""" + d = {"PROCEDURE_CD": ["99213"], "REVENUE_CD": ["0456"]} + self.assertFalse(code_funcs._has_unmatched_codes(d, self.constants)) + + @patch("src.codes.code_funcs.code_explicit") + def test_extract_codes_explicit_unmatched(self, mock_explicit): + """Explicit path sets Unmatched when any format-valid code not in mapping; no retry.""" + mock_explicit.return_value = { + "PROCEDURE_CD": ["12345", "99999"], + "REVENUE_CD": [], + } + with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"): + with patch("src.codes.code_funcs.fill_bill_type", return_value={}): + with patch( + "src.codes.code_funcs.get_implicit_runs", + return_value={"PROCEDURE_CD": True, "REVENUE_CD": True}, + ): + result = code_funcs.extract_codes_from_service( + {"SERVICE_TERM": "TEST"}, self.constants + ) + self.assertEqual(result["CODE_METHODOLOGY"], "Explicit - Unmatched") + self.assertIn("12345", str(result["PROCEDURE_CD"])) + self.assertIn("99999", str(result["PROCEDURE_CD"])) + self.assertEqual(mock_explicit.call_count, 1) + + @patch("src.codes.code_funcs.code_explicit") + def test_extract_codes_explicit_description_filled(self, mock_explicit): + """Test explicit path fills PROCEDURE_CD_DESC from mapping.""" + mock_explicit.return_value = { + "PROCEDURE_CD": ["99213"], + "REVENUE_CD": ["0456"], + } + with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"): + with patch("src.codes.code_funcs.fill_bill_type", return_value={}): + with patch( + "src.codes.code_funcs.get_implicit_runs", + return_value={"PROCEDURE_CD": True, "REVENUE_CD": True}, + ): + result = code_funcs.extract_codes_from_service( + {"SERVICE_TERM": "Office visit"}, self.constants + ) + self.assertEqual(result["CODE_METHODOLOGY"], "Explicit") + self.assertIn("99213", str(result["PROCEDURE_CD"])) + self.assertIn("Office visit", str(result["PROCEDURE_CD_DESC"])) + self.assertIn("Advanced Revenue", str(result["REVENUE_CD_DESC"])) + + @patch("src.codes.code_funcs.code_implicit_arbitration") + @patch("src.codes.code_funcs.build_implicit_candidates") + @patch("src.codes.code_funcs.code_explicit") + def test_extract_codes_implicit_arbitration_unmatched( + self, mock_explicit, mock_build_candidates, mock_arbitration + ): + """Implicit path appends ' - Unmatched' when chosen candidate has unmapped code.""" + mock_explicit.return_value = { + "PROCEDURE_CD": ["Category: J"], + "REVENUE_CD": [], + } + mock_build_candidates.return_value = [ + { + "source": "RAG", + "code_answer_dict": { + "PROCEDURE_CD": ["99999"], + "REVENUE_CD": [], + }, + } + ] + chosen = { + "PROCEDURE_CD": ["99999"], + "REVENUE_CD": [], + "CODE_METHODOLOGY": "Implicit - Arbitration (RAG)", + } + mock_arbitration.return_value = chosen + with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"): + with patch("src.codes.code_funcs.fill_bill_type", return_value={}): + with patch( + "src.codes.code_funcs.get_implicit_runs", + return_value={"PROCEDURE_CD": True, "REVENUE_CD": True}, + ): + result = code_funcs.extract_codes_from_service( + {"SERVICE_TERM": "J CODES"}, self.constants + ) + self.assertEqual( + result["CODE_METHODOLOGY"], "Implicit - Arbitration (RAG) - Unmatched" + ) + self.assertIn("99999", str(result["PROCEDURE_CD"])) + + @patch("src.codes.code_funcs.code_implicit_arbitration") + @patch("src.codes.code_funcs.build_implicit_candidates") + @patch("src.codes.code_funcs.code_explicit") + def test_extract_codes_implicit_arbitration_all_mapped( + self, mock_explicit, mock_build_candidates, mock_arbitration + ): + """Implicit path does not append ' - Unmatched' when all codes are in mapping.""" + mock_explicit.return_value = { + "PROCEDURE_CD": ["Category: J"], + "REVENUE_CD": [], + } + mock_build_candidates.return_value = [ + { + "source": "RAG", + "code_answer_dict": { + "PROCEDURE_CD": ["99213"], + "REVENUE_CD": ["0456"], + }, + } + ] + chosen = { + "PROCEDURE_CD": ["99213"], + "REVENUE_CD": ["0456"], + "CODE_METHODOLOGY": "Implicit - Arbitration (RAG)", + } + mock_arbitration.return_value = chosen + with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"): + with patch("src.codes.code_funcs.fill_bill_type", return_value={}): + with patch( + "src.codes.code_funcs.get_implicit_runs", + return_value={"PROCEDURE_CD": True, "REVENUE_CD": True}, + ): + result = code_funcs.extract_codes_from_service( + {"SERVICE_TERM": "J CODES"}, self.constants + ) + self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Arbitration (RAG)") + self.assertIn("99213", str(result["PROCEDURE_CD"])) + self.assertIn("0456", str(result["REVENUE_CD"])) + @patch("src.codes.code_funcs.extract_codes_from_service") def test_code_breakout(self, mock_extract_codes): """Test the code_breakout function for processing DataFrame records.""" - # Create test data + # BILL_TYPE_CD_DESC set so pre-pass skips fill_bill_type (avoids LLM on CI). test_df = pd.DataFrame( - [{"SERVICE_TERM": "SERVICE A"}, {"SERVICE_TERM": "SERVICE B"}] + [ + { + "SERVICE_TERM": "SERVICE A", + "BILL_TYPE_CD_DESC": "Inpatient Hospital", + }, + { + "SERVICE_TERM": "SERVICE B", + "BILL_TYPE_CD_DESC": "Inpatient Hospital", + }, + ] ) - # Setup mocks - extract_codes_from_service is called directly - mock_extract_codes.side_effect = [ - {"SERVICE_TERM": "SERVICE A", "PROCEDURE_CD": ["12345"]}, - {"SERVICE_TERM": "SERVICE B", "PROCEDURE_CD": ["67890"]}, - ] + # Return value by SERVICE_TERM so order is correct under ThreadPoolExecutor. + def return_by_service(rep, constants): + st = rep.get("SERVICE_TERM", "") + if st == "SERVICE A": + return {"SERVICE_TERM": "SERVICE A", "PROCEDURE_CD": ["12345"]} + return {"SERVICE_TERM": "SERVICE B", "PROCEDURE_CD": ["67890"]} + + mock_extract_codes.side_effect = return_by_service # Test function result_df = code_funcs.code_breakout(test_df, self.constants) diff --git a/src/tests/test_prompt_caching.py b/src/tests/test_prompt_caching.py index dd2ac2d..5ff4a2c 100644 --- a/src/tests/test_prompt_caching.py +++ b/src/tests/test_prompt_caching.py @@ -175,6 +175,30 @@ class TestPromptTemplatesReturnStrings(unittest.TestCase): self.assertIn("Lab Services", prompt_str) self.assertIn("Medicare Fee Schedule", prompt_str) + def test_code_implicit_arbitration_returns_string(self): + """Test CODE_IMPLICIT_ARBITRATION returns a (prompt_string, parser) tuple. + Instruction is in CODE_IMPLICIT_ARBITRATION_INSTRUCTION() for caching. + """ + candidates_text = 'Candidate 0 (Source: Category):\n PROCEDURE_CD: ["J0001"]' + result = prompt_templates.CODE_IMPLICIT_ARBITRATION( + "Office visit", candidates_text + ) + + self.assertIsInstance(result, tuple) + self.assertEqual(len(result), 2) + prompt_str, parser = result + self.assertIsInstance(prompt_str, str) + self.assertIn("Office visit", prompt_str) + self.assertIn("CANDIDATES", prompt_str) + self.assertIn("Candidate 0 (Source: Category)", prompt_str) + self.assertTrue(callable(parser)) + parsed = parser('{"chosen_index": 0}') + self.assertIsInstance(parsed, dict) + self.assertIn("chosen_index", parsed) + parsed_no_match = parser('{"no_match": true}') + self.assertIsInstance(parsed_no_match, dict) + self.assertIn("no_match", parsed_no_match) + class TestInstructionFunctions(unittest.TestCase): """Test that _INSTRUCTION() functions return static instruction text.""" @@ -322,6 +346,14 @@ class TestInstructionFunctions(unittest.TestCase): self.assertIn("REVENUE_CD", result) self.assertIn("DIAG_CD", result) + def test_code_implicit_arbitration_instruction_returns_string(self): + """Test CODE_IMPLICIT_ARBITRATION_INSTRUCTION returns a string for caching.""" + result = prompt_templates.CODE_IMPLICIT_ARBITRATION_INSTRUCTION() + self.assertIsInstance(result, str) + self.assertIn("OBJECTIVE", result) + self.assertIn("chosen_index", result) + self.assertIn("no_match", result) + class TestInstructionFieldContent(unittest.TestCase): """Test that _INSTRUCTION() functions contain expected static field definitions.""" @@ -439,6 +471,7 @@ class TestGetCacheableInstructions(unittest.TestCase): "DERIVED_TERM_DATE", "CHECK_PROVIDER_NAME_MATCH", "SPECIAL_CASE_ASSIGNMENT", + "CODE_IMPLICIT_ARBITRATION", ] for key in expected_keys: