Merged in feature/tin-enhancements (pull request #697)

Feature/tin enhancements * Add OCR-correctable TIN patterns and substitution mapping * Merge remote-tracking branch 'origin/main' into feature/tin-enhancements * Refactor OCR-correctable TIN pattern and enhance documentation for clarity * Enhance get_all_matches function with type hints and improve documentation; add get_all_matches_with_ocr for OCR-correctable matches * Add attempt_ocr_correction function for TIN OCR error handling and validation * Merge remote-tracking branch 'origin/main' into feature/tin-enhancements * Integrate OCR correction for TIN extraction in run_provider_info_fields; log corrections for monitoring * Refactor OCR-correctable patterns in regex_patterns.py; add NPI_OCR_PATTERN and enhance documentation for clarity * Enhance OCR correction in attempt_ocr_correction to support both TIN and NPI; update validation logic for identifier lengths * Integrate NPI OCR correction in run_provider_info_fields; log corrections for monitoring * Enhance logging in run_provider_info_fields to summarize TIN and NPI extraction quality; include counts of exact and OCR-corrected identifiers * Add unit and integration tests for OCR correction * isort * Merge remote-tracking branch 'origin/main' into feature/tin-enhancements * Merged main into feature/tin-enhancements Approved-by: Katon Minhas
2025-09-05 15:15:39 +00:00
parent 457249e768
commit 306034820c
3 changed files with 390 additions and 5 deletions
@@ -6,3 +6,63 @@ TRIPLE_BACKTICK_PATTERN = r"```([^`]+)```"
 # TIN and NPI Regex Patterns
 TIN_PATTERN = r"\b\d{9}\b|\b\d{2}[-.]\d{7}\b|\b\d{3}[-.]\d{2}[-.]\d{4}\b"
 NPI_PATTERN = r"\b\d{10}\b"
+
+# OCR-Correctable Patterns
+# Handles common OCR misinterpretations where letters are mistaken for digits.
+# Common OCR substitutions: O/o→0, l/I→1, S/s→5, G/g→6, B/b→8
+
+# OCR-Correctable TIN Pattern
+# Strategy: Match 9-character sequences OR formatted sequences (with hyphens/dots)
+# containing digits and OCR-prone letters, then apply corrections post-match.
+#
+# Examples that WILL be captured and corrected:
+# - "12345678O" → "123456780" (unformatted, O at end)
+# - "O23456789" → "023456789" (unformatted, O at start)
+# - "1234O6789" → "123406789" (unformatted, O in middle)
+# - "12-345678O" → "12-3456780" (formatted XX-XXXXXXX, O at end)
+# - "O2-3456789" → "02-3456789" (formatted XX-XXXXXXX, O at start)
+# - "12O-456789" → "120-456789" (formatted XX-XXXXXXX, O in first segment)
+# - "123-4S-6789" → "123-45-6789" (formatted XXX-XX-XXXX, S in middle segment)
+#
+# Examples that will NOT match:
+# - "12345678" (only 8 characters)
+# - "1234567890" (10 characters)
+# - "123ABC789" (non-OCR letters like A, C)
+#
+# Note: This pattern is intentionally permissive. Validation happens post-match
+# by applying OCR_SUBSTITUTIONS and checking that result has exactly 9 digits.
+TIN_OCR_PATTERN = r"\b[O0lI1S5G6B8gbo\d]{9}\b|\b[O0lI1S5G6B8gbo\d]{2}[-.][O0lI1S5G6B8gbo\d]{7}\b|\b[O0lI1S5G6B8gbo\d]{3}[-.][O0lI1S5G6B8gbo\d]{2}[-.][O0lI1S5G6B8gbo\d]{4}\b"
+
+# OCR-Correctable NPI Pattern
+# Strategy: Match 10-character sequences containing digits and OCR-prone letters,
+# then apply corrections post-match to validate it becomes a proper 10-digit NPI.
+#
+# Examples that WILL be captured and corrected:
+# - "123456789O" → "1234567890" (O at end)
+# - "l234567890" → "1234567890" (l at start)
+# - "12345O7890" → "1234507890" (O in middle)
+# - "123456789l" → "1234567891" (l at end)
+#
+# Examples that will NOT match:
+# - "123456789" (only 9 characters)
+# - "12345678901" (11 characters)
+# - "123ABC7890" (non-OCR letters like A, C)
+#
+# Note: NPIs are always unformatted 10-digit numbers, so no hyphen/dot variations.
+# Validation happens post-match by applying OCR_SUBSTITUTIONS and checking
+# that result is exactly 10 digits.
+NPI_OCR_PATTERN = r"\b[O0lI1S5G6B8gbo\d]{10}\b"
+
+# OCR Substitution Mapping
+OCR_SUBSTITUTIONS = {
+    "O": "0",  # Letter O to zero
+    "o": "0",
+    "l": "1",  # Letter l/I to one
+    "I": "1",
+    "S": "5",  # Letter S to five
+    "s": "5",
+    "G": "6",  # Letter G to six
+    "g": "6",
+    "B": "8",  # Letter B to eight
+    "b": "8",
+}
@@ -12,17 +12,16 @@ from constants.delimiters import Delimiter
 from src.prompts.fieldset import Field, FieldSet


-def get_all_matches(text, pattern):
+def get_all_matches(text: str, pattern: str) -> list[str]:
    """
    Extracts all unique matches of a given pattern from the provided text.

    Args:
        text (str): The input text to search for matches.
        pattern (str): The regular expression pattern to match against the text.
-        filename (str): The name of the file being processed (not used in the function logic).

    Returns:
-        list: A list of unique matches found in the text. If the pattern is None or empty, returns an empty list.
+        list[str]: A list of unique matches found in the text. If the pattern is None or empty, returns an empty list.
    """
    if string_utils.is_empty(pattern):  # check if pattern is None or empty
        return []
@@ -30,6 +29,91 @@ def get_all_matches(text, pattern):
    return list(set(matches))  # Remove duplicates


+def get_all_matches_with_ocr(
+    text: str, exact_pattern: str, ocr_pattern: str, identifier_type: str = "TIN"
+) -> list[tuple]:
+    """Extracts all matches for a specific identifier type from the text, using both exact and OCR patterns.
+
+    Args:
+        text (str): The input text to search for matches.
+        exact_pattern (str): The regular expression pattern for exact matches.
+        ocr_pattern (str): The regular expression pattern for OCR-correctable matches.
+        identifier_type (str, optional): The type of identifier being searched for (e.g., "TIN" or "NPI"). Defaults to "TIN".
+
+    Returns:
+        list[tuple]: A list of (value, match_quality) tuples where match_quality is "EXACT" or "OCR_CORRECTED".
+    """
+    results = []
+    seen_values = set()
+
+    # Get exact matches
+    exact_matches = get_all_matches(text, exact_pattern)
+    for match in exact_matches:
+        results.append((match, "EXACT"))
+        seen_values.add(match)
+
+    # Get OCR-correctable candidates (even if we found exact matches)
+    ocr_candidates = get_all_matches(text, ocr_pattern)
+    for candidate in ocr_candidates:
+        corrected, is_valid = attempt_ocr_correction(candidate, identifier_type)
+        if is_valid and corrected not in seen_values:
+            results.append((corrected, "OCR_CORRECTED"))
+            seen_values.add(corrected)
+
+    return results
+
+
+def attempt_ocr_correction(candidate: str, identifier_type: str) -> tuple[str, bool]:
+    """Attempt OCR correction on a candidate string
+
+    Args:
+        candidate (str): The candidate string that may contain OCR errors.
+        identifier_type (str): The type of identifier (e.g., "TIN" or "NPI").
+
+    Returns:
+        tuple[str, bool]:   A tuple containing the corrected string and a boolean
+                            indicating if the correction was successful:
+                            (corrected_value, is_valid_after_correction)
+    """
+    if identifier_type not in ["TIN", "NPI"]:
+        return (
+            candidate,
+            False,
+        )  # Currently only TIN and NPI OCR correction is supported
+
+    corrected = candidate
+    corrections_made = []
+
+    # Apply OCR substitutions
+    for bad_char, good_char in regex_patterns.OCR_SUBSTITUTIONS.items():
+        if bad_char in corrected:
+            corrected = corrected.replace(bad_char, good_char)
+            corrections_made.append(f"{bad_char}→{good_char}")
+
+    # Remove formatting (hyphens, dots) and keep only digits
+    digits_only = "".join(c for c in corrected if c.isdigit())
+
+    # Validate based on identifier type
+    if identifier_type == "TIN":
+        # TIN must be exactly 9 digits
+        expected_length = 9
+    elif identifier_type == "NPI":
+        # NPI must be exactly 10 digits
+        expected_length = 10
+    else:
+        return candidate, False  # Unsupported identifier type
+
+    if len(digits_only) == expected_length:
+        if corrections_made:
+            logging.info(
+                f"OCR correction applied: {candidate} -> {corrected} ({', '.join(corrections_made)})"
+            )
+        return digits_only, True
+
+    # If not valid, return original candidate
+    return candidate, False
+
+
 def chunk_on_matches(matches: list[str], text_dict: dict) -> str:
    """
    Extracts and concatenates text from a dictionary of pages where any of the specified matches are found.
@@ -459,8 +543,41 @@ def run_provider_info_fields(
        list[dict]: A cleaned and standardized list containing provider information extracted
                    from the contract text.
    """
-    all_tins = get_all_matches(contract_text, regex_patterns.TIN_PATTERN)
-    all_npis = get_all_matches(contract_text, regex_patterns.NPI_PATTERN)
+    tin_matches = get_all_matches_with_ocr(
+        contract_text,
+        regex_patterns.TIN_PATTERN,
+        regex_patterns.TIN_OCR_PATTERN,
+        identifier_type="TIN",
+    )
+
+    npi_matches = get_all_matches_with_ocr(
+        contract_text,
+        regex_patterns.NPI_PATTERN,
+        regex_patterns.NPI_OCR_PATTERN,  # no OCR correction for NPI right now
+        identifier_type="NPI",
+    )
+
+    # Extract just the values from the (value, match_quality) tuples for downstream processing
+    all_tins = [tin for tin, quality in tin_matches]
+    all_npis = [npi for npi, quality in npi_matches]
+
+    # Log OCR corrections for monitoring
+    exact_tins = [tin for tin, quality in tin_matches if quality == "EXACT"]
+    ocr_corrected_tins = [
+        tin for tin, quality in tin_matches if quality == "OCR_CORRECTED"
+    ]
+
+    exact_npis = [npi for npi, quality in npi_matches if quality == "EXACT"]
+    ocr_corrected_npis = [
+        npi for npi, quality in npi_matches if quality == "OCR_CORRECTED"
+    ]
+
+    logging.info(
+        f"TIN extraction summary for {filename}: {len(exact_tins)} exact, {len(ocr_corrected_tins)} OCR-corrected"
+    )
+    logging.info(
+        f"NPI extraction summary for {filename}: {len(exact_npis)} exact, {len(ocr_corrected_npis)} OCR-corrected"
+    )

    # If there are no identifiers, return early with default values
    if not all_tins and not all_npis:
@@ -503,6 +620,12 @@ def run_provider_info_fields(
    one_to_one_results["PROV_INFO_JSON_FORMATTED"] = "\n".join(
        [json.dumps(provider) for provider in deduplicated_provider_info]
    )
+    one_to_one_results["TIN_EXTRACTION_QUALITY"] = (
+        f"{len(exact_tins)} exact, {len(ocr_corrected_tins)} OCR-corrected"
+    )
+    one_to_one_results["NPI_EXTRACTION_QUALITY"] = (
+        f"{len(exact_npis)} exact, {len(ocr_corrected_npis)} OCR-corrected"
+    )
    one_to_one_results = merge_provider_info(
        one_to_one_results, deduplicated_provider_info
    )  # Add Group TIN
@@ -1,8 +1,10 @@
 import json
+import re
 from unittest.mock import MagicMock, patch

 import pytest

+import constants.regex_patterns as regex_patterns
 import src.investment.tin_npi_funcs as tin_npi_funcs
 from src.prompts.fieldset import FieldSet

@@ -141,3 +143,203 @@ class TestTinNpiFuncs:
        assert "PROV_INFO_JSON" in results
        assert "PROV_INFO_JSON_FORMATTED" in results
        assert isinstance(results["PROV_INFO_JSON"], str)
+
+    def test_get_all_matches_with_ocr_exact_matches(self):
+        """Test that exact matches are found and marked correctly"""
+        text = "TIN: 123456789 and NPI: 1234567890"
+
+        tin_results = tin_npi_funcs.get_all_matches_with_ocr(
+            text, r"\b\d{9}\b", r"\b[O0lI1S5G6B8gbo\d]{9}\b", "TIN"
+        )
+
+        assert len(tin_results) == 1
+        assert tin_results[0] == ("123456789", "EXACT")
+
+    def test_get_all_matches_with_ocr_corrected_matches(self):
+        """Test that OCR-corrected matches work properly"""
+        text = "TIN: 12345678O and Provider ID: l234567890"
+
+        tin_results = tin_npi_funcs.get_all_matches_with_ocr(
+            text, r"\b\d{9}\b", r"\b[O0lI1S5G6B8gbo\d]{9}\b", "TIN"
+        )
+
+        npi_results = tin_npi_funcs.get_all_matches_with_ocr(
+            text, r"\b\d{10}\b", r"\b[O0lI1S5G6B8gbo\d]{10}\b", "NPI"
+        )
+
+        assert len(tin_results) == 1
+        assert tin_results[0] == ("123456780", "OCR_CORRECTED")
+
+        assert len(npi_results) == 1
+        assert npi_results[0] == ("1234567890", "OCR_CORRECTED")
+
+    def test_get_all_matches_with_ocr_mixed_matches(self):
+        """Test that both exact and OCR matches are found in same text"""
+        text = "Clean TIN: 123456789 and corrupted TIN: 98765432O"
+
+        results = tin_npi_funcs.get_all_matches_with_ocr(
+            text, r"\b\d{9}\b", r"\b[O0lI1S5G6B8gbo\d]{9}\b", "TIN"
+        )
+
+        assert len(results) == 2
+
+        # Sort results to ensure consistent order
+        results.sort(key=lambda x: x[0])
+        assert results[0] == ("123456789", "EXACT")
+        assert results[1] == ("987654320", "OCR_CORRECTED")
+
+    def test_get_all_matches_with_ocr_no_duplicates(self):
+        """Test that OCR correction doesn't create duplicates of exact matches"""
+        text = "TIN: 123456789 and same TIN with OCR error: 12345G789"
+
+        results = tin_npi_funcs.get_all_matches_with_ocr(
+            text, r"\b\d{9}\b", r"\b[O0lI1S5G6B8gbo\d]{9}\b", "TIN"
+        )
+
+        # Should only get one result (the exact match wins)
+        assert len(results) == 1
+        assert results[0] == ("123456789", "EXACT")
+
+    def test_attempt_ocr_correction_tin_success(self):
+        """Test successful TIN OCR correction"""
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("12345678O", "TIN")
+        assert is_valid is True
+        assert result == "123456780"
+
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("l23456789", "TIN")
+        assert is_valid is True
+        assert result == "123456789"
+
+    def test_attempt_ocr_correction_npi_success(self):
+        """Test successful NPI OCR correction"""
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("123456789O", "NPI")
+        assert is_valid is True
+        assert result == "1234567890"
+
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("l234567890", "NPI")
+        assert is_valid is True
+        assert result == "1234567890"
+
+    def test_attempt_ocr_correction_failure_cases(self):
+        """Test cases where OCR correction should fail"""
+        # Wrong length after correction
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("12345678", "TIN")
+        assert is_valid is False
+        assert result == "12345678"
+
+        # Invalid characters that aren't in OCR substitution map
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("123ABC789", "TIN")
+        assert is_valid is False
+        assert result == "123ABC789"
+
+        # Unsupported identifier type
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("123456789", "UNKNOWN")
+        assert is_valid is False
+        assert result == "123456789"
+
+    def test_attempt_ocr_correction_formatted_tins(self):
+        """Test OCR correction on formatted TINs"""
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("12-345678O", "TIN")
+        assert is_valid is True
+        assert result == "123456780"
+
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("123-4S-6789", "TIN")
+        assert is_valid is True
+        assert result == "123456789"
+
+    def test_attempt_ocr_correction_multiple_substitutions(self):
+        """Test OCR correction with multiple character substitutions"""
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("l2345678O", "TIN")
+        assert is_valid is True
+        assert result == "123456780"
+
+        result, is_valid = tin_npi_funcs.attempt_ocr_correction("lS345678O", "TIN")
+        assert is_valid is True
+        assert result == "153456780"
+
+    @patch("src.utils.llm_utils.invoke_claude")
+    def test_run_provider_info_fields_with_ocr_quality_tracking(
+        self, mock_invoke_claude, sample_text_dict
+    ):
+        """Test that quality tracking fields are added to results"""
+        mock_invoke_claude.return_value = json.dumps(
+            [
+                {
+                    "TIN": "123456789",
+                    "NPI": "1234567890",
+                    "NAME": "Test Provider",
+                    "ON_SIGNATURE_PAGE": "N",
+                }
+            ]
+        )
+
+        one_to_one_fields = FieldSet(relationship="one_to_one")
+        # Include both exact and OCR-correctable TINs
+        contract_text = "Contract with TIN 12-3456789 and corrupted TIN 98765432O"
+
+        results, fields = tin_npi_funcs.run_provider_info_fields(
+            contract_text, one_to_one_fields, sample_text_dict, "test.pdf"
+        )
+
+        # Check that quality tracking fields are present
+        assert "TIN_EXTRACTION_QUALITY" in results
+        assert "NPI_EXTRACTION_QUALITY" in results
+
+        # Check format of quality strings
+        assert "exact" in results["TIN_EXTRACTION_QUALITY"]
+        assert "OCR-corrected" in results["TIN_EXTRACTION_QUALITY"]
+
+    def test_ocr_regex_patterns(self):
+        """Test that OCR regex patterns match expected cases"""
+
+        # Test TIN OCR pattern
+        tin_ocr_cases = [
+            "12345678O",  # O at end
+            "O23456789",  # O at start
+            "1234O6789",  # O in middle
+            "12-345678O",  # Formatted with O
+            "123-4S-6789",  # Formatted with S
+        ]
+
+        for case in tin_ocr_cases:
+            matches = re.findall(regex_patterns.TIN_OCR_PATTERN, case)
+            assert len(matches) == 1, f"TIN OCR pattern should match '{case}'"
+
+        # Test NPI OCR pattern
+        npi_ocr_cases = [
+            "123456789O",  # O at end
+            "l234567890",  # l at start
+            "12345O7890",  # O in middle
+        ]
+
+        for case in npi_ocr_cases:
+            matches = re.findall(regex_patterns.NPI_OCR_PATTERN, case)
+            assert len(matches) == 1, f"NPI OCR pattern should match '{case}'"
+
+    @patch("src.utils.llm_utils.invoke_claude")
+    def test_run_provider_info_fields_no_exact_matches_but_ocr_matches(
+        self, mock_invoke_claude, sample_text_dict
+    ):
+        """Test scenario where no exact matches found but OCR matches are"""
+        mock_invoke_claude.return_value = json.dumps(
+            [
+                {
+                    "TIN": "123456780",  # Corrected from 12345678O
+                    "NPI": "1234567890",  # Corrected from 123456789O
+                    "NAME": "Test Provider",
+                    "ON_SIGNATURE_PAGE": "N",
+                }
+            ]
+        )
+
+        one_to_one_fields = FieldSet(relationship="one_to_one")
+        # Only OCR-corrupted identifiers, no clean ones
+        contract_text = "Contract with corrupted TIN 12345678O and NPI 123456789O"
+
+        results, fields = tin_npi_funcs.run_provider_info_fields(
+            contract_text, one_to_one_fields, sample_text_dict, "test.pdf"
+        )
+
+        # Should find OCR-corrected values
+        assert "0 exact, 1 OCR-corrected" in results["TIN_EXTRACTION_QUALITY"]
+        assert "0 exact, 1 OCR-corrected" in results["NPI_EXTRACTION_QUALITY"]