jacob-mathison/doczyai-pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Merged in feature/code_optimization (pull request #903)

Browse Source 
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
This commit is contained in:
Praneel Panchigar
2026-03-11 17:39:31 +00:00
committed by Katon Minhas
parent ea1f4a2c8c
commit 863c9db9a1
5 changed files with 1500 additions and 279 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/codes/code_funcs.py
+767 -246
View File
File diff suppressed because it is too large Load Diff
src/constants/constants.py
+20 -1
View File
@@ -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()
src/prompts/prompt_templates.py
+78 -5
View File
@@ -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]
src/tests/test_code_funcs.py
+602 -27
View File
@@ -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)
src/tests/test_prompt_caching.py
+33
View File
@@ -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: