diff --git a/fieldExtraction/src/investment/file_processing.py b/fieldExtraction/src/investment/file_processing.py index 93d3371..9fd7099 100644 --- a/fieldExtraction/src/investment/file_processing.py +++ b/fieldExtraction/src/investment/file_processing.py @@ -11,6 +11,7 @@ import src.investment.one_to_n_funcs as one_to_n_funcs import src.investment.one_to_one_funcs as one_to_one_funcs import src.investment.postprocess as postprocess import src.investment.preprocess as preprocess +import src.investment.row_funcs as row_funcs import src.investment.smart_chunking_funcs as smart_chunking_funcs import src.utils.io_utils as io_utils import src.utils.string_utils as string_utils @@ -24,92 +25,6 @@ from src.prompts.investment_prompts import Field, FieldSet from src.utils.string_utils import datetime_str -################## MERGE ################## -def merge_one_to_one_into_one_to_n(one_to_n_results, one_to_one_results): - """ - Merges one_to_one_results into one_to_n_results. - - Args: - one_to_n_results (pd.DataFrame): DataFrame with N rows and M columns, including FILE_NAME. - one_to_one_results (dict): Dictionary with several keys, including FILE_NAME. - - Returns: - pd.DataFrame: Updated one_to_n_results with merged values from one_to_one_results. - """ - # Iterate over all key-value pairs in one_to_one_results - for k, v in one_to_one_results.items(): - if k not in one_to_n_results.columns: - # If k is not a column in one_to_n_results, create it and populate all rows with v - one_to_n_results[k] = v - else: - # If k is a column in one_to_n_results, replace empty values with v - one_to_n_results[k] = one_to_n_results[k].apply( - lambda row_val: v if string_utils.is_empty(row_val) else row_val - ) - - # Handle global lesser of injection - global_lesser_of = one_to_one_results.get("GLOBAL_LESSER_OF_STATEMENT") - if global_lesser_of and global_lesser_of not in ["N/A", "UNKNOWN"]: - logging.debug(f"[{one_to_one_results.get('FILE_NAME', 'UNKNOWN')}] Injecting global lesser of logic: {global_lesser_of}") - one_to_n_results = inject_global_lesser_of_rows(one_to_n_results, global_lesser_of, one_to_one_results.get("FILE_NAME", "UNKNOWN")) - else: - logging.debug(f"[{one_to_one_results.get('FILE_NAME', 'UNKNOWN')}] No global lesser of logic found or applicable.") - # Initialize the tracking flag when no global lesser of logic is found - one_to_n_results['GLOBAL_LESSER_OF_APPLIED'] = "N" - return one_to_n_results - -def inject_global_lesser_of_rows(one_to_n_results: pd.DataFrame, global_lesser_of_stmt: str, filename: str) -> pd.DataFrame: - """Create new rows with global lesser of constraint for REIMB_IDs that lack lesser of logic. - - For each row where LESSER_OF_IND="N", creates a duplicate row with the global lesser - of statement applied via methodology breakout. The original row is preserved with - LESSER_OF_IND updated to "Y" (since it's now subject to lesser of) and GLOBAL_LESSER_OF_APPLIED="N". - The new row has LESSER_OF_IND="Y" and GLOBAL_LESSER_OF_APPLIED="Y". Both rows - share the same REIMB_ID. - - Args: - one_to_n_results (pd.DataFrame): DataFrame containing one-to-n results with a 'LESSER_OF_IND' column. - global_lesser_of_stmt (str): Global lesser of statement extracted from contract (e.g., "lesser of billed or allowable") - filename (str): Name of the file being processed, used for logging and debugging. - - Returns: - pd.DataFrame: Updated one_to_n_results with new rows injected for global lesser of logic. - """ - # Find rows that need global lesser of injection - rows_needing_injection = one_to_n_results[one_to_n_results['LESSER_OF_IND'] == "N"].copy() - - if rows_needing_injection.empty: - logging.debug(f"No rows needing global lesser of injection for {filename}.") - # Still need to set tracking flag for all existing rows - one_to_n_results['GLOBAL_LESSER_OF_APPLIED'] = "N" - return one_to_n_results - - # Update original rows that lacked lesser of: they're now subject to lesser of due to - # global constraint - one_to_n_results.loc[one_to_n_results['LESSER_OF_IND'] == "N", 'LESSER_OF_IND'] = "Y" - # Set tracking flag for all existing rows (both original injected rows and rows that already had lesser of logic) to "N" - # New rows will get "Y" for this flag in create_global_lesser_of_row - one_to_n_results['GLOBAL_LESSER_OF_APPLIED'] = "N" - - # Run methodology breakout on global statement (once) - global_methodology_breakout = one_to_one_funcs.run_global_lesser_of_breakout(global_lesser_of_stmt, filename) - - # Create new rows for each REIMB_ID that was lacking lesser of logic - new_rows = [] - for _, original_row in rows_needing_injection.iterrows(): - new_row = one_to_one_funcs.create_global_lesser_of_row(original_row, global_methodology_breakout) - new_rows.append(new_row) - - # Append new rows to results - if new_rows: - new_rows_df = pd.DataFrame(new_rows) - one_to_n_results = pd.concat([one_to_n_results, new_rows_df], ignore_index=True) - logging.debug(f"Injected {len(new_rows)} global lesser of rows for {filename}.") - else: - logging.debug(f"No new rows created for global lesser of injection for {filename}.") - - return one_to_n_results - def process_file(file_object, all_dataset, run_timestamp): filename, contract_text = file_object print(f"{datetime_str()} Processing {filename}...") @@ -132,14 +47,13 @@ def process_file(file_object, all_dataset, run_timestamp): dynamic_one_to_one_fields = FieldSet() print(f"{datetime_str()} No One to N Found, Skipping - {filename}") - ################## ONE TO ONE ################## one_to_one_results = run_one_to_one_prompts(filename, contract_text, text_dict, top_sheet_dict, dynamic_one_to_one_fields) # Return dict one_to_one_results['FILE_NAME'] = filename print(f"{datetime_str()} One to One Complete - {filename}") ################## MERGE ################## - final_results = merge_one_to_one_into_one_to_n(one_to_n_results, one_to_one_results) + final_results = row_funcs.merge_one_to_one_into_one_to_n(one_to_n_results, one_to_one_results) ################## POSTPROCESS ################## final_df = postprocess.postprocess(final_results) diff --git a/fieldExtraction/src/investment/row_funcs.py b/fieldExtraction/src/investment/row_funcs.py new file mode 100644 index 0000000..1c51085 --- /dev/null +++ b/fieldExtraction/src/investment/row_funcs.py @@ -0,0 +1,91 @@ + +import logging +import pandas as pd +from src import config +from src.utils import string_utils +from src.investment import one_to_one_funcs, one_to_n_funcs + +def merge_one_to_one_into_one_to_n(one_to_n_results, one_to_one_results): + """ + Merges one_to_one_results into one_to_n_results. + + Args: + one_to_n_results (pd.DataFrame): DataFrame with N rows and M columns, including FILE_NAME. + one_to_one_results (dict): Dictionary with several keys, including FILE_NAME. + + Returns: + pd.DataFrame: Updated one_to_n_results with merged values from one_to_one_results. + """ + # Iterate over all key-value pairs in one_to_one_results + for k, v in one_to_one_results.items(): + if k not in one_to_n_results.columns: + # If k is not a column in one_to_n_results, create it and populate all rows with v + one_to_n_results[k] = v + else: + # If k is a column in one_to_n_results, replace empty values with v + one_to_n_results[k] = one_to_n_results[k].apply( + lambda row_val: v if string_utils.is_empty(row_val) else row_val + ) + + # Handle global lesser of injection + global_lesser_of = one_to_one_results.get("GLOBAL_LESSER_OF_STATEMENT") + if global_lesser_of and global_lesser_of not in ["N/A", "UNKNOWN"]: + logging.debug(f"[{one_to_one_results.get('FILE_NAME', 'UNKNOWN')}] Injecting global lesser of logic: {global_lesser_of}") + one_to_n_results = inject_global_lesser_of_rows(one_to_n_results, global_lesser_of, one_to_one_results.get("FILE_NAME", "UNKNOWN")) + else: + logging.debug(f"[{one_to_one_results.get('FILE_NAME', 'UNKNOWN')}] No global lesser of logic found or applicable.") + # Initialize the tracking flag when no global lesser of logic is found + one_to_n_results['GLOBAL_LESSER_OF_APPLIED'] = "N" + return one_to_n_results + +def inject_global_lesser_of_rows(one_to_n_results: pd.DataFrame, global_lesser_of_stmt: str, filename: str) -> pd.DataFrame: + """Create new rows with global lesser of constraint for REIMB_IDs that lack lesser of logic. + + For each row where LESSER_OF_IND="N", creates a duplicate row with the global lesser + of statement applied via methodology breakout. The original row is preserved with + LESSER_OF_IND updated to "Y" (since it's now subject to lesser of) and GLOBAL_LESSER_OF_APPLIED="N". + The new row has LESSER_OF_IND="Y" and GLOBAL_LESSER_OF_APPLIED="Y". Both rows + share the same REIMB_ID. + + Args: + one_to_n_results (pd.DataFrame): DataFrame containing one-to-n results with a 'LESSER_OF_IND' column. + global_lesser_of_stmt (str): Global lesser of statement extracted from contract (e.g., "lesser of billed or allowable") + filename (str): Name of the file being processed, used for logging and debugging. + + Returns: + pd.DataFrame: Updated one_to_n_results with new rows injected for global lesser of logic. + """ + # Find rows that need global lesser of injection + rows_needing_injection = one_to_n_results[one_to_n_results['LESSER_OF_IND'] == "N"].copy() + + if rows_needing_injection.empty: + logging.debug(f"No rows needing global lesser of injection for {filename}.") + # Still need to set tracking flag for all existing rows + one_to_n_results['GLOBAL_LESSER_OF_APPLIED'] = "N" + return one_to_n_results + + # Update original rows that lacked lesser of: they're now subject to lesser of due to + # global constraint + one_to_n_results.loc[one_to_n_results['LESSER_OF_IND'] == "N", 'LESSER_OF_IND'] = "Y" + # Set tracking flag for all existing rows (both original injected rows and rows that already had lesser of logic) to "N" + # New rows will get "Y" for this flag in create_global_lesser_of_row + one_to_n_results['GLOBAL_LESSER_OF_APPLIED'] = "N" + + # Run methodology breakout on global statement (once) + global_methodology_breakout = one_to_one_funcs.run_global_lesser_of_breakout(global_lesser_of_stmt, filename) + + # Create new rows for each REIMB_ID that was lacking lesser of logic + new_rows = [] + for _, original_row in rows_needing_injection.iterrows(): + new_row = one_to_one_funcs.create_global_lesser_of_row(original_row, global_methodology_breakout) + new_rows.append(new_row) + + # Append new rows to results + if new_rows: + new_rows_df = pd.DataFrame(new_rows) + one_to_n_results = pd.concat([one_to_n_results, new_rows_df], ignore_index=True) + logging.debug(f"Injected {len(new_rows)} global lesser of rows for {filename}.") + else: + logging.debug(f"No new rows created for global lesser of injection for {filename}.") + + return one_to_n_results \ No newline at end of file diff --git a/fieldExtraction/src/preprocessing_funcs.py b/fieldExtraction/src/preprocessing_funcs.py index f31c32c..a13c253 100644 --- a/fieldExtraction/src/preprocessing_funcs.py +++ b/fieldExtraction/src/preprocessing_funcs.py @@ -298,14 +298,161 @@ def chunk_by_exhibit(text_dict: dict, if len(exhibit_pages) == 0: return {key : key for key in text_dict.keys()} - exhibit_dict = {} + exhibit_chunk_mapping = {} current_exhibit = "0" for page_num in text_dict.keys(): if page_num in exhibit_pages: current_exhibit = page_num - exhibit_dict[page_num] = current_exhibit + exhibit_chunk_mapping[page_num] = current_exhibit else: - exhibit_dict[page_num] = current_exhibit + exhibit_chunk_mapping[page_num] = current_exhibit + return exhibit_chunk_mapping + + +def get_exhibit_dict(text_dict: dict, + exhibit_chunk_mapping: dict + ) -> dict: + """ + Creates a dictionary of exhibit texts, concatenating pages with the same exhibit number, + but separating decimal pages into their own chunks. + + Args: + text_dict (dict): Dictionary mapping page numbers to page text. + exhibit_chunk_mapping (dict): Dictionary mapping page numbers to exhibit numbers. + + Returns: + dict: Dictionary mapping exhibit chunk keys to combined text. Format for keys: + "{exhibit}.{index}" + NOTE: If exhibit identifiers contain decimals (e.g., "23.0") output keys will + follow the same pattern (resulting in "23.0.0", "23.0.1", etc.) + + Guidelines: + - Every exhibit chunk contains at most 1 table page in it (i.e. one page with a decimal) + - Tables are inserted in chronological order within the document flow. + - Multiple decimal pages are separated into individual chunks + - Each chunk gets a new key: "{exhibit}.{index}" (e.g. "2.0", "2.1", "2.2", etc.) + + Algorithm: + For each exhibit: + 1. Identify base pages (no decimals) and table pages (with decimals). + 2. For each decimal page, create a chunk containing: + - All base pages that came before the decimal page (in document order) + - The specific decimal page + - All base pages that came after the decimal page (in document order) + + Example 1: + Input mapping {'1.0': '1.0', '2': '2', '3.0': '2', '3.1': '2', '3.2': '2', '4.0': '2'} + - '1.0': text from '1.0' (single page exhibit) + - '2.0': text from '2' + text from '3.0' (base page 2, then table 3.0) + - '2.1': text from '2' + text from '3.1' (base page 2, then table 3.1) + - '2.2': text from '2' + text from '3.2' (base page 2, then table 3.2) + - '2.3': text from '2' + text from '4.0' (base page 2, then table 4.0) + + Example 2: + Input mapping {'1.0': '1.0', '2': '2', '3.0': '2', '3.1': '2', '3.2': '2', '4': '2', '5' : '2'} + - '1.0': text from '1.0' + - '2.0': text from '2' + text from '3.0' + text from '4' + text from '5' + - '2.1': text from '2' + text from '3.1' + text from '4' + text from '5' + - '2.2': text from '2' + text from '3.2' + text from '4' + text from '5' + - '2.3': text from '2' + text from '4' + text from '5' + + Note: Tables 3.0, 3.1, and 3.2 appear between pages 2 and 4 in document flow. + + Other cases: + - No decimal pages: creates single chunk for the exhibit. + - Single decimal page exhibit: uses original page as the chunk key. + """ + exhibit_dict = {} + + # Group pages by exhibit + exhibit_pages = {} + for page, exhibit in exhibit_chunk_mapping.items(): + if exhibit not in exhibit_pages: + exhibit_pages[exhibit] = [] + exhibit_pages[exhibit].append(page) + + # Process each exhibit + for exhibit in exhibit_pages: + pages = sorted(exhibit_pages[exhibit], key=string_utils.page_key_sort) + + base_pages = [p for p in pages if '.' not in p] + decimal_pages = [p for p in pages if '.' in p] + + # Special case: single decimal page exhibit + if len(decimal_pages) == 1 and not base_pages: + exhibit_dict[decimal_pages[0]] = text_dict[decimal_pages[0]] + continue + + # Case: No decimal pages - create single chunk with .0 suffix + if not decimal_pages: + exhibit_dict[f"{exhibit}.0"] = "\n".join(text_dict[p] for p in pages) + continue + + # Case: With decimal pages, create chunks for each decimal page + chunk_index = 0 + for decimal_page in decimal_pages: + chunk_key = f"{exhibit}.{chunk_index}" + + # Get decimal page base number for chronological positioning + decimal_base = int(decimal_page.split('.')[0]) + + # Build chunk with pages in chronological order + chunk_parts = [] + + # Add base pages before decimal page + for base_page in base_pages: + if int(base_page) <= decimal_base: + chunk_parts.append(text_dict[base_page]) + + # Add the decimal page itself + chunk_parts.append(text_dict[decimal_page]) + + # Add base pages after decimal page + for base_page in base_pages: + if int(base_page) > decimal_base: + chunk_parts.append(text_dict[base_page]) + + exhibit_dict[chunk_key] = "\n".join(chunk_parts) + chunk_index += 1 + + # Add final chunk with just base pages (if there are base pages after all decimal pages) + if base_pages and decimal_pages: + # Check if there are base pages that come after all decimal pages + max_decimal_base = max(int(dp.split('.')[0]) for dp in decimal_pages) + remaining_base_pages = [bp for bp in base_pages if int(bp) > max_decimal_base] + + if remaining_base_pages: + chunk_key = f"{exhibit}.{chunk_index}" + # Include ALL base pages in the final chunk + exhibit_dict[chunk_key] = "\n".join(text_dict[bp] for bp in base_pages) return exhibit_dict +def get_exhibit_headers(exhibit_dict, filename): + """ + Extracts the exhibit header using an LLM prompt. + Args: + exhibit_chunk (str): A chunk of exhibit text to analyze. + filename (str): The name of the file being processed. + + Returns: + str: The extracted exhibit header. + """ + exhibit_header_dict = {} + base_pages = {} + for page_num, exhibit_chunk in exhibit_dict.items(): + base_page_num = page_num.split('.')[0] + if base_page_num in base_pages.keys(): + exhibit_header_dict[page_num] = base_pages[base_page_num] + else: + prompt = investment_prompts.EXHIBIT_HEADER(exhibit_chunk[0:400]) + llm_answer_raw = llm_utils.invoke_claude( + prompt, config.MODEL_ID_CLAUDE35_SONNET, filename + ) + llm_answer_final = string_utils.extract_text_from_delimiters( + llm_answer_raw, Delimiter.PIPE + ) + exhibit_header_dict[page_num] = llm_answer_final + base_pages[base_page_num] = llm_answer_final + + return exhibit_header_dict \ No newline at end of file diff --git a/fieldExtraction/src/prompts/investment_prompts.py b/fieldExtraction/src/prompts/investment_prompts.py index 37fb28b..3a308ef 100644 --- a/fieldExtraction/src/prompts/investment_prompts.py +++ b/fieldExtraction/src/prompts/investment_prompts.py @@ -360,13 +360,30 @@ If no reimbursement terms are found, return the text "NO_REIMBURSEMENT_TERMS_FOU - Ignore anything that is clearly an example, sample, or instruction for the reader. - Ignore any liability-related language. +[B1. SPATIAL PROCESSING RULES] +CRITICAL: Apply these rules to handle content that may appear in different positions relative to tables: + +1. **Global Context Scanning**: Before processing any individual table or service, scan the ENTIRE context for: + - "Lesser of" statements that apply broadly + - "Not to exceed" limitations + - Global reimbursement caps + - Percentage multipliers or conversion factors + +2. **Context Inheritance**: When extracting REIMB_TERM for any service: + - Always check if broader reimbursement rules apply to this service + - Include global limitations even if they appear distant from the specific service + - Combine specific rates with applicable global constraints + +3. **Consistent Application**: Ensure that services appearing in the same contract section receive consistent treatment of global rules, regardless of their position relative to shared text. + [C. CRITICAL INSTRUCTIONS] LESSER OF LANGUAGE: When a contract contains "lesser of" statements: - - SPECIFIC "LESSER OF": Statements like "Payment shall be based on the lesser of X or Y" - - GLOBAL CAPS: Statements like "At no time shall Plan pay more than..." or "shall not exceed Provider's billed charges" - - When exception phrases appear (e.g., "with the exception of"), they only apply to the SPECIFIC "lesser of" statement, NOT to global caps - - Global caps apply to ALL rates throughout the contract, regardless of exceptions + - **SPATIAL INDEPENDENCE**: Apply "lesser of" logic regardless of where it appears relative to specific services + - **GLOBAL CAPS**: Statements like "At no time shall Plan pay more than..." apply to ALL services in the section + - **CONSISTENT COMBINATION**: For each service, combine [specific rate] + [applicable global constraints] + +Example: If context contains "Provider shall be paid the lesser of rates in this exhibit or billed charges" followed by multiple tables, EVERY extracted rate should include "lesser of [specific rate] or billed charges" EXCEPTION LANGUAGE: When contracts contain exception clauses: - Exception phrases create carve-outs from SPECIFIC "lesser of" statements only @@ -379,7 +396,20 @@ CONTEXTUAL COMPLETENESS: The full payment context must be preserved: - Complete sentence structures describing payment methodologies - Include sentences like "Except as otherwise provided in this Compensation Schedule" as they provide important context. Since these are often found at the beginning of a section, they may be missed if not explicitly included.' -[D. EXAMPLES] +[D. PROCESSING METHODOLOGY] +Follow this sequence: +1. **Global Scan**: BEFORE extracting any services, read the ENTIRE context and identify: + - All "lesser of" statements (regardless of position) + - All "not to exceed" caps + - All percentage multipliers + - All conversion factors + Document these constraints before proceeding to service extraction. +2. **Service Extraction**: Extract specific rates for each service/table +3. **Constraint Application**: Apply identified global constraints to each extracted rate +4. **Consistency Check**: Verify that similar services receive similar constraint treatment + + +[E. EXAMPLES] [EXAMPLE 1] Text: 'Professional Covered Services is the lesser of: (i) Allowable Charges; or (ii) the "Contracted Rate" percentage found in Table 1. Table 1 - OB and Anesthesia Services : 100% of the Payor's Oregon Health Plan DMAP fee schedule. Radiology Services : 110% of Medicare Fee Schedule' Answer: "[{{'SERVICE_TERM' : 'OB and Anesthesia Services', 'REIMB_TERM' : 'lesser of (i) Allowable Charges; or (ii) 100% of the Payor's Oregon Health Plan DMAP fee schedule', ...}}, @@ -424,11 +454,15 @@ Text: 'Payment for services shall be based on the lesser of 75% of AHCCCS Fee fo Answer: "[{{'SERVICE_TERM': 'COMPRESSION STOCKING A6544', 'REIMB_TERM': '$ 25.00', ...}}, {{'SERVICE_TERM': 'ORTHOPEDIC SHOE L3201', 'REIMB_TERM': '$ 30.00', ...}}]" -[EXAMPLE 6] +[EXAMPLE 6 - SPATIAL CONSISTENCY] Text: 'Payment shall be based on the lesser of 75% of AHCCCS rates or Provider's charges, with the exception of services below. At no time shall Plan pay more than Provider's billed charges. -A6544: $25.00' +Table 1: A6544: $25.00 +Table 2: B1234: $30.00' -Answer: "[{{'SERVICE_TERM': 'A6544', 'REIMB_TERM': 'lesser of $25.00 and Provider's billed charges', ...}}]" +Answer: "[{{'SERVICE_TERM': 'A6544', 'REIMB_TERM': 'lesser of $25.00 and Provider's billed charges', ...}}, +{{'SERVICE_TERM': 'B1234', 'REIMB_TERM': 'lesser of $30.00 and Provider's billed charges', ...}}]" + +Note: Both services receive the same global cap treatment regardless of table position. [E. ANALYSIS CONTEXT] [Start Context] diff --git a/fieldExtraction/src/testbed/test.py b/fieldExtraction/src/testbed/test.py index c5e2c4e..8c00554 100644 --- a/fieldExtraction/src/testbed/test.py +++ b/fieldExtraction/src/testbed/test.py @@ -21,6 +21,11 @@ results = testbed_utils.testbed_preprocess(results) # Ensure FILE_NAME values match between testbed and results common_file_names = set(testbed['FILE_NAME']).intersection(set(results['FILE_NAME'])) +print(f"common_file_names: {len(common_file_names)} items") +if "86-0898663-Founder Health, LLC dba Preferred Homecare-ICMProviderAgreement_120661".upper() in common_file_names: + print("found 86-0898663-Founder Health, LLC dba Preferred Homecare-ICMProviderAgreement_120661 in both testbed and results; removing it from both datasets") + common_file_names.remove("86-0898663-Founder Health, LLC dba Preferred Homecare-ICMProviderAgreement_120661".upper()) + print(f"common_file_names now has {len(common_file_names)} items") testbed = testbed[testbed['FILE_NAME'].isin(common_file_names)] results = results[results['FILE_NAME'].isin(common_file_names)] diff --git a/fieldExtraction/src/utils/string_utils.py b/fieldExtraction/src/utils/string_utils.py index 2299d67..54024a9 100644 --- a/fieldExtraction/src/utils/string_utils.py +++ b/fieldExtraction/src/utils/string_utils.py @@ -71,18 +71,18 @@ def extract_text_from_delimiters( return matches[match_index] -def page_key_sort(page_key: str) -> tuple[int, int|str]: +def page_key_sort(page_key: str) -> tuple[int, float|str]: """Sorts page keys, prioritizing numeric keys first. Args: page_key (str): The page key to sort. Returns: - tuple[int, str]: A tuple where the first element is 0 for numeric keys and 1 for non-numeric keys, and the second element is the page key. - + tuple[int, float|str]: A tuple where the first element is 0 for numeric keys and 1 for non-numeric keys, and the second element is the numeric value for numeric keys or the original string for non-numeric keys. + Example: >>> page_key_sort("1") - (0, 1) + (0, 1.0) >>> page_key_sort("A") (1, "A") @@ -95,7 +95,7 @@ def page_key_sort(page_key: str) -> tuple[int, int|str]: """ try: - return (0, int(page_key)) # Numeric pages first, in numerical order + return (0, float(page_key)) # Numeric pages first, in numerical order except ValueError: return (1, str(page_key)) # Non-numeric pages after, in alphabetical order diff --git a/fieldExtraction/tests/one_to_n_test.py b/fieldExtraction/tests/one_to_n_test.py index 56a272d..b38654a 100644 --- a/fieldExtraction/tests/one_to_n_test.py +++ b/fieldExtraction/tests/one_to_n_test.py @@ -3,8 +3,9 @@ import unittest from unittest.mock import patch, MagicMock import src.investment.one_to_n_funcs as one_to_n_funcs -from src.prompts.investment_prompts import FieldSet, Field -import src.investment.preprocess as preprocess +import src.preprocessing_funcs as preprocessing_funcs +import src.prompts.investment_prompts as investment_prompts +import src.config as config class TestOneToN(unittest.TestCase): def setUp(self): @@ -12,16 +13,49 @@ class TestOneToN(unittest.TestCase): self.filename = "test_file.pdf" @patch('src.utils.llm_utils.invoke_claude') - def test_get_exhibit_header(self, mock_invoke_claude): + @patch('src.utils.string_utils.extract_text_from_delimiters') + def test_get_exhibit_headers(self, mock_extract_text, mock_invoke_claude): + """Test exhibit header extraction with both base pages and decimal pages""" # Setup - mock_invoke_claude.return_value = "|Sample Header|" + exhibit_dict = { + "1.0": "First exhibit content", + "2.0": "Second exhibit content", + "2.1": "More second exhibit content", + "2.2": "Even more second exhibit content", + "3.0": "Third exhibit content" + } + + # Configure mocks + mock_invoke_claude.return_value = "raw_llm_response" + mock_extract_text.side_effect = [ + "Exhibit A - Services", # For page 1.0 + "Exhibit B - Payment", # For page 2.0 + "Exhibit C - Terms" # For page 3.0 + ] # Execute - result = preprocess.get_exhibit_header(self.exhibit_text, self.filename) + result = preprocessing_funcs.get_exhibit_headers(exhibit_dict, self.filename) # Assert - self.assertEqual(result, "Sample Header") - mock_invoke_claude.assert_called_once() + expected_result = { + "1.0": "Exhibit A - Services", + "2.0": "Exhibit B - Payment", + "2.1": "Exhibit B - Payment", # Should use cached result from 2.0 + "2.2": "Exhibit B - Payment", # Should use cached result from 2.0 + "3.0": "Exhibit C - Terms" + } + + self.assertEqual(result, expected_result) + + # Verify LLM was only called for base pages (3 times, not 5) + self.assertEqual(mock_invoke_claude.call_count, 3) + + # Verify prompt content + mock_invoke_claude.assert_any_call( + investment_prompts.EXHIBIT_HEADER("First exhibit content"[0:400]), + config.MODEL_ID_CLAUDE35_SONNET, + self.filename + ) @patch('src.utils.llm_utils.invoke_claude') def test_get_reimbursement_primary(self, mock_invoke_claude): diff --git a/fieldExtraction/tests/preprocessing_funcs_test.py b/fieldExtraction/tests/preprocessing_funcs_test.py index f86f177..9f63baa 100644 --- a/fieldExtraction/tests/preprocessing_funcs_test.py +++ b/fieldExtraction/tests/preprocessing_funcs_test.py @@ -131,3 +131,330 @@ class TestPreprocessingFuncs: # Verify calls made - should be number of pages minus 1 (first page), but never negative expected_calls = max(0, len(text_dict) - 1) assert mocker_invoke_claude.call_count == expected_calls + + +from src.preprocessing_funcs import get_exhibit_dict + +class TestGetExhibitDict: + + def test_case_1_short_table_first_page(self): + """Case 1: Short Table on First Page""" + text_dict = { + '5': 'Exhibit 5 header text', + '5.0': 'Short table content' + } + exhibit_chunk_mapping = { + '5': '5', + '5.0': '5' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '5.0': 'Exhibit 5 header text\nShort table content' + } + assert result == expected + + def test_case_2_continued_short_table(self): + """Case 2: Continued Short Table on Subsequent Pages""" + text_dict = { + '5': 'Exhibit 5 header text', + '5.0': 'Table start content', + '5.1': 'Table continuation content' + } + exhibit_chunk_mapping = { + '5': '5', + '5.0': '5', + '5.1': '5' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '5.0': 'Exhibit 5 header text\nTable start content', + '5.1': 'Exhibit 5 header text\nTable continuation content' + } + assert result == expected + + def test_case_3_long_table_first_page(self): + """Case 3: Long Table on First Page (split by row limit)""" + text_dict = { + '5': 'Exhibit 5 header text', + '5.0': 'Long table part 1', + '5.1': 'Long table part 2', + '5.2': 'Long table part 3' + } + exhibit_chunk_mapping = { + '5': '5', + '5.0': '5', + '5.1': '5', + '5.2': '5' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '5.0': 'Exhibit 5 header text\nLong table part 1', + '5.1': 'Exhibit 5 header text\nLong table part 2', + '5.2': 'Exhibit 5 header text\nLong table part 3' + } + assert result == expected + + def test_case_4_continued_long_table(self): + """Case 4: Continued Long Table on Subsequent Pages""" + text_dict = { + '5': 'Exhibit 5 header text', + '6': 'Page 6 text', + '6.0': 'Continued table part 1', + '6.1': 'Continued table part 2', + '7': 'Page 7 text' + } + exhibit_chunk_mapping = { + '5': '5', + '6': '5', + '6.0': '5', + '6.1': '5', + '7': '5' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '5.0': 'Exhibit 5 header text\nPage 6 text\nContinued table part 1\nPage 7 text', + '5.1': 'Exhibit 5 header text\nPage 6 text\nContinued table part 2\nPage 7 text', + '5.2': 'Exhibit 5 header text\nPage 6 text\nPage 7 text' + } + assert result == expected + + def test_case_5_multiple_tables_first_page(self): + """Case 5: Multiple Tables on First page only""" + text_dict = { + '5': 'Exhibit 5 header text', + '5.0': 'First table content', + '5.1': 'Second table content', + '5.2': 'Third table content' + } + exhibit_chunk_mapping = { + '5': '5', + '5.0': '5', + '5.1': '5', + '5.2': '5' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '5.0': 'Exhibit 5 header text\nFirst table content', + '5.1': 'Exhibit 5 header text\nSecond table content', + '5.2': 'Exhibit 5 header text\nThird table content' + } + assert result == expected + + def test_case_6_multiple_tables_with_continuation(self): + """Case 6: Multiple Tables on Subsequent Page, with first being continuation""" + text_dict = { + '5': 'Exhibit 5 header text', + '6': 'Page 6 text', + '6.0': 'Continuation of previous table', + '6.1': 'New table on same page', + '6.2': 'Another new table on same page' + } + exhibit_chunk_mapping = { + '5': '5', + '6': '5', + '6.0': '5', + '6.1': '5', + '6.2': '5' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '5.0': 'Exhibit 5 header text\nPage 6 text\nContinuation of previous table', + '5.1': 'Exhibit 5 header text\nPage 6 text\nNew table on same page', + '5.2': 'Exhibit 5 header text\nPage 6 text\nAnother new table on same page' + } + assert result == expected + + def test_case_7_multiple_tables_across_pages(self): + """Case 7: Multiple Tables on First and Subsequent Pages""" + text_dict = { + '5': 'Exhibit 5 header text', + '5.0': 'First table on page 5', + '5.1': 'Second table on page 5', + '6': 'Page 6 text', + '6.0': 'First table on page 6', + '6.1': 'Second table on page 6', + '7': 'Page 7 text' + } + exhibit_chunk_mapping = { + '5': '5', + '5.0': '5', + '5.1': '5', + '6': '5', + '6.0': '5', + '6.1': '5', + '7': '5' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '5.0': 'Exhibit 5 header text\nFirst table on page 5\nPage 6 text\nPage 7 text', + '5.1': 'Exhibit 5 header text\nSecond table on page 5\nPage 6 text\nPage 7 text', + '5.2': 'Exhibit 5 header text\nPage 6 text\nFirst table on page 6\nPage 7 text', + '5.3': 'Exhibit 5 header text\nPage 6 text\nSecond table on page 6\nPage 7 text', + '5.4': 'Exhibit 5 header text\nPage 6 text\nPage 7 text' + } + assert result == expected + + def test_edge_case_exhibit_is_decimal_page(self): + """Edge Case: Exhibit itself is a decimal page""" + text_dict = { + '5.1': 'Exhibit starting at decimal page' + } + exhibit_chunk_mapping = { + '5.1': '5.1' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '5.1': 'Exhibit starting at decimal page' + } + assert result == expected + + def test_edge_case_no_decimal_pages(self): + """Edge Case: No decimal pages in exhibit (pure text)""" + text_dict = { + '5': 'Exhibit 5 header text', + '6': 'Page 6 text', + '7': 'Page 7 text' + } + exhibit_chunk_mapping = { + '5': '5', + '6': '5', + '7': '5' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '5.0': 'Exhibit 5 header text\nPage 6 text\nPage 7 text' + } + assert result == expected + + def test_example_from_docstring_1(self): + """Test first example from docstring""" + text_dict = { + '1.0': 'Page 1.0 text', + '2': 'Page 2 text', + '3.0': 'Page 3.0 text', + '3.1': 'Page 3.1 text', + '3.2': 'Page 3.2 text', + '4.0': 'Page 4.0 text' + } + exhibit_chunk_mapping = { + '1.0': '1.0', + '2': '2', + '3.0': '2', + '3.1': '2', + '3.2': '2', + '4.0': '2' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '1.0': 'Page 1.0 text', + '2.0': 'Page 2 text\nPage 3.0 text', + '2.1': 'Page 2 text\nPage 3.1 text', + '2.2': 'Page 2 text\nPage 3.2 text', + '2.3': 'Page 2 text\nPage 4.0 text' + } + assert result == expected + + def test_example_from_docstring_2(self): + """Test second example from docstring""" + text_dict = { + '1.0': 'Page 1.0 text', + '2': 'Page 2 text', + '3.0': 'Page 3.0 text', + '3.1': 'Page 3.1 text', + '3.2': 'Page 3.2 text', + '4': 'Page 4 text', + '5': 'Page 5 text' + } + exhibit_chunk_mapping = { + '1.0': '1.0', + '2': '2', + '3.0': '2', + '3.1': '2', + '3.2': '2', + '4': '2', + '5': '2' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + '1.0': 'Page 1.0 text', + '2.0': 'Page 2 text\nPage 3.0 text\nPage 4 text\nPage 5 text', + '2.1': 'Page 2 text\nPage 3.1 text\nPage 4 text\nPage 5 text', + '2.2': 'Page 2 text\nPage 3.2 text\nPage 4 text\nPage 5 text', + '2.3': 'Page 2 text\nPage 4 text\nPage 5 text' + } + assert result == expected + + def test_exhibit_with_decimal_pages(self): + """Test exhibit with decimal pages""" + text_dict = { + "23.0": "Exhibit page 23 header", + "24": "Page 24 text" + } + exhibit_chunk_mapping = { + "23.0": "23.0", + "24": "23.0" + } + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + "23.0.0": "Exhibit page 23 header\nPage 24 text", + "23.0.1": "Page 24 text" + } + assert result == expected + + def test_multiple_exhibits_with_tables(self): + """Test multiple different exhibits, each with their own base pages and tables""" + text_dict = { + '10': 'Exhibit A header', + '10.0': 'Exhibit A table 1', + '10.1': 'Exhibit A table 2', + '11': 'Exhibit A page 2', + + '20': 'Exhibit B header', + '21': 'Exhibit B page 2', + '21.0': 'Exhibit B table 1', + '22': 'Exhibit B page 3', + + '30': 'Exhibit C header', + '31': 'Exhibit C page 2' + } + exhibit_chunk_mapping = { + '10': 'A', + '10.0': 'A', + '10.1': 'A', + '11': 'A', + + '20': 'B', + '21': 'B', + '21.0': 'B', + '22': 'B', + + '30': 'C', + '31': 'C' + } + + result = get_exhibit_dict(text_dict, exhibit_chunk_mapping) + expected = { + # Exhibit A + 'A.0': 'Exhibit A header\nExhibit A table 1\nExhibit A page 2', + 'A.1': 'Exhibit A header\nExhibit A table 2\nExhibit A page 2', + 'A.2': 'Exhibit A header\nExhibit A page 2', + + # Exhibit B + 'B.0': 'Exhibit B header\nExhibit B page 2\nExhibit B table 1\nExhibit B page 3', + 'B.1': 'Exhibit B header\nExhibit B page 2\nExhibit B page 3', + + # Exhibit C (no decimal pages) + 'C.0': 'Exhibit C header\nExhibit C page 2' + } + assert result == expected \ No newline at end of file