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/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)