Refactor - sub functions

src/pipelines/shared/extraction/tin_npi_funcs.py

@@ -565,61 +565,7 @@ def get_provider_info(
         return [{"TIN": [], "NPI": [], "NAME": ["PARSING_ERROR"]}]
 
 
-def run_provider_info_fields(
+def get_relevant_pages(all_tins, all_npis, text_dict, filename):
-    contract_text: str, text_dict: dict, filename: str, payer_name: str
-):
-    """
-    Extracts provider information from a contract text using regex-based identifier extraction,
-    processes the relevant text chunks, and cleans the results.
-    Args:
-        contract_text (str): The full text of the contract to process.
-        text_dict (dict): A dictionary containing text data, organized by pages
-        filename (str): The name of the file being processed, used for logging or reference.
-        payer_name (str): The name of the payer extracted from the full context answers.
-            This is not being used right now, but it might be useful for filtering providers later.
-    Returns:
-        list[dict]: A cleaned and standardized list containing provider information extracted
-                    from the contract text.
-    """
-    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,
-        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]
-
-    logging.debug(f"Extracted TINs: {all_tins}")
-    logging.debug(f"Extracted NPIs: {all_npis}")
-
-    # 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.debug(
-        f"TIN extraction summary for {filename}: {len(exact_tins)} exact, {len(ocr_corrected_tins)} OCR-corrected"
-    )
-    logging.debug(
-        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:
         # if there are no TINs or NPIs, we need to extract signature pages
@@ -640,7 +586,6 @@ def run_provider_info_fields(
                 for npi in all_npis
             )
         ]
-
         # Stopgap for when the TIN is written in the header or footer of every page (or nearly every page)
         if (
             len(relevant_pages) >= len(text_dict) - 1
@@ -648,9 +593,11 @@ def run_provider_info_fields(
             and (len(all_npis) == 1 or len(all_npis) == 0)
         ):
             relevant_pages = [relevant_pages[0]]
+    return relevant_pages
 
+def get_prov_info_json(relevant_pages, text_dict, filename, payer_name):
     if not relevant_pages:
-        deduplicated_provider_info = [
+        prov_info_json_clean = [
             {
                 "TIN": ["UNKNOWN"],
                 "NPI": ["UNKNOWN"],
@@ -659,23 +606,93 @@ def run_provider_info_fields(
         ]
     else:
         # Process each page separately and combine results
-        combined_provider_info = []
+        prov_info_json = []
         for page_num in relevant_pages:
             page_provider_info = get_provider_info(
                 text_dict, page_num, filename, payer_name
             )
-            combined_provider_info.extend(page_provider_info)
+            prov_info_json.extend(page_provider_info)
 
         # Clean and standardize results
-        cleaned_provider_info = clean_provider_info(combined_provider_info)
+        prov_info_json_clean = clean_provider_info(prov_info_json)
         # Deduplicate providers based on TIN and NPI
-        deduplicated_provider_info = deduplicate_providers(cleaned_provider_info)
+        prov_info_json_clean = deduplicate_providers(prov_info_json_clean)
+    return prov_info_json_clean
+
+def get_all_tins_and_npis(contract_text):
+    
+    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,
+        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]
+    logging.debug(f"Extracted TINs: {all_tins}")
+    logging.debug(f"Extracted NPIs: {all_npis}")
+    
+    return all_tins, all_npis, tin_matches, npi_matches
+
+def get_tin_extraction_quality(tin_matches, npi_matches, filename):
+    # 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.debug(
+        f"TIN extraction summary for {filename}: {len(exact_tins)} exact, {len(ocr_corrected_tins)} OCR-corrected"
+    )
+    logging.debug(
+        f"NPI extraction summary for {filename}: {len(exact_npis)} exact, {len(ocr_corrected_npis)} OCR-corrected"
+    )
+    return exact_tins, exact_npis, ocr_corrected_tins, ocr_corrected_npis
+
+
+def run_provider_info_fields(
+    contract_text: str, text_dict: dict, filename: str, payer_name: str
+):
+    """
+    Extracts provider information from a contract text using regex-based identifier extraction,
+    processes the relevant text chunks, and cleans the results.
+    Args:
+        contract_text (str): The full text of the contract to process.
+        text_dict (dict): A dictionary containing text data, organized by pages
+        filename (str): The name of the file being processed, used for logging or reference.
+        payer_name (str): The name of the payer extracted from the full context answers.
+            This is not being used right now, but it might be useful for filtering providers later.
+    Returns:
+        list[dict]: A cleaned and standardized list containing provider information extracted
+                    from the contract text.
+    """
+
+    all_tins, all_npis, tin_matches, npi_matches = get_all_tins_and_npis(contract_text)
+
+    exact_tins, exact_npis, ocr_corrected_tins, ocr_corrected_npis = get_tin_extraction_quality(tin_matches, npi_matches, filename)
+    
+    relevant_pages = get_relevant_pages(all_tins, all_npis, text_dict, filename)
+    
+    prov_info_json_clean = get_prov_info_json(relevant_pages, text_dict, filename, payer_name)
 
     # Create columns
     one_to_one_results = {}
-    one_to_one_results["PROV_INFO_JSON"] = json.dumps(deduplicated_provider_info)
+    one_to_one_results["PROV_INFO_JSON"] = json.dumps(prov_info_json_clean)
     one_to_one_results["PROV_INFO_JSON_FORMATTED"] = "\n".join(
-        [json.dumps(provider) for provider in deduplicated_provider_info]
+        [json.dumps(provider) for provider in prov_info_json_clean]
     )
     one_to_one_results["TIN_EXTRACTION_QUALITY"] = (
         f"{len(exact_tins)} exact, {len(ocr_corrected_tins)} OCR-corrected"
@@ -688,6 +705,7 @@ def run_provider_info_fields(
     filename_tin = get_all_matches(filename, regex_patterns.FILE_NAME_TIN_PATTERN)
     one_to_one_results["FILENAME_TIN"] = filename_tin[0] if filename_tin else "N/A"
 
+    print("run_provider_info_fields: ", one_to_one_results)
     return one_to_one_results