From 786bab611821bb489ae69c5f195957cbb323e952 Mon Sep 17 00:00:00 2001 From: Karan Desai Date: Thu, 5 Mar 2026 19:57:18 +0000 Subject: [PATCH] Merged in feature/exhibit-smart-chunking (pull request #883) Feature/exhibit smart chunking * exhibit processing per page * header dict deduplication * dedup prompt refinment * refinment for header extraction proecess * Merge remote-tracking branch 'origin/DEV' into feature/exhibit-smart-chunking * merge updates * minor fix * prompt fix for reimb type * COB defenition for clear understanding * black formatting * remove quit statement * pipiline test * pipeline test * black formatting * Merge remote-tracking branch 'origin/DEV' into feature/exhibit-smart-chunking * black formating * Merge remote-tracking branch 'origin/DEV' into feature/exhibit-smart-chunking * typo * PR comment fixes * exhibit funcs refactored * black formatting * Refactor exhibit chunking config into dedicated class Created ExhibitChunkingConfig class to centralize exhibit smart chunking configuration parameters (DEFAULT_SUBCHUNK_SIZE, MIN_PARENT_CHUNK_SIZE, CHUNK_RELEVANCE_THRESHOLD). This improves code organization by consolidating related constants and makes configuration more maintainable. Changes: - Created ExhibitChunkingConfig class with ESC_CONFIG instance - Moved CHUNK_RELEVANCE_THRESHOLD from config.py to ExhibitChunkingConfig - Updated all constant references to use ESC_CONFIG prefix - Added missing EXHIBIT_HEADER_MARKERS parameter documentation * Fix logging levels and refactor imports for exhibit chunking - Upgrade logging from WARNING to ERROR for embedding and semantic search failures - Remove unused constant imports from exhibit_funcs.py - Update test imports to use ESC_CONFIG pattern for configuration constants - Add warning when no exhibit headers found during deduplication - Expand mypy type checking by removing s3_utilities from exclude list * Merged DEV into feature/exhibit-smart-chunking * Merged DEV into feature/exhibit-smart-chunking Approved-by: Siddhant Medar --- src/config.py | 1 - src/constants/delimiters.py | 5 + src/constants/regex_patterns.py | 26 + src/pipelines/saas/file_processing.py | 245 +++--- src/pipelines/saas/prompts/prompt_calls.py | 82 +- .../shared/extraction/exhibit_funcs.py | 566 ++++++++++++- src/pipelines/shared/extraction/page_funcs.py | 5 +- .../exhibit_smart_chunking_funcs.py | 795 ++++++++++++++++++ .../shared/preprocessing/preprocess.py | 191 +++-- .../preprocessing/preprocessing_funcs.py | 210 ++++- src/prompts/prompt_templates.py | 117 ++- src/tests/test_exhibit_funcs.py | 636 +++++++++++++- src/tests/test_prompt_calls.py | 12 +- src/utils/rag_utils.py | 9 + src/utils/string_utils.py | 18 + 15 files changed, 2707 insertions(+), 211 deletions(-) create mode 100644 src/pipelines/shared/preprocessing/exhibit_smart_chunking_funcs.py diff --git a/src/config.py b/src/config.py index 2fdb804..235b268 100644 --- a/src/config.py +++ b/src/config.py @@ -323,7 +323,6 @@ def resolve_model_id(model_identifier: str) -> str: ENABLE_ANSWER_TRACKING = False ANSWER_TRACKING_DICT = {} - ######################################## POSTPROCESSING SETTINGS ######################################## FUZZY_MATCH_THRESHOLD = 0.8 FILE_NAME_COLUMN = "Contract Name" diff --git a/src/constants/delimiters.py b/src/constants/delimiters.py index e2948ef..f95270d 100644 --- a/src/constants/delimiters.py +++ b/src/constants/delimiters.py @@ -11,3 +11,8 @@ class Delimiter(Enum): PIPE = "|" # Deprecated: Use JSON format instead BACKTICK = "`" TRIPLE_BACKTICK = "```" + + +# Text parsing markers inserted by Textract table extraction +TABLE_START_MARKER = "-------Table Start--------" +TABLE_END_MARKER = "-------Table End--------" diff --git a/src/constants/regex_patterns.py b/src/constants/regex_patterns.py index a3082d3..ccb816a 100644 --- a/src/constants/regex_patterns.py +++ b/src/constants/regex_patterns.py @@ -76,3 +76,29 @@ OCR_SUBSTITUTIONS = { # dba patterns DBA_PATTERNS = [r"\bD/B/A\b", r"\bDBA\b", r"\bDOING BUSINESS AS\b", r"\bD B A\b"] + +# Exhibit Smart Chunking — Section Header Patterns +# Used to identify section/exhibit header boundaries when splitting exhibit text into chunks. +# Matches numbered sections, lettered sections, roman numerals, ARTICLE/SECTION/EXHIBIT keywords, +# legal preambles, and attachment/addendum/schedule markers. +SECTION_HEADER_PATTERNS = [ + # Numbered sections: 1.1, 4.2.1, 1.1.1.1 (at start of line or after newline) + r"(?:^|\n)\s*(\d+(?:\.\d+)*\.?)\s+", + # Single letters with period: A., B., a., b. (uppercase or lowercase) + r"(?:^|\n)\s*([A-Za-z]\.)\s+", + # Roman numerals: I., II., III., IV., etc. + r"(?:^|\n)\s*([IVXivx]+\.)\s+", + # ARTICLE headers: ARTICLE ONE, ARTICLE 1, ARTICLE I + r"(?:^|\n)\s*(ARTICLE\s+(?:ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN|\d+|[IVX]+)(?:\s*[-–—]\s*[A-Z][A-Z\s]+)?)", + # SECTION headers: SECTION 1, Section 1.1 + r"(?:^|\n)\s*(SECTION\s+\d+(?:\.\d+)*)", + # Legal preambles + r"(?:^|\n)\s*(WHEREAS,?)", + r"(?:^|\n)\s*(NOW,?\s*THEREFORE,?)", + r"(?:^|\n)\s*(IN\s+WITNESS\s+WHEREOF,?)", + # EXHIBIT/ATTACHMENT headers within text + r"(?:^|\n)\s*(EXHIBIT\s+[A-Z0-9]+)", + r"(?:^|\n)\s*(ATTACHMENT\s+[A-Z0-9]+)", + r"(?:^|\n)\s*(ADDENDUM\s+[A-Z0-9]+)", + r"(?:^|\n)\s*(SCHEDULE\s+[A-Z0-9]+)", +] diff --git a/src/pipelines/saas/file_processing.py b/src/pipelines/saas/file_processing.py index eb1e276..733da6c 100644 --- a/src/pipelines/saas/file_processing.py +++ b/src/pipelines/saas/file_processing.py @@ -1,6 +1,6 @@ import concurrent.futures import logging -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Dict, Optional import pandas as pd import src.codes.code_funcs as code_funcs @@ -9,6 +9,7 @@ from src.pipelines.shared.preprocessing import ( preprocess, hybrid_smart_chunking_funcs, ) +from src.pipelines.shared.preprocessing.exhibit_smart_chunking_funcs import ESC_CONFIG from src.pipelines.shared.postprocessing import ( aarete_derived, postprocess, @@ -22,7 +23,7 @@ from src.pipelines.shared.extraction import ( tin_npi_funcs, exhibit_funcs, ) -from src.pipelines.shared.extraction.exhibit_funcs import Exhibit +from src.pipelines.shared.extraction.exhibit_funcs import Exhibit, ExhibitChunk from src.utils import io_utils, logging_utils, string_utils, timing_utils from src.constants.constants import Constants from src import config @@ -53,8 +54,7 @@ def process_file(file_object, constants: Constants, run_timestamp): with timing_utils.timed_block("preprocess", context=filename): contract_text = preprocess.clean_text(contract_text) text_dict, top_sheet_dict = preprocess.split_text(contract_text) - text_dict, header, footer = preprocess.find_headers_and_footers(text_dict) - + text_dict, removal_metadata = preprocess.clean_header_footer(text_dict) # ONE TO N PROCESSING one_to_n_results = pd.DataFrame() # Initialize empty DataFrame for one_to_n_results if process_one_to_n: @@ -69,21 +69,19 @@ def process_file(file_object, constants: Constants, run_timestamp): with timing_utils.timed_block("one_to_n_exhibit_chunking", context=filename): # Use pages_dict for exhibit chunking (preferred) or fall back to text_dict - exhibit_chunk_mapping, all_exhibit_headers = ( - preprocess.one_to_n_exhibit_chunking( - pages_dict=pages_dict, - text_dict=text_dict, - EXHIBIT_HEADER_MARKERS=constants.EXHIBIT_HEADER_MARKERS, - filename=filename, - ) + exhibit_header_dict = preprocess.one_to_n_exhibit_chunking( + pages_dict=pages_dict, + text_dict=text_dict, + EXHIBIT_HEADER_MARKERS=constants.EXHIBIT_HEADER_MARKERS, + filename=filename, ) - logging.debug(f"{datetime_str()} Preprocessing Complete - {filename}") + logging.info(f"{datetime_str()} Preprocessing Complete - {filename}") one_to_n_results = pd.DataFrame([{"FILE_NAME": filename}]) # Initialize here if string_utils.contains_reimbursement(contract_text): - logging.debug( - f"{datetime_str()} Starting One-to-N extraction for {len(exhibit_chunk_mapping)} exhibits - {filename}" + logging.info( + f"{datetime_str()} Starting One-to-N extraction for {sum(len(value) for value in exhibit_header_dict.values() if isinstance(value, list))} exhibits - {filename}" ) with timing_utils.timed_block("one_to_n_extraction", context=filename): # Get specific fields to extract (if configured) @@ -95,8 +93,7 @@ def process_file(file_object, constants: Constants, run_timestamp): ) = run_one_to_n_prompts( pages_dict=pages_dict, text_dict=text_dict, - exhibit_chunk_mapping=exhibit_chunk_mapping, - all_exhibit_headers=all_exhibit_headers, + exhibit_header_dict=exhibit_header_dict, constants=constants, filename=filename, specific_fields=specific_fields, @@ -301,11 +298,16 @@ def process_page_reimbursements( constants: Optional[Constants] = None, filename: Optional[str] = None, specific_fields: Optional[list] = None, + relevant_chunks: Optional[list[ExhibitChunk]] = None, ): """ STEP 1 HELPER: Extract reimbursements from a single page (without exhibit-level fields). Runs: reimbursement_level → carveout_and_special_case → lesser_of_distribution → breakout + Each relevant chunk on this page is processed individually through reimbursement_level, + while downstream steps (lesser_of, etc.) receive page-based simplified exhibit context. + Pages without relevant chunks are skipped (returns [], []). + Args: exhibit: Exhibit object containing the page page_num: Page identifier (can be "27" or "27.1" for sub-pages) @@ -317,66 +319,40 @@ def process_page_reimbursements( specific_fields: Optional list of specific field names to extract. If provided with exhibit-level-only fields (like claim_type), skip detailed reimbursement extraction. + relevant_chunks: List of relevant chunks on this page. Each chunk is processed + individually through reimbursement_level. If empty/None, page is skipped. """ - # Get page text - prefer using exhibit.get_page_text() if available - if exhibit.pages is not None: - page_text = exhibit.get_page_text(page_num) - exhibit_page_nums = exhibit.exhibit_page_nums - elif pages_dict is not None: - # Handle sub-pages: "27.1" means get sub-page "1" from page "27" - if "." in page_num and not page_num.startswith("."): - base_page, sub_page = page_num.rsplit(".", 1) - if base_page in pages_dict: - page_text = pages_dict[base_page].get_text(sub_page) - else: - page_text = "" - elif page_num in pages_dict: - page_text = pages_dict[page_num].get_text() - else: - page_text = "" - exhibit_page_nums = ( - exhibit.exhibit_page_nums - if exhibit_page_nums is None - else exhibit_page_nums - ) - elif text_dict is not None: - page_text = text_dict.get(page_num, "") - exhibit_page_nums = ( - exhibit_page_nums - if exhibit_page_nums is not None - else exhibit.exhibit_page_nums - ) - else: - raise ValueError( - "Either pages_dict, text_dict, or exhibit.pages must be provided" - ) - # Simplify exhibit text - prefer using exhibit object + # No relevant chunks on this page — skip reimbursement extraction + if not relevant_chunks: + logging.debug( + f"{datetime_str()} Page {page_num}: No relevant chunks, skipping - {filename}" + ) + return [], [] + + ############################### Reimbursement Primary ############################### + # Run reimbursement_level per-chunk: each relevant chunk is processed individually + # so the LLM sees focused, relevant text one chunk at a time + reimbursement_level_answers = [] + for chunk in relevant_chunks: + chunk_answers = one_to_n_funcs.reimbursement_level( + chunk.text, constants, filename + ) + if chunk_answers: + reimbursement_level_answers.extend(chunk_answers) + + if not reimbursement_level_answers: + logging.debug( + f"{datetime_str()} Page {page_num}: No reimbursement answers found, skipping - {filename}" + ) + return [], [] + + # Page-level simplification of exhibit for downstream steps exhibit_text_simplified = preprocessing_funcs.simplify_exhibit( exhibit=exhibit, current_page_num=page_num, ) - ############################### Reimbursement Primary ############################### - reimbursement_level_answers = one_to_n_funcs.reimbursement_level( - page_text, constants, filename - ) - - if not reimbursement_level_answers: - return [], [] - - # If specific_fields is set and doesn't include reimbursement-level fields, - # skip detailed extraction (carveouts, lesser_of, breakout) - just return - # minimal info to trigger exhibit_level processing - if specific_fields is not None: - # Check if any reimbursement-level fields are requested - reimbursement_fields = {"SERVICE_TERM", "REIMB_TERM", "REIMB_PAGE"} - if not any(f in reimbursement_fields for f in specific_fields): - # Just return minimal reimbursement info to indicate this page has reimbursements - for answer_dict in reimbursement_level_answers: - answer_dict["REIMB_PAGE"] = page_num - return reimbursement_level_answers, [] - ################################ Carveouts and Special Case ################################ reimbursement_level_answers, special_case_answers = ( one_to_n_funcs.carveout_and_special_case( @@ -392,7 +368,7 @@ def process_page_reimbursements( page_num, constants, filename, - exhibit, # ← Added this parameter + exhibit, ) ################################ Get Breakouts ############################### @@ -413,8 +389,7 @@ def process_page_reimbursements( def run_one_to_n_prompts( pages_dict: Optional[dict[str, "Page"]] = None, text_dict: Optional[dict[str, str]] = None, - exhibit_chunk_mapping: Optional[dict[str, list[str]]] = None, - all_exhibit_headers: Optional[dict[str, str]] = None, + exhibit_header_dict: Optional[dict[str, list[str]]] = None, constants: Optional[Constants] = None, filename: Optional[str] = None, specific_fields: Optional[list] = None, @@ -428,76 +403,107 @@ def run_one_to_n_prompts( Args: pages_dict: Dictionary mapping page numbers to Page objects (preferred) text_dict: Dictionary mapping page numbers to text strings (for backward compatibility) - exhibit_chunk_mapping: Dictionary mapping exhibit pages to lists of page numbers - all_exhibit_headers: Dictionary mapping exhibit pages to their headers + exhibit_header_dict: Dictionary mapping page numbers to lists of headers starting on that page + Example: {'2': ['ARTICLE ONE', 'ARTICLE TWO'], '4': ['ARTICLE THREE']} constants: Constants object filename: Name of the file being processed specific_fields: Optional list of specific field names to extract. If provided, only these fields will be extracted at exhibit level. """ - if exhibit_chunk_mapping is None: - exhibit_chunk_mapping = {} - if all_exhibit_headers is None: - all_exhibit_headers = {} + if exhibit_header_dict is None: + exhibit_header_dict = {} - total_exhibits = len(exhibit_chunk_mapping) + total_exhibits = sum(len(headers) for headers in exhibit_header_dict.values()) logging.debug( f"{datetime_str()} Processing {total_exhibits} exhibits with NEW 3-step approach - {filename}" ) # Create Exhibit objects in order (maintains exhibit chain via prev_exhibit) - exhibits = exhibit_funcs.get_exhibit_list( + # Uses header-based text splitting for precise exhibit boundaries + exhibits = exhibit_funcs.create_exhibits_from_header_dict( + exhibit_header_dict=exhibit_header_dict, pages_dict=pages_dict, text_dict=text_dict, - exhibit_chunk_mapping=exhibit_chunk_mapping, - all_exhibit_headers=all_exhibit_headers, - ) # <- Returns list[Exhibit] + ) + logging.info( + f"{datetime_str()} Created {len(exhibits)} exhibit(s) from header dict - {filename}" + ) # Process exhibits serially; within each exhibit, process pages in parallel - total_pages = sum(len(exhibit.exhibit_page_nums) for exhibit in exhibits) - completed_pages = 0 + total_chunks_processed = 0 one_to_n_results = [] first_reimbursement_page = "1" for exhibit in exhibits: - exhibit_pages = exhibit.exhibit_page_nums - if exhibit_pages: - with concurrent.futures.ThreadPoolExecutor( - max_workers=min(len(exhibit_pages), 20) - ) as executor: - page_futures = { - executor.submit( - process_page_reimbursements, - exhibit, - page_num, - pages_dict, - text_dict, - exhibit_pages, - constants, - filename, - specific_fields, - ): page_num - for page_num in exhibit_pages - } + # Search for reimbursement-related chunks, then group by page + relevant_chunks = exhibit.get_relevant_chunks( + threshold=ESC_CONFIG.CHUNK_RELEVANCE_THRESHOLD + ) + page_chunks_map = ( + exhibit.group_chunks_by_page(relevant_chunks) if relevant_chunks else {} + ) - for future in concurrent.futures.as_completed(page_futures): - try: - page_num = page_futures[future] - reimbursement_rows, special_case_rows = future.result() - exhibit.add_reimbursement_rows( - reimbursement_rows, special_case_rows + logging.debug( + f"{datetime_str()} Processing exhibit {exhibits.index(exhibit) + 1}/{len(exhibits)}: " + f"page={exhibit.exhibit_page}, header='{exhibit.exhibit_header}' - {filename}" + ) + + pages_to_process = exhibit.exhibit_page_nums + + with concurrent.futures.ThreadPoolExecutor( + max_workers=min(len(pages_to_process), 20) + ) as executor: + page_futures = { + executor.submit( + process_page_reimbursements, + exhibit, + page_num, + pages_dict, + text_dict, + exhibit.exhibit_page_nums, + constants, + filename, + specific_fields, + relevant_chunks=page_chunks_map.get( + page_num + ), # None for pages without relevant chunks + ): page_num + for page_num in pages_to_process + } + + total_pages_in_exhibit = len(pages_to_process) + completed_pages = 0 + for future in concurrent.futures.as_completed(page_futures): + try: + page_num = page_futures[future] + reimbursement_rows, special_case_rows = future.result() + exhibit.add_reimbursement_rows( + reimbursement_rows, special_case_rows + ) + completed_pages += 1 + total_chunks_processed += 1 + if ( + completed_pages % 20 == 0 + or completed_pages == total_pages_in_exhibit + ): + logging.debug( + f"{datetime_str()} Page progress for exhibit {exhibit.exhibit_page}: {completed_pages}/{total_pages_in_exhibit} - {filename}" ) - completed_pages += 1 - if completed_pages % 20 == 0 or completed_pages == total_pages: - logging.debug( - f"{datetime_str()} Page progress: {completed_pages}/{total_pages} - {filename}" - ) - except Exception as e: - logging.error(f"Error processing page: {str(e)}") + except Exception as e: + logging.error(f"Error processing page {page_num}: {str(e)}") if not exhibit.has_reimbursements: + logging.debug( + f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: No reimbursements found, skipping steps 2-3 - {filename}" + ) continue + logging.debug( + f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: " + f"{len(exhibit.reimbursement_rows)} reimbursement row(s), " + f"{len(exhibit.special_case_rows)} special case row(s) - {filename}" + ) + # STEP 2: exhibit_level for this exhibit exhibit_level_answers, dynamic_reimbursement_fields = ( one_to_n_funcs.exhibit_level( @@ -533,6 +539,7 @@ def run_one_to_n_prompts( minimal_rows_with_crosswalk = aarete_derived.get_crosswalk_fields( [minimal_row], constants ) + exhibit.final_rows = minimal_rows_with_crosswalk one_to_n_results.extend(minimal_rows_with_crosswalk) else: # STEP 3: dynamic assignment & combine for this exhibit @@ -578,6 +585,10 @@ def run_one_to_n_prompts( exhibit.final_rows = combined_rows one_to_n_results.extend(combined_rows) + logging.debug( + f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: " + f"{len(combined_rows)} final row(s) after combine & clean - {filename}" + ) # Track first reimbursement page if first_reimbursement_page == "1" and exhibit.has_reimbursements: diff --git a/src/pipelines/saas/prompts/prompt_calls.py b/src/pipelines/saas/prompts/prompt_calls.py index eb06b3e..5b69311 100644 --- a/src/pipelines/saas/prompts/prompt_calls.py +++ b/src/pipelines/saas/prompts/prompt_calls.py @@ -5,8 +5,10 @@ import src.config as config import src.prompts.prompt_templates as prompt_templates from src.constants.constants import Constants from src.constants.investment_columns import FIELD_FORMAT_MAPPING +from src.constants.delimiters import Delimiter from src.prompts.fieldset import Field, FieldSet from src.utils import llm_utils, string_utils +from src.pipelines.shared.preprocessing import preprocessing_funcs from src.utils.formatting_utils import normalize_field_value from src.pipelines.shared.extraction import tin_npi_funcs import json @@ -555,7 +557,7 @@ def prompt_exhibit_linkage(page_header, previous_header, filename): return exhibit_is_different -def prompt_exhibit_header(page_content, filename): +def prompt_exhibit_header(page_content, EXHIBIT_HEADER_MARKERS, filename): """Extract exhibit header identifier from page content using pattern matching. Uses LLM to identify exhibit headers in page content based on predefined markers @@ -564,12 +566,18 @@ def prompt_exhibit_header(page_content, filename): Args: page_content (str): First 400 characters of page content to analyze. + EXHIBIT_HEADER_MARKERS (list): List of exhibit header markers to look for. filename (str): Source filename for logging and LLM attribution. Returns: - str: Extracted exhibit header identifier or marker indicating no header found. + dict: Dictionary containing exhibit header information with keys like 'exhibit_header', 'exhibit_title', etc. + Returns empty dict if parsing fails. """ - prompt, _parser = prompt_templates.EXHIBIT_HEADER(page_content[0:400]) + section_boundaries = preprocessing_funcs.identify_section_boundaries(page_content) + prompt = prompt_templates.EXHIBIT_HEADER_NEW( + section_boundaries, EXHIBIT_HEADER_MARKERS + ) + llm_answer_raw = llm_utils.invoke_claude( prompt, "sonnet_latest", @@ -579,11 +587,69 @@ def prompt_exhibit_header(page_content, filename): instruction=prompt_templates.EXHIBIT_HEADER_INSTRUCTION(), usage_label="EXHIBIT_HEADER", ) - llm_answer_extracted = _parser(llm_answer_raw) - # Parser returns a list, extract first element - if isinstance(llm_answer_extracted, list) and len(llm_answer_extracted) > 0: - return llm_answer_extracted[0] - return llm_answer_extracted + + try: + exhibit_header_dict = string_utils.universal_json_load(llm_answer_raw) + return exhibit_header_dict + except ValueError as e: + logging.error(f"Error parsing exhibit header response: {e}") + logging.error(f"Raw LLM output: {llm_answer_raw}") + return {} + + +def prompt_header_deduplication( + exhibit_header_dict: dict[str, list[str]], filename: str +) -> dict[str, list[str]]: + """Deduplicate exhibit headers across pages using LLM to handle fuzzy matches. + + Sends the full header dictionary to an LLM which identifies headers that belong + to the same section (handling case variations, appended sub-lines, etc.) and + returns only the first occurrence of each unique section. + + Args: + exhibit_header_dict: Dictionary mapping page numbers to lists of headers. + filename: Source filename for logging and LLM attribution. + + Returns: + dict[str, list[str]]: Deduplicated dictionary with only the first page + of each unique section. Falls back to original dict on parse failure. + """ + if not exhibit_header_dict: + logging.warning(f"No exhibit headers found in {filename} for deduplication.") + return exhibit_header_dict + + header_dict_str = json.dumps(exhibit_header_dict, indent=2) + + prompt = prompt_templates.EXHIBIT_HEADER_DEDUP(header_dict_str) + + llm_answer_raw = llm_utils.invoke_claude( + prompt, + "sonnet_latest", + filename, + max_tokens=2048, + cache=True, + instruction=prompt_templates.EXHIBIT_HEADER_DEDUP_INSTRUCTION(), + usage_label="EXHIBIT_HEADER_DEDUP", + ) + + try: + llm_answer_final = string_utils.extract_text_from_delimiters( + llm_answer_raw, Delimiter.PIPE + ) + deduped_dict = json.loads(llm_answer_final) + + if not isinstance(deduped_dict, dict) or not deduped_dict: + logging.warning( + f"EXHIBIT_HEADER_DEDUP returned invalid structure, using original dict" + ) + return exhibit_header_dict + + return deduped_dict + + except (json.JSONDecodeError, ValueError) as e: + logging.error(f"Failed to parse EXHIBIT_HEADER_DEDUP response: {e}") + logging.error(f"Raw output: {llm_answer_raw}") + return exhibit_header_dict def prompt_date_fix(date: str) -> str: diff --git a/src/pipelines/shared/extraction/exhibit_funcs.py b/src/pipelines/shared/extraction/exhibit_funcs.py index 24a95c0..2470408 100644 --- a/src/pipelines/shared/extraction/exhibit_funcs.py +++ b/src/pipelines/shared/extraction/exhibit_funcs.py @@ -4,11 +4,32 @@ Exhibit class for tracking exhibit state during one-to-n processing. This module provides the Exhibit class which maintains all state and field values for a single exhibit during field extraction, including a linked list structure for propagating dynamic fields from exhibit n-1 to exhibit n. + +Chunking infrastructure (SubChunk, ExhibitChunk, chunking functions) lives in: + src/pipelines/shared/preprocessing/exhibit_smart_chunking_funcs.py """ -from typing import TYPE_CHECKING, Optional +import re +from typing import TYPE_CHECKING, Optional, List from src.prompts.fieldset import FieldSet +from src.constants.delimiters import TABLE_START_MARKER, TABLE_END_MARKER +from src.utils.rag_utils import get_embeddings, HSC_CONFIG, cosine_similarity +from src.pipelines.shared.preprocessing.exhibit_smart_chunking_funcs import ( + SubChunk, + ExhibitChunk, + REIMBURSEMENT_QUERIES, + REIMBURSEMENT_KEYWORDS, + _split_text_at_sentence_boundaries, + _find_table_regions, + _chunk_non_table_text, + _chunk_exhibit_text, + _merge_header_only_chunks, + _is_section_header_chunk, + _merge_small_parent_chunks, + _compute_chunk_embeddings, + search_reimbursement_chunks as _search_reimbursement_chunks, +) if TYPE_CHECKING: from src.pipelines.shared.extraction.page_funcs import Page @@ -94,6 +115,13 @@ class Exhibit: # NEW: Lesser-of answer tracking self.lesser_of_answers = [] # List[dict] - answer dicts from LESSER_OF_CHECK + # Chunks: lazy-loaded sections/paragraphs within this exhibit + self._chunks: Optional[list[ExhibitChunk]] = None + + # Page boundary tracking: maps each page's character range within exhibit_text + # Each tuple is (page_num, start_pos, end_pos) + self._page_char_ranges: Optional[list[tuple[str, int, int]]] = None + @property def pages(self) -> Optional[dict[str, "Page"]]: """ @@ -109,14 +137,20 @@ class Exhibit: """ Get the full exhibit text. - Computed from pages if available, otherwise returns stored exhibit_text. - This property provides backward compatibility. + Priority: + 1. Returns stored _exhibit_text if explicitly set (e.g., by create_exhibits_from_header_dict + which computes text with proper header boundaries) + 2. Computes from pages if available + 3. Returns empty string if neither available Returns: Full text content of the exhibit """ - if self._pages is not None: - # Compute from pages + # Prioritize explicitly-set exhibit_text (has correct boundaries) + if self._exhibit_text is not None: + return self._exhibit_text + elif self._pages is not None: + # Compute from pages (fallback when _exhibit_text not set) text_parts = [] for page_id in self.exhibit_page_nums: # Handle sub-pages: "27.1" means get sub-page "1" from page "27" @@ -129,11 +163,264 @@ class Exhibit: page_text = self._pages[page_id].get_text() text_parts.append(page_text) return "\n".join(text_parts) - elif self._exhibit_text is not None: - return self._exhibit_text else: return "" + @property + def chunks(self) -> list[ExhibitChunk]: + """ + Get the chunks (sections/paragraphs) within this exhibit. + + Chunks are computed on first access by parsing the exhibit text + for section headers like "1.1", "A.", "ARTICLE ONE", etc. + Embeddings are computed eagerly for each chunk. + + Returns: + List of ExhibitChunk objects representing logical sections with embeddings + """ + if self._chunks is None: + self._chunks = _chunk_exhibit_text(self.exhibit_text) + # Compute embeddings eagerly + self._chunks = _compute_chunk_embeddings(self._chunks) + return self._chunks + + @property + def page_char_ranges(self) -> list[tuple[str, int, int]]: + """ + Get character position ranges for each page within exhibit_text. + + Returns: + List of (page_num, start_pos, end_pos) tuples. + start_pos/end_pos are character positions in self.exhibit_text. + """ + if self._page_char_ranges is None: + self._page_char_ranges = self._compute_page_char_ranges() + return self._page_char_ranges + + def _compute_page_char_ranges(self) -> list[tuple[str, int, int]]: + """Fallback: reconstruct page boundaries from exhibit_text and page texts. + + Note: This uses get_page_text() which returns untrimmed page text, so + boundaries may be slightly off for first/last pages of trimmed exhibits. + The primary path stores exact ranges via create_exhibits_from_header_dict. + """ + if self._pages is None: + return [] + ranges = [] + pos = 0 + for page_id in self.exhibit_page_nums: + page_text = self.get_page_text(page_id) + ranges.append((page_id, pos, pos + len(page_text))) + pos += len(page_text) + 1 # +1 for "\n" join separator + return ranges + + def group_chunks_by_page( + self, chunks: list[ExhibitChunk] + ) -> dict[str, list[ExhibitChunk]]: + """ + Group chunks by the page they fall on, using character position mapping. + + A chunk is assigned to the page whose character range in exhibit_text + contains the chunk's start_pos. Chunks spanning a page boundary are + assigned to the page where they start. + + Args: + chunks: List of ExhibitChunk objects to group (e.g., relevant_chunks) + + Returns: + Dict mapping page_num -> list of chunks on that page, ordered by chunk_id. + Pages with no matching chunks are excluded. + """ + if not chunks or not self.page_char_ranges: + return {} + + page_chunks: dict[str, list[ExhibitChunk]] = {} + + for chunk in chunks: + assigned_page = None + for page_num, page_start, page_end in self.page_char_ranges: + if page_start <= chunk.start_pos < page_end: + assigned_page = page_num + break + + # Fallback: if chunk.start_pos is at a "\n" separator or beyond, + # assign to last page + if assigned_page is None and self.page_char_ranges: + assigned_page = self.page_char_ranges[-1][0] + + if assigned_page is not None: + if assigned_page not in page_chunks: + page_chunks[assigned_page] = [] + page_chunks[assigned_page].append(chunk) + + # Sort chunks within each page by chunk_id + for page_num in page_chunks: + page_chunks[page_num].sort(key=lambda c: int(c.chunk_id)) + + return page_chunks + + def get_chunk_by_id(self, chunk_id: str) -> Optional[ExhibitChunk]: + """ + Get a specific chunk by its ID. + + Args: + chunk_id: The chunk identifier (e.g., "1", "2", "3") + + Returns: + The ExhibitChunk with the given ID, or None if not found + """ + for chunk in self.chunks: + if chunk.chunk_id == chunk_id: + return chunk + return None + + def get_chunks_by_header_pattern(self, pattern: str) -> list[ExhibitChunk]: + """ + Get chunks whose section header matches a regex pattern. + + Args: + pattern: Regex pattern to match against section headers + + Returns: + List of matching ExhibitChunk objects + """ + matching = [] + regex = re.compile(pattern, re.IGNORECASE) + for chunk in self.chunks: + if regex.search(chunk.section_header): + matching.append(chunk) + return matching + + def search_chunks(self, query: str, top_k: int = 5) -> list[ExhibitChunk]: + """ + Search chunks by semantic similarity to a query. + + For chunks with sub-chunks, searches the sub-chunks and returns the parent + chunk if any sub-chunk matches. Table chunks are ALWAYS included in results. + Final results are ordered by chunk_id to maintain document order. + + Args: + query: The search query text + top_k: Number of top non-table results to include (table chunks always included) + + Returns: + List of ExhibitChunk objects ordered by chunk_id (document order). + Includes all table chunks plus top_k matching non-table chunks. + """ + if not self.chunks: + return [] + + # Separate table chunks from non-table chunks + table_chunks = [c for c in self.chunks if c.is_table] + non_table_chunks = [c for c in self.chunks if not c.is_table] + + if not non_table_chunks and not table_chunks: + logging.warning("No chunks available for search") + return [] + + matched_chunk_ids: set[str] = set() + + # Collect all searchable items: (parent_chunk_id, embedding, text) + # For chunks with sub-chunks: search sub-chunks + # For chunks without sub-chunks: search chunk directly + search_items: List[tuple[str, List[float], str]] = [] + + for chunk in non_table_chunks: + if chunk.has_sub_chunks(): + # Search sub-chunks, map back to parent + for sub_chunk in chunk.sub_chunks: + if sub_chunk.embedding is not None: + search_items.append( + (chunk.chunk_id, sub_chunk.embedding, sub_chunk.text) + ) + elif chunk.embedding is not None: + # Search chunk directly + search_items.append((chunk.chunk_id, chunk.embedding, chunk.text)) + + if search_items: + try: + # Get embedding for query + embeddings_model = get_embeddings( + HSC_CONFIG.EMBEDDING_MODEL_ID, + HSC_CONFIG.BEDROCK_REGION, + ) + query_embedding = embeddings_model.embed_query(query) + + # Compute similarities and track best score per parent chunk + chunk_best_scores: dict[str, float] = {} + for parent_id, embedding, _ in search_items: + similarity = cosine_similarity(query_embedding, embedding) + if ( + parent_id not in chunk_best_scores + or similarity > chunk_best_scores[parent_id] + ): + chunk_best_scores[parent_id] = similarity + + # Sort by best score and take top_k parent chunks + sorted_chunks = sorted( + chunk_best_scores.items(), key=lambda x: x[1], reverse=True + ) + for parent_id, _ in sorted_chunks[:top_k]: + matched_chunk_ids.add(parent_id) + + except Exception as e: + logging.error(f"Error during chunk search: {e}") + # On error, still return table chunks + + # Collect all table chunk IDs (always included) + table_chunk_ids = {c.chunk_id for c in table_chunks} + + # Combine: all table chunks + matched non-table chunks + all_selected_ids = table_chunk_ids | matched_chunk_ids + + # Get chunks and sort by chunk_id to maintain document order + result_chunks = [c for c in self.chunks if c.chunk_id in all_selected_ids] + result_chunks.sort(key=lambda c: int(c.chunk_id)) + + return result_chunks + + def get_relevant_chunks(self, threshold: float = 0.30) -> list[ExhibitChunk]: + """ + Get chunks that are relevant to reimbursements using hybrid search. + + Uses a two-stage hybrid approach: + 1. Keyword matching: Chunks containing REIMBURSEMENT_KEYWORDS are auto-included + 2. Semantic search: Additional chunks matching REIMBURSEMENT_QUERIES above threshold + + Table chunks are ALWAYS included since they typically contain reimbursement data. + + Args: + threshold: Minimum similarity score (0-1) for semantic matching. Default 0.30. + + Returns: + List of ExhibitChunk objects ordered by chunk_id (document order). + Includes: all table chunks + keyword matches + semantic matches. + """ + return self.search_reimbursement_chunks(threshold=threshold) + + def search_reimbursement_chunks( + self, + threshold: float = 0.30, + queries: Optional[List[str]] = None, + keywords: Optional[List[str]] = None, + use_semantic: bool = True, + use_keywords: bool = True, + ) -> list[ExhibitChunk]: + """ + Retrieve chunks related to reimbursements using hybrid keyword + semantic search. + + Delegates to search_reimbursement_chunks in exhibit_smart_chunking_funcs.py. + See that function for full documentation. + """ + return _search_reimbursement_chunks( + self.chunks, + threshold=threshold, + queries=queries, + keywords=keywords, + use_semantic=use_semantic, + use_keywords=use_keywords, + ) + def get_page_text(self, page_id: str) -> str: """ Get text for a specific page or sub-page. @@ -363,12 +650,266 @@ class Exhibit: def __repr__(self): """String representation for debugging.""" + chunk_count = len(self._chunks) if self._chunks is not None else "not loaded" return ( f"Exhibit(page={self.exhibit_page}, pages={len(self.exhibit_page_nums)}, " - f"has_reimb={self.has_reimbursements}, reimb_count={len(self.reimbursement_rows)})" + f"chunks={chunk_count}, has_reimb={self.has_reimbursements}, " + f"reimb_count={len(self.reimbursement_rows)})" ) +def _find_header_in_text(text: str, header: str) -> int: + """ + Find header position in text, trying multiple matching strategies. + + Args: + text: The text to search in + header: The header to find + + Returns: + Character position of header start, or -1 if not found + """ + if not header or not text: + return -1 + + # Try exact match first + pos = text.find(header) + if pos != -1: + return pos + + # Try with normalized whitespace (header may have \n that appears as space in text) + normalized_header = " ".join(header.split()) + normalized_text = " ".join(text.split()) + + # Find in normalized, then map back to original position + norm_pos = normalized_text.find(normalized_header) + if norm_pos != -1: + # Find approximate position in original text + # Count how many chars into the normalized text, find similar position + words_before = normalized_text[:norm_pos].split() + # Find the start of this word sequence in original text + search_start = 0 + for word in words_before: + found = text.find(word, search_start) + if found != -1: + search_start = found + len(word) + + # Now search for first word of header from this position + first_header_word = ( + normalized_header.split()[0] if normalized_header else header + ) + pos = text.find(first_header_word, max(0, search_start - 50)) + if pos != -1: + return pos + + # Try case-insensitive match + text_lower = text.lower() + header_lower = header.lower() + pos = text_lower.find(header_lower) + if pos != -1: + return pos + + # Try case-insensitive with normalized whitespace + normalized_header_lower = " ".join(header_lower.split()) + normalized_text_lower = " ".join(text_lower.split()) + norm_pos = normalized_text_lower.find(normalized_header_lower) + if norm_pos != -1: + # Map back to original position using first significant word + first_words = normalized_header_lower.split()[:3] # Use first 3 words + for word in first_words: + if len(word) > 3: # Skip short words like "the", "and" + pos = text_lower.find(word) + if pos != -1: + return pos + + # Last resort: find the first significant word of the header + header_words = header.split() + for word in header_words: + if len(word) > 4: # Skip short words + pos = text.find(word) + if pos != -1: + logging.debug( + f"_find_header_in_text: Found header '{header}' via word '{word}' at pos {pos}" + ) + return pos + + logging.warning( + f"_find_header_in_text: Could not find header '{header[:50]}...' in text" + ) + return -1 + + +def create_exhibits_from_header_dict( + exhibit_header_dict: dict[str, list[str]], + pages_dict: Optional[dict[str, "Page"]] = None, + text_dict: Optional[dict[str, str]] = None, +) -> list[Exhibit]: + """ + Create Exhibit objects directly from exhibit_header_dict. + + Each header in exhibit_header_dict becomes a separate Exhibit. + Exhibit text starts from its header and ends just before the next header. + Headers are searched within the page text to determine exact boundaries. + + Args: + exhibit_header_dict: Mapping of page numbers to list of headers starting on that page. + Example: {'2': ['ARTICLE ONE', 'ARTICLE TWO'], '4': ['ARTICLE THREE']} + pages_dict: Dictionary mapping page numbers to Page objects (preferred) + text_dict: Dictionary mapping page numbers to text strings (backward compatibility) + + Returns: + List of Exhibit objects ordered by appearance in document + """ + if not exhibit_header_dict: + return [] + + if pages_dict is None and text_dict is None: + raise ValueError("Either pages_dict or text_dict must be provided") + + # Get sorted list of all page numbers + source_dict = pages_dict if pages_dict is not None else text_dict + assert source_dict is not None # guaranteed by the check above + all_page_nums = sorted(source_dict.keys(), key=lambda x: int(x.split(".")[0])) + + # Helper to get page text + def get_page_text(page_num: str) -> str: + if pages_dict is not None and page_num in pages_dict: + return pages_dict[page_num].get_text() + elif text_dict is not None and page_num in text_dict: + return text_dict[page_num] + return "" + + # Flatten exhibit_header_dict into ordered list of (page_num, header) + exhibit_entries: list[tuple[str, str]] = [] + sorted_header_pages = sorted(exhibit_header_dict.keys(), key=lambda x: int(x)) + for page_num in sorted_header_pages: + headers = exhibit_header_dict[page_num] + for header in headers: + exhibit_entries.append((page_num, header)) + + if not exhibit_entries: + return [] + + # Build exhibits with proper text boundaries + exhibits: list[Exhibit] = [] + prev_exhibit: Optional[Exhibit] = None + + for i, (start_page, header) in enumerate(exhibit_entries): + # Determine the next header info (if any) + next_header = None + next_header_page = None + if i + 1 < len(exhibit_entries): + next_header_page, next_header = exhibit_entries[i + 1] + + # Determine end page for this exhibit + if next_header_page is not None: + try: + next_page_idx = all_page_nums.index(next_header_page) + # End page is the page containing the next header + # (we'll trim the text at the header position) + end_page = next_header_page + except ValueError: + end_page = all_page_nums[-1] + else: + end_page = all_page_nums[-1] + + # Collect page numbers from start_page to end_page + try: + start_idx = all_page_nums.index(start_page) + end_idx = all_page_nums.index(end_page) + except ValueError: + continue + + exhibit_page_nums = all_page_nums[start_idx : end_idx + 1] + + # Build exhibit text with proper boundaries + text_parts = [] + for j, page_num in enumerate(exhibit_page_nums): + page_text = get_page_text(page_num) + + if j == 0: + # First page: start from header position + header_pos = _find_header_in_text(page_text, header) + if header_pos != -1: + page_text = page_text[header_pos:] + # If header not found, use full page (fallback) + + if ( + page_num == end_page + and next_header is not None + and next_header_page == end_page + ): + # Last page AND next header is on this same page: trim at next header + next_header_pos = _find_header_in_text(page_text, next_header) + if next_header_pos != -1: + page_text = page_text[:next_header_pos] + + text_parts.append(page_text) + + exhibit_text = "\n".join(text_parts) + + # Adjust exhibit_page_nums to exclude the end page if exhibit text doesn't + # actually include content from that page (when next header is at start of page) + # Check if the last page contributed any text + if len(exhibit_page_nums) > 1 and next_header_page == end_page: + last_page_text = get_page_text(end_page) + next_header_pos = ( + _find_header_in_text(last_page_text, next_header) if next_header else -1 + ) + if next_header_pos == 0 or ( + next_header_pos != -1 and last_page_text[:next_header_pos].strip() == "" + ): + # Next header is at the very start of the page, exclude this page + exhibit_page_nums = exhibit_page_nums[:-1] + end_page = exhibit_page_nums[-1] if exhibit_page_nums else start_page + # Rebuild text without the last page + text_parts = [] + for page_num in exhibit_page_nums: + page_text = get_page_text(page_num) + if page_num == start_page: + header_pos = _find_header_in_text(page_text, header) + if header_pos != -1: + page_text = page_text[header_pos:] + text_parts.append(page_text) + exhibit_text = "\n".join(text_parts) + + # Compute page character ranges within exhibit_text + # (uses the final text_parts and exhibit_page_nums after any rebuild) + page_char_ranges = [] + pos = 0 + for pn, tp in zip(exhibit_page_nums, text_parts): + page_char_ranges.append((pn, pos, pos + len(tp))) + pos += len(tp) + 1 # +1 for "\n" join separator + + if pages_dict is not None: + exhibit_pages_dict = { + pn: pages_dict[pn] for pn in exhibit_page_nums if pn in pages_dict + } + exhibit = Exhibit( + exhibit_page=start_page, + exhibit_page_nums=exhibit_page_nums, + exhibit_header=header, + pages=exhibit_pages_dict, + prev_exhibit=prev_exhibit, + ) + # Override the computed text since we have custom boundaries + exhibit._exhibit_text = exhibit_text + else: + exhibit = Exhibit( + exhibit_page=start_page, + exhibit_page_nums=exhibit_page_nums, + exhibit_header=header, + exhibit_text=exhibit_text, + prev_exhibit=prev_exhibit, + ) + + exhibit._page_char_ranges = page_char_ranges + exhibits.append(exhibit) + prev_exhibit = exhibit + + return exhibits + + def get_exhibit_list( pages_dict: Optional[dict[str, "Page"]] = None, text_dict: Optional[dict[str, str]] = None, @@ -376,6 +917,12 @@ def get_exhibit_list( all_exhibit_headers: Optional[dict[str, str]] = None, ) -> list[Exhibit]: """ + DEPRECATED: Use create_exhibits_from_header_dict() instead. + + This function assigns whole pages to exhibits without precise text boundary splitting. + The new create_exhibits_from_header_dict() function splits text at header boundaries + for more accurate exhibit text extraction. + Create a list of Exhibit objects from pages or text dictionary. This function supports both the new Page-based approach and the legacy @@ -394,6 +941,9 @@ def get_exhibit_list( Either pages_dict or text_dict must be provided. If both are provided, pages_dict takes precedence. """ + logging.warning( + "get_exhibit_list is deprecated. Use create_exhibits_from_header_dict() instead." + ) if exhibit_chunk_mapping is None: exhibit_chunk_mapping = {} if all_exhibit_headers is None: diff --git a/src/pipelines/shared/extraction/page_funcs.py b/src/pipelines/shared/extraction/page_funcs.py index d052639..df5efe3 100644 --- a/src/pipelines/shared/extraction/page_funcs.py +++ b/src/pipelines/shared/extraction/page_funcs.py @@ -12,10 +12,7 @@ import re from typing import Optional from src import config - - -TABLE_START_MARKER = "-------Table Start--------" -TABLE_END_MARKER = "-------Table End--------" +from src.constants.delimiters import TABLE_START_MARKER, TABLE_END_MARKER class Page: diff --git a/src/pipelines/shared/preprocessing/exhibit_smart_chunking_funcs.py b/src/pipelines/shared/preprocessing/exhibit_smart_chunking_funcs.py new file mode 100644 index 0000000..bbd8e04 --- /dev/null +++ b/src/pipelines/shared/preprocessing/exhibit_smart_chunking_funcs.py @@ -0,0 +1,795 @@ +""" +Exhibit Smart Chunking - chunking infrastructure for 1:n reimbursement extraction. + +Follows the same pattern as hybrid_smart_chunking_funcs.py (for 1:1 fields). +Provides chunking, embedding, and hybrid retrieval functions for exhibit text analysis. +""" + +import re +import logging +from typing import Optional, List + +from src.utils.rag_utils import get_embeddings, HSC_CONFIG, cosine_similarity +from src.utils.string_utils import keyword_match +from src.constants.regex_patterns import SECTION_HEADER_PATTERNS +from src.constants.delimiters import TABLE_START_MARKER, TABLE_END_MARKER + + +class ExhibitChunkingConfig: + """Configuration for Exhibit Smart Chunking (reimbursement extraction)""" + + DEFAULT_SUBCHUNK_SIZE: int = 500 # default sub-chunk size in characters + MIN_PARENT_CHUNK_SIZE: int = 100 # minimum parent chunk size in characters + CHUNK_RELEVANCE_THRESHOLD: float = ( + 0.30 # default relevance threshold for semantic search + ) + + +ESC_CONFIG = ExhibitChunkingConfig() + +# Reimbursement-related queries for multi-query RAG retrieval +REIMBURSEMENT_QUERIES = [ + # Primary reimbursement/payment terms + "reimbursement payment rates fees maximum compensation amount percentage flat rate", + # Service descriptions + "covered services medical procedures cpt codes service category ancillary services", + # Payment methodologies and schedules + "Medicare fee schedule DRG per diem per unit per visit payment methodology", + # Dollar amounts and financial terms + "dollar($) charges costs pricing schedule rates amounts maximum percent(%) of lesser of", + # Special cases (stop loss, outliers, adjustments) + "stop loss outlier threshold discount escalator premium adjustment capped", + # Contract-specific rate language + "paid at percent (%) of billed charges allowable fee schedule responsible for payment EOP EOB", + # Provider/facility specific + "inpatient outpatient hospital physician professional facility rates", +] + +# Keywords for hybrid search - chunks containing these get automatic inclusion +REIMBURSEMENT_KEYWORDS = ["$", "%", "percentage", "lesser of"] + + +class SubChunk: + """ + Represents a sub-chunk of a parent ExhibitChunk for fine-grained embedding search. + + Sub-chunks are created by splitting large chunks (~500 chars) at sentence boundaries. + When a sub-chunk matches a search query, the parent chunk is returned. + + Attributes: + sub_chunk_id (str): Identifier like "15.1", "15.2" where 15 is parent chunk_id + parent_chunk_id (str): The parent chunk's ID (e.g., "15") + text (str): The sub-chunk text content + embedding (Optional[List[float]]): Vector embedding for this sub-chunk + """ + + def __init__( + self, + sub_chunk_id: str, + parent_chunk_id: str, + text: str, + embedding: Optional[List[float]] = None, + ): + self.sub_chunk_id = sub_chunk_id + self.parent_chunk_id = parent_chunk_id + self.text = text + self.embedding = embedding + + def __repr__(self): + text_preview = ( + self.text[:50].replace("\n", " ") + "..." + if len(self.text) > 50 + else self.text.replace("\n", " ") + ) + has_embedding = "yes" if self.embedding is not None else "no" + return f"SubChunk(id={self.sub_chunk_id}, parent={self.parent_chunk_id}, embedding={has_embedding}, text='{text_preview}')" + + +def _split_text_at_sentence_boundaries( + text: str, max_size: int = ESC_CONFIG.DEFAULT_SUBCHUNK_SIZE +) -> List[str]: + """ + Split text into chunks at sentence boundaries, targeting max_size characters. + + Args: + text: Text to split + max_size: Target maximum size for each chunk (default 500) + + Returns: + List of text chunks split at sentence boundaries + """ + if not text or len(text) <= max_size: + return [text] if text else [] + + # Sentence-ending patterns + sentence_endings = re.compile(r"(?<=[.!?])\s+") + + sentences = sentence_endings.split(text) + chunks = [] + current_chunk = "" + + for sentence in sentences: + # If adding this sentence exceeds max_size and we have content, start new chunk + if current_chunk and len(current_chunk) + len(sentence) + 1 > max_size: + chunks.append(current_chunk.strip()) + current_chunk = sentence + else: + current_chunk = ( + current_chunk + " " + sentence if current_chunk else sentence + ) + + # Add final chunk + if current_chunk.strip(): + chunks.append(current_chunk.strip()) + + return chunks + + +class ExhibitChunk: + """ + Represents a logical chunk/section within an exhibit. + + Chunks are paragraphs or subsections identified by section headers like: + - Numbered sections: "1.1", "4.2.1", "1." + - Lettered sections: "A.", "a.", "C." + - Article headers: "ARTICLE ONE", "ARTICLE 1" + - Legal preambles: "WHEREAS", "NOW, THEREFORE" + - Tables: Content between table markers + + Large chunks (>500 chars) are split into sub-chunks for fine-grained embedding search. + When a sub-chunk matches a query, the full parent chunk is returned. + + Attributes: + chunk_id (str): Unique identifier for this chunk (e.g., "1", "2", "3") + section_header (str): The section header/identifier (e.g., "1.1", "A.", "ARTICLE ONE", "TABLE") + text (str): The full text content of this chunk (including header) + start_pos (int): Starting character position in parent exhibit text + end_pos (int): Ending character position in parent exhibit text + embedding (Optional[List[float]]): Vector embedding of the chunk text (for small chunks) + is_table (bool): True if this chunk represents a table + sub_chunks (List[SubChunk]): Sub-chunks for large non-table chunks + """ + + def __init__( + self, + chunk_id: str, + section_header: str, + text: str, + start_pos: int, + end_pos: int, + embedding: Optional[List[float]] = None, + is_table: bool = False, + ): + self.chunk_id = chunk_id + self.section_header = section_header + self.text = text + self.start_pos = start_pos + self.end_pos = end_pos + self.embedding = embedding + self.is_table = is_table + self._sub_chunks: Optional[List[SubChunk]] = None + + @property + def sub_chunks(self) -> List[SubChunk]: + """ + Get sub-chunks for this chunk. Created lazily on first access. + + Table chunks and small chunks (<= 500 chars) return empty list. + Large non-table chunks are split at sentence boundaries. + """ + if self._sub_chunks is None: + self._sub_chunks = self._create_sub_chunks() + return self._sub_chunks + + def _create_sub_chunks(self) -> List[SubChunk]: + """Create sub-chunks by splitting text at sentence boundaries.""" + # Don't split table chunks or small chunks + if self.is_table or len(self.text) <= ESC_CONFIG.DEFAULT_SUBCHUNK_SIZE: + return [] + + text_parts = _split_text_at_sentence_boundaries( + self.text, ESC_CONFIG.DEFAULT_SUBCHUNK_SIZE + ) + + # If splitting resulted in only 1 part, no need for sub-chunks + if len(text_parts) <= 1: + return [] + + sub_chunks = [] + for i, part in enumerate(text_parts, start=1): + sub_chunk_id = f"{self.chunk_id}.{i}" + sub_chunks.append( + SubChunk( + sub_chunk_id=sub_chunk_id, + parent_chunk_id=self.chunk_id, + text=part, + ) + ) + + return sub_chunks + + def has_sub_chunks(self) -> bool: + """Check if this chunk has sub-chunks.""" + return len(self.sub_chunks) > 0 + + def __repr__(self): + header_preview = self.section_header if self.section_header else "(no header)" + text_preview = ( + self.text[:50].replace("\n", " ") + "..." + if len(self.text) > 50 + else self.text.replace("\n", " ") + ) + has_embedding = "yes" if self.embedding is not None else "no" + table_flag = ", TABLE" if self.is_table else "" + sub_chunk_info = ( + f", sub_chunks={len(self.sub_chunks)}" if self.has_sub_chunks() else "" + ) + return f"ExhibitChunk(id={self.chunk_id}, header='{header_preview}', embedding={has_embedding}{table_flag}{sub_chunk_info}, text='{text_preview}')" + + +def _find_table_regions(exhibit_text: str) -> list[tuple[int, int]]: + """ + Find all table regions in the text delimited by table markers. + + Args: + exhibit_text: The full text content of an exhibit + + Returns: + List of (start_pos, end_pos) tuples for each table region. + The positions include the markers themselves. + """ + table_regions = [] + search_start = 0 + + while True: + # Find next table start marker + start_marker_pos = exhibit_text.find(TABLE_START_MARKER, search_start) + if start_marker_pos == -1: + break + + # Find corresponding table end marker + end_marker_pos = exhibit_text.find(TABLE_END_MARKER, start_marker_pos) + if end_marker_pos == -1: + # No end marker found, treat rest of text as table + table_regions.append((start_marker_pos, len(exhibit_text))) + break + + # Include the end marker in the region + end_pos = end_marker_pos + len(TABLE_END_MARKER) + table_regions.append((start_marker_pos, end_pos)) + + # Continue searching after this table + search_start = end_pos + + return table_regions + + +def _chunk_non_table_text(text: str, offset: int = 0) -> list[ExhibitChunk]: + """ + Parse non-table text into logical chunks based on section headers. + + Args: + text: Text content without tables + offset: Position offset to add to chunk positions (for correct global positioning) + + Returns: + List of ExhibitChunk objects representing sections/paragraphs + """ + if not text or not text.strip(): + return [] + + # Find all section header matches with their positions + matches = [] + for pattern in SECTION_HEADER_PATTERNS: + for match in re.finditer(pattern, text, re.IGNORECASE | re.MULTILINE): + match_start = match.start() + header = match.group(1).strip() if match.group(1) else "" + matches.append((match_start, header, match.end())) + + # Sort matches by position + matches.sort(key=lambda x: x[0]) + + # Remove duplicates (same position) - keep the longest header + unique_matches = [] + for match in matches: + if not unique_matches or match[0] != unique_matches[-1][0]: + unique_matches.append(match) + elif len(match[1]) > len(unique_matches[-1][1]): + unique_matches[-1] = match + + matches = unique_matches + + # If no section headers found, return entire text as single chunk + if not matches: + stripped_text = text.strip() + if stripped_text: + return [ + ExhibitChunk( + chunk_id="", # Will be assigned later + section_header="", + text=stripped_text, + start_pos=offset, + end_pos=offset + len(text), + ) + ] + return [] + + chunks = [] + + # Handle text before first section header (if any meaningful content) + first_match_pos = matches[0][0] + if first_match_pos > 0: + preamble_text = text[:first_match_pos].strip() + if preamble_text and len(preamble_text) > 10: + chunks.append( + ExhibitChunk( + chunk_id="", + section_header="", + text=preamble_text, + start_pos=offset, + end_pos=offset + first_match_pos, + ) + ) + + # Create chunks from section headers + for i, (start_pos, header, header_end) in enumerate(matches): + if i + 1 < len(matches): + end_pos = matches[i + 1][0] + else: + end_pos = len(text) + + chunk_text = text[start_pos:end_pos].strip() + + if chunk_text: + chunks.append( + ExhibitChunk( + chunk_id="", + section_header=header, + text=chunk_text, + start_pos=offset + start_pos, + end_pos=offset + end_pos, + ) + ) + + return chunks + + +def _chunk_exhibit_text(exhibit_text: str) -> list[ExhibitChunk]: + """ + Parse exhibit text into logical chunks based on section headers and tables. + + Tables (delimited by TABLE_START_MARKER and TABLE_END_MARKER) are treated as + their own separate chunks. Non-table text is chunked based on section headers + like "1.1", "A.", "ARTICLE ONE", etc. + + Args: + exhibit_text: The full text content of an exhibit + + Returns: + List of ExhibitChunk objects representing sections/paragraphs/tables + """ + if not exhibit_text or not exhibit_text.strip(): + return [] + + # Find all table regions + table_regions = _find_table_regions(exhibit_text) + + # If no tables, use the original section header chunking + if not table_regions: + chunks = _chunk_non_table_text(exhibit_text, offset=0) + chunks = _merge_header_only_chunks(chunks) + chunks = _merge_small_parent_chunks(chunks) # Enforce 100-char minimum + for i, chunk in enumerate(chunks): + chunk.chunk_id = str(i + 1) + return chunks + + # Process text segments and tables in order + all_chunks = [] + current_pos = 0 + + for table_start, table_end in table_regions: + # Process text before this table + if current_pos < table_start: + text_before = exhibit_text[current_pos:table_start] + text_chunks = _chunk_non_table_text(text_before, offset=current_pos) + all_chunks.extend(text_chunks) + + # Create a chunk for the table itself + table_text = exhibit_text[table_start:table_end].strip() + if table_text: + table_chunk = ExhibitChunk( + chunk_id="", + section_header="TABLE", + text=table_text, + start_pos=table_start, + end_pos=table_end, + is_table=True, + ) + all_chunks.append(table_chunk) + + current_pos = table_end + + # Process any remaining text after the last table + if current_pos < len(exhibit_text): + text_after = exhibit_text[current_pos:] + text_chunks = _chunk_non_table_text(text_after, offset=current_pos) + all_chunks.extend(text_chunks) + + # Post-process: merge chunks where header == text (empty content chunks) + all_chunks = _merge_header_only_chunks(all_chunks) + + # Post-process: merge parent chunks smaller than 100 characters + all_chunks = _merge_small_parent_chunks(all_chunks) + + # Re-number chunk IDs + for i, chunk in enumerate(all_chunks): + chunk.chunk_id = str(i + 1) + + return all_chunks + + +def _merge_header_only_chunks(chunks: list[ExhibitChunk]) -> list[ExhibitChunk]: + """ + Merge chunks where the header is essentially the same as the entire text. + + This handles cases like: + Chunk 6: header="1.4", text="1.4" + Chunk 7: header="Schedule of", text="Schedule of Payment. This Agreement..." + + These should be merged into: + Chunk 6: header="1.4", text="1.4\nSchedule of Payment. This Agreement..." + + Args: + chunks: List of ExhibitChunk objects + + Returns: + List of merged ExhibitChunk objects + """ + if not chunks: + return chunks + + merged = [] + i = 0 + + while i < len(chunks): + current = chunks[i] + + # Check if current chunk is "header-only" (header == text, ignoring whitespace) + current_text_stripped = current.text.strip() + current_header_stripped = current.section_header.strip() + + is_header_only = ( + current_header_stripped and current_text_stripped == current_header_stripped + ) + + if is_header_only and i + 1 < len(chunks): + # Merge with next chunk + next_chunk = chunks[i + 1] + merged_text = current.text + "\n" + next_chunk.text + merged_header = ( + current.section_header + ) # Keep the original header (e.g., "1.4") + + merged_chunk = ExhibitChunk( + chunk_id=current.chunk_id, + section_header=merged_header, + text=merged_text, + start_pos=current.start_pos, + end_pos=next_chunk.end_pos, + ) + merged.append(merged_chunk) + i += 2 # Skip both current and next + else: + merged.append(current) + i += 1 + + return merged + + +def _is_section_header_chunk(chunk: ExhibitChunk) -> bool: + """ + Check if a chunk is primarily a section header (should merge forward, not backward). + + Detects patterns like: + - Roman numerals: "VI.", "IV. TITLE", "III. SERVICES" + - ARTICLE/SECTION headers: "ARTICLE ONE", "SECTION 5" + - Short all-caps titles: "RELATIONSHIP OF PARTIES" + - Exhibit/Schedule headers: "EXHIBIT A", "SCHEDULE 1" + + Args: + chunk: ExhibitChunk to check + + Returns: + True if chunk appears to be a section header + """ + text = chunk.text.strip() + + # Roman numeral headers: "VI.", "IV. TITLE", "III. SERVICES", etc. + if re.match(r"^[IVX]+\.?\s*[A-Z\s]*$", text, re.IGNORECASE): + return True + + # ARTICLE/SECTION/EXHIBIT/SCHEDULE headers + if re.match( + r"^(ARTICLE|SECTION|EXHIBIT|SCHEDULE|ADDENDUM|ATTACHMENT)\s+", + text, + re.IGNORECASE, + ): + return True + + # Short all-caps titles (< 60 chars, mostly uppercase, no lowercase content) + if len(text) < 60 and text.isupper() and len(text.split()) <= 6: + return True + + return False + + +def _merge_small_parent_chunks(chunks: list[ExhibitChunk]) -> list[ExhibitChunk]: + """ + Merge parent ExhibitChunks smaller than ESC_CONFIG.MIN_PARENT_CHUNK_SIZE (100 chars). + + Rules: + - First tiny chunk: merge with next chunk + - Header chunks (Roman numerals, ARTICLE, etc.): merge with NEXT chunk + - Other middle/end tiny chunks: merge with previous chunk + - All chunks tiny: return as single combined chunk + - Table chunks are never merged (they're treated as separate entities) + + Args: + chunks: List of ExhibitChunk objects + + Returns: + List of ExhibitChunk objects where no non-table chunk is < 100 chars + """ + if not chunks: + return chunks + + if len(chunks) == 1: + return chunks + + # Separate table chunks (never merged) from non-table chunks + table_chunks = [(i, c) for i, c in enumerate(chunks) if c.is_table] + non_table_chunks = [(i, c) for i, c in enumerate(chunks) if not c.is_table] + + if not non_table_chunks: + return chunks # All tables, nothing to merge + + # Check if all non-table chunks are tiny + if all(len(c.text) < ESC_CONFIG.MIN_PARENT_CHUNK_SIZE for _, c in non_table_chunks): + # Combine all non-table chunks into one + combined_text = "\n".join(c.text for _, c in non_table_chunks) + combined_chunk = ExhibitChunk( + chunk_id="", + section_header=non_table_chunks[0][1].section_header, + text=combined_text, + start_pos=non_table_chunks[0][1].start_pos, + end_pos=non_table_chunks[-1][1].end_pos, + ) + # Reconstruct list with tables in original positions + result = [combined_chunk] + for _, table in table_chunks: + result.append(table) + result.sort(key=lambda c: c.start_pos) + return result + + # Process non-table chunks, merging small ones + merged_non_table = [] + i = 0 + non_table_list = [c for _, c in non_table_chunks] + + while i < len(non_table_list): + current = non_table_list[i] + + if len(current.text) >= ESC_CONFIG.MIN_PARENT_CHUNK_SIZE: + merged_non_table.append(current) + i += 1 + else: + # Current chunk is too small, need to merge + is_header = _is_section_header_chunk(current) + has_next = i + 1 < len(non_table_list) + + if not merged_non_table or (is_header and has_next): + # First chunk OR header chunk with next available - merge FORWARD + combined_text = current.text + combined_start = current.start_pos + combined_header = current.section_header + i += 1 + while ( + i < len(non_table_list) + and len(combined_text) < ESC_CONFIG.MIN_PARENT_CHUNK_SIZE + ): + combined_text = combined_text + "\n" + non_table_list[i].text + i += 1 + combined_end = ( + non_table_list[i - 1].end_pos if i > 0 else current.end_pos + ) + merged_non_table.append( + ExhibitChunk( + chunk_id="", + section_header=combined_header, + text=combined_text, + start_pos=combined_start, + end_pos=combined_end, + ) + ) + else: + # Middle/end non-header chunk - merge with PREVIOUS + prev_chunk = merged_non_table[-1] + merged_text = prev_chunk.text + "\n" + current.text + merged_non_table[-1] = ExhibitChunk( + chunk_id=prev_chunk.chunk_id, + section_header=prev_chunk.section_header, + text=merged_text, + start_pos=prev_chunk.start_pos, + end_pos=current.end_pos, + ) + i += 1 + + # Combine merged non-table chunks with table chunks + result = merged_non_table + [c for _, c in table_chunks] + result.sort(key=lambda c: c.start_pos) + + return result + + +def _compute_chunk_embeddings(chunks: list[ExhibitChunk]) -> list[ExhibitChunk]: + """ + Compute embeddings for all chunks using AWS Bedrock Titan. + + For chunks with sub-chunks (large non-table chunks), embeddings are computed + for the sub-chunks instead of the parent chunk. This enables fine-grained + semantic search while returning the full parent chunk on match. + + For small chunks (<= 500 chars) and table chunks, the chunk itself is embedded. + + Args: + chunks: List of ExhibitChunk objects without embeddings + + Returns: + Same list of chunks with embeddings populated (on chunk or sub-chunks) + """ + if not chunks: + return chunks + + try: + embeddings_model = get_embeddings( + HSC_CONFIG.EMBEDDING_MODEL_ID, + HSC_CONFIG.BEDROCK_REGION, + ) + + direct_embed_chunks = [] + sub_chunks_to_embed = [] + + for chunk in chunks: + if chunk.is_table: + # Table chunks: embed the full table + direct_embed_chunks.append(chunk) + elif chunk.has_sub_chunks(): + # Large chunk: collect sub-chunks for embedding + sub_chunks_to_embed.extend(chunk.sub_chunks) + else: + # Small chunk: embed directly + direct_embed_chunks.append(chunk) + + # Batch embed direct chunks + if direct_embed_chunks: + direct_texts = [c.text for c in direct_embed_chunks] + direct_embeddings = embeddings_model.embed_documents(direct_texts) + for chunk, embedding in zip(direct_embed_chunks, direct_embeddings): + chunk.embedding = embedding + + # Batch embed sub-chunks + if sub_chunks_to_embed: + sub_texts = [sc.text for sc in sub_chunks_to_embed] + sub_embeddings = embeddings_model.embed_documents(sub_texts) + for sub_chunk, embedding in zip(sub_chunks_to_embed, sub_embeddings): + sub_chunk.embedding = embedding + + except Exception as e: + logging.error(f"Failed to embed chunks: {e}") + + return chunks + + +def search_reimbursement_chunks( + chunks: list[ExhibitChunk], + threshold: float = 0.30, + queries: Optional[List[str]] = None, + keywords: Optional[List[str]] = None, + use_semantic: bool = True, + use_keywords: bool = True, +) -> list[ExhibitChunk]: + """ + Retrieve chunks related to reimbursements using hybrid keyword + semantic search. + + Uses a two-stage approach for maximum recall: + 1. Keyword matching: Chunks containing reimbursement keywords are automatically included + 2. Semantic search: Searches sub-chunks (for large chunks) or chunks directly, + returns parent chunk if any sub-chunk matches above threshold + + Table chunks are ALWAYS included in results since tables typically contain + reimbursement data. Final results are ordered by chunk_id to maintain document order. + + Args: + chunks: List of ExhibitChunk objects to search + threshold: Minimum similarity score (0-1) for semantic match. Default 0.30. + queries: Optional custom semantic queries. If None, uses REIMBURSEMENT_QUERIES. + keywords: Optional custom keywords. If None, uses REIMBURSEMENT_KEYWORDS. + use_semantic: Whether to use semantic search in addition to keywords. Default True. + use_keywords: Whether to use keyword matching. Default True. + + Returns: + List of ExhibitChunk objects ordered by chunk_id (document order). + Includes: all table chunks + keyword-matched chunks + semantic-matched chunks. + """ + if not chunks: + return [] + + # Separate table chunks from non-table chunks + table_chunks = [c for c in chunks if c.is_table] + non_table_chunks = [c for c in chunks if not c.is_table] + + if not non_table_chunks and not table_chunks: + logging.warning("No chunks available for reimbursement search") + return [] + + search_keywords = keywords if keywords is not None else REIMBURSEMENT_KEYWORDS + search_queries = queries if queries is not None else REIMBURSEMENT_QUERIES + + # Track matched chunk_ids from different sources + keyword_matched_ids: set[str] = set() + semantic_matched_ids: set[str] = set() + + # Stage 1: Keyword matching (fast, no API calls) + if use_keywords: + for chunk in non_table_chunks: + if keyword_match(chunk.text, search_keywords): + keyword_matched_ids.add(chunk.chunk_id) + + # Stage 2: Semantic search with sub-chunk support + if use_semantic: + # Collect all searchable items: (parent_chunk_id, embedding) + search_items: List[tuple[str, List[float]]] = [] + + for chunk in non_table_chunks: + if chunk.has_sub_chunks(): + # Search sub-chunks, map back to parent + for sub_chunk in chunk.sub_chunks: + if sub_chunk.embedding is not None: + search_items.append((chunk.chunk_id, sub_chunk.embedding)) + elif chunk.embedding is not None: + # Search chunk directly + search_items.append((chunk.chunk_id, chunk.embedding)) + + if search_items: + try: + embeddings_model = get_embeddings( + HSC_CONFIG.EMBEDDING_MODEL_ID, + HSC_CONFIG.BEDROCK_REGION, + ) + + # Batch embed all queries in a single API call + query_embeddings = embeddings_model.embed_documents(search_queries) + + for query_embedding in query_embeddings: + # Compute similarities for all search items + for parent_id, embedding in search_items: + similarity = cosine_similarity(query_embedding, embedding) + + # Include parent chunk if any sub-chunk/chunk exceeds threshold + if similarity >= threshold: + semantic_matched_ids.add(parent_id) + + except Exception as e: + logging.error(f"Failed to compute semantic similarity: {e}") + + # Collect all table chunk IDs (always included) + table_chunk_ids = {c.chunk_id for c in table_chunks} + + # Combine: all table chunks + keyword matches + semantic matches + all_selected_ids = table_chunk_ids | keyword_matched_ids | semantic_matched_ids + + # Get chunks and sort by chunk_id to maintain document order + result_chunks = [c for c in chunks if c.chunk_id in all_selected_ids] + result_chunks.sort(key=lambda c: int(c.chunk_id)) + + return result_chunks diff --git a/src/pipelines/shared/preprocessing/preprocess.py b/src/pipelines/shared/preprocessing/preprocess.py index 7b9b2d8..35b3cc7 100644 --- a/src/pipelines/shared/preprocessing/preprocess.py +++ b/src/pipelines/shared/preprocessing/preprocess.py @@ -1,7 +1,10 @@ import logging +from typing import Dict, Tuple +from collections import Counter, defaultdict from typing import TYPE_CHECKING from src import config +from src.pipelines.saas.prompts import prompt_calls from src.pipelines.shared.preprocessing import preprocessing_funcs from src.utils import timing_utils @@ -78,12 +81,13 @@ def one_to_n_exhibit_chunking( pages_dict=None, EXHIBIT_HEADER_MARKERS=None, filename=None, -) -> tuple[dict, dict]: +) -> dict[str, list[str]]: """ - Process the text dictionary or pages dictionary to identify and chunk exhibits. + Process the text dictionary or pages dictionary to identify exhibit headers. - This function extracts exhibit pages, maps them to their respective chunks, - and identifies reimbursement exhibits using LLM. + This function extracts exhibit headers and their page locations, returning + a dictionary that maps page numbers to lists of headers starting on that page. + Consecutive pages with the same header are merged into a single exhibit. Args: text_dict: A dictionary where keys are page numbers and values are the text @@ -94,11 +98,9 @@ def one_to_n_exhibit_chunking( filename: The name of the file being processed. Returns: - tuple: A tuple containing: - - dict: A dictionary where keys are exhibit page numbers and values are - lists of page numbers that belong to that exhibit. - - dict: A dictionary where keys are exhibit page numbers and values are - the exhibit headers. + dict[str, list[str]]: A dictionary where keys are page numbers and values are + lists of exhibit headers starting on that page. + Example: {'2': ['ARTICLE ONE', 'ARTICLE TWO'], '4': ['ARTICLE THREE']} Note: Either text_dict or pages_dict must be provided. If both are provided, @@ -106,54 +108,32 @@ def one_to_n_exhibit_chunking( """ # Determine which dictionary to use if pages_dict is not None: - with timing_utils.timed_block( - "exhibit_chunking.get_exhibit_pages", context=filename - ): - exhibit_pages, all_exhibit_headers = preprocessing_funcs.get_exhibit_pages( - pages_dict=pages_dict, - EXHIBIT_HEADER_MARKERS=EXHIBIT_HEADER_MARKERS, - filename=filename, - ) - with timing_utils.timed_block( - "exhibit_chunking.link_exhibit_pages", context=filename - ): - exhibit_pages, all_exhibit_headers = preprocessing_funcs.link_exhibit_pages( - all_exhibit_headers, pages_dict=pages_dict, filename=filename - ) - with timing_utils.timed_block( - "exhibit_chunking.chunk_by_exhibit", context=filename - ): - exhibit_chunk_mapping = preprocessing_funcs.chunk_by_exhibit( - pages_dict=pages_dict, exhibit_pages=exhibit_pages - ) + _, exhibit_header_dict = preprocessing_funcs.get_exhibit_pages_new( + pages_dict=pages_dict, + EXHIBIT_HEADER_MARKERS=EXHIBIT_HEADER_MARKERS, + filename=filename, + ) elif text_dict is not None: - with timing_utils.timed_block( - "exhibit_chunking.get_exhibit_pages", context=filename - ): - exhibit_pages, all_exhibit_headers = preprocessing_funcs.get_exhibit_pages( - text_dict=text_dict, - EXHIBIT_HEADER_MARKERS=EXHIBIT_HEADER_MARKERS, - filename=filename, - ) - with timing_utils.timed_block( - "exhibit_chunking.link_exhibit_pages", context=filename - ): - exhibit_pages, all_exhibit_headers = preprocessing_funcs.link_exhibit_pages( - all_exhibit_headers, text_dict=text_dict, filename=filename - ) - with timing_utils.timed_block( - "exhibit_chunking.chunk_by_exhibit", context=filename - ): - exhibit_chunk_mapping = preprocessing_funcs.chunk_by_exhibit( - text_dict=text_dict, exhibit_pages=exhibit_pages - ) + _, exhibit_header_dict = preprocessing_funcs.get_exhibit_pages_new( + text_dict=text_dict, + EXHIBIT_HEADER_MARKERS=EXHIBIT_HEADER_MARKERS, + filename=filename, + ) else: raise ValueError("Either text_dict or pages_dict must be provided") - logging.debug( - f"Exhibit chunking complete: {len(exhibit_chunk_mapping)} exhibits identified - {filename}" + # Deduplicate headers across pages using LLM (handles case variations, + # appended sub-lines, and other fuzzy duplicates) + exhibit_header_dict = prompt_calls.prompt_header_deduplication( + exhibit_header_dict, filename ) - return exhibit_chunk_mapping, all_exhibit_headers + + # Sort keys numerically (they are strings like "1", "4", "14") + exhibit_header_dict = dict( + sorted(exhibit_header_dict.items(), key=lambda x: int(x[0].split(".")[0])) + ) + + return exhibit_header_dict def find_headers_and_footers(text_dict: dict[str, str]) -> tuple[dict, str, str]: @@ -230,3 +210,110 @@ def find_headers_and_footers(text_dict: dict[str, str]) -> tuple[dict, str, str] modified_text_dict[page_num] = page_text return modified_text_dict, header_marker, footer_marker + + +def clean_header_footer( + text_dict: Dict[str, str], +) -> Tuple[Dict[str, str], Dict[str, str]]: + """ + Advanced document preprocessing function that: + 1. Identifies and removes common headers/footers + 2. Removes lines that appear across all pages + 3. Provides detailed metadata about removed content + + Args: + text_dict: Dictionary with page numbers as keys and page text as values + + Returns: + Tuple containing: + 1. Cleaned document dictionary + 2. Metadata about removed content + """ + if not text_dict or len(text_dict) == 0: + return text_dict, {} + + # Preprocessing: Split pages into lines and words + pages_lines = {} + pages_words = {} + for page_num, text in text_dict.items(): + # Strip leading/trailing whitespace + text = text.strip() + pages_lines[page_num] = text.split("\n") + pages_words[page_num] = text.split() + + # Metadata dictionary to track removed content + removal_metadata = { + "common_lines": set(), + "header_markers": [], + "footer_markers": [], + } + + # 1. Identify common headers (word-based approach) + header_words = [] + min_word_count = min(len(page_words) for page_words in pages_words.values()) + for i in range(min_word_count): + word = list(pages_words.values())[0][i] + if all(page_words[i] == word for page_words in pages_words.values()): + header_words.append(word) + else: + break + header_marker = " ".join(header_words) + if header_marker: + removal_metadata["header_markers"] = header_words + + # 2. Identify common footers + footer_words = [] + for i in range(1, min_word_count + 1): + word = list(pages_words.values())[0][-i] + if all(page_words[-i] == word for page_words in pages_words.values()): + footer_words.insert(0, word) + else: + break + footer_marker = " ".join(footer_words) + if footer_marker: + removal_metadata["footer_markers"] = footer_words + + # 3. Identify lines that appear in ALL pages + line_page_count = Counter() + total_pages = len(text_dict) + + for page_lines in pages_lines.values(): + # Use set to count each unique line only once per page + unique_lines_in_page = set(line.strip() for line in page_lines if line.strip()) + for line in unique_lines_in_page: + line_page_count[line] += 1 + + # Find lines that appear in ALL pages (minimum 15 characters, exclude "signature") + common_lines = { + line + for line, count in line_page_count.items() + if count == total_pages + and len(line.strip()) >= 15 + and line.strip().lower() != "signature" + } + removal_metadata["common_lines"] = common_lines + + # 4. Clean the document + cleaned_dict = {} + first_page_num = list(text_dict.keys())[0] + for page_num, lines in pages_lines.items(): + # Remove header (except for first page) and footer (all pages, except "signature") + if header_marker and page_num != first_page_num: + lines = [line for line in lines if not line.startswith(header_marker)] + if footer_marker: + lines = [ + line + for line in lines + if not line.endswith(footer_marker) + or line.strip().lower() == "signature" + ] + + # Remove common lines (skip page 1 entirely) + if page_num == first_page_num: + cleaned_lines = lines + else: + cleaned_lines = [line for line in lines if line.strip() not in common_lines] + + cleaned_dict[page_num] = "\n".join(cleaned_lines).strip() + + return cleaned_dict, removal_metadata diff --git a/src/pipelines/shared/preprocessing/preprocessing_funcs.py b/src/pipelines/shared/preprocessing/preprocessing_funcs.py index f89b901..5e94fd8 100644 --- a/src/pipelines/shared/preprocessing/preprocessing_funcs.py +++ b/src/pipelines/shared/preprocessing/preprocessing_funcs.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Optional import src.utils.string_utils as string_utils from src.pipelines.saas.prompts import prompt_calls from src import config +from src.constants.delimiters import TABLE_START_MARKER, TABLE_END_MARKER if TYPE_CHECKING: from src.pipelines.shared.extraction.page_funcs import Page @@ -376,14 +377,15 @@ def simplify_exhibit( exhibit: Exhibit object (preferred - uses exhibit.pages and exhibit.exhibit_page_nums) exhibit_page_nums: List of page numbers that belong to this exhibit (required if exhibit is not provided) - current_page_num: The current page being processed (will NOT be simplified) - (required if exhibit is not provided) + current_page_num: The current page being processed (will NOT be simplified). + If None, ALL pages will be simplified (useful for chunk-based processing). Returns: str: The exhibit text with tables simplified on all pages except current_page_num + (or all pages if current_page_num is None) Note: - Either (exhibit) or (text_dict/pages_dict + exhibit_page_nums + current_page_num) must be provided. + Either (exhibit) or (text_dict/pages_dict + exhibit_page_nums) must be provided. If exhibit is provided, it takes precedence. """ # Check if any parameter that should be a dict is actually an Exhibit object @@ -409,10 +411,8 @@ def simplify_exhibit( exhibit_page_nums = exhibit.exhibit_page_nums # For current_page_num, we need to determine which base page the current_page_num belongs to # If current_page_num is a sub-page like "27.1", the base page is "27" - if current_page_num is None: - # Try to infer from exhibit_page_nums - use first page as default - current_page_num = exhibit_page_nums[0] if exhibit_page_nums else "" - else: + # If None, we'll simplify all pages uniformly (for chunk-based processing) + if current_page_num is not None: # Extract base page if it's a sub-page if "." in current_page_num and not current_page_num.startswith("."): current_page_num = current_page_num.rsplit(".", 1)[0] @@ -422,8 +422,6 @@ def simplify_exhibit( raise ValueError( "exhibit_page_nums must be provided if exhibit is not provided" ) - if current_page_num is None: - raise ValueError("current_page_num must be provided if exhibit is not provided") # Determine which dictionary to use use_pages = False @@ -480,8 +478,13 @@ def simplify_exhibit( table_counter = 1 # Extract base page from current_page_num if it's a sub-page + # If None, we'll simplify all pages uniformly current_base_page = current_page_num - if "." in current_page_num and not current_page_num.startswith("."): + if ( + current_page_num is not None + and "." in current_page_num + and not current_page_num.startswith(".") + ): current_base_page = current_page_num.rsplit(".", 1)[0] for page_id in exhibit_page_nums: @@ -507,7 +510,8 @@ def simplify_exhibit( page_text = page_dict.get(page_id, "") # Check if this is the current page (match base page) - if base_page == current_base_page: + # If current_base_page is None, simplify all pages + if current_base_page is not None and base_page == current_base_page: # Don't simplify the current page - keep it as-is simplified_exhibit_pages.append(page_text) # But still need to count tables to keep numbering consistent @@ -516,7 +520,8 @@ def simplify_exhibit( page_text, table_counter ) else: - # Simplify other pages by replacing tables with placeholders + # Simplify pages by replacing tables with placeholders + # (either other pages, or ALL pages if current_base_page is None) if "-------Table Start--------" in page_text: simplified_page, table_counter = replace_tables_with_placeholder( page_text, table_counter @@ -576,3 +581,184 @@ def split_large_tables( pages_dict[page_num] = page return pages_dict + + +def identify_section_boundaries(page_text: str, context_chars: int = 100) -> str: + import re + + """ + Identifies potential section/exhibit headers in a contract page and returns + a compact string with context around each potential boundary. + + Args: + page_text: Text content from a single page + context_chars: Number of characters to include before/after each header + + Returns: + A string with markers and context for each potential section boundary + """ + + # Healthcare contract section patterns (order matters - more specific first) + section_patterns = [ + (r"^\s*EXHIBIT\s+[A-Z0-9]+", "EXHIBIT"), + ( + r"^\s*ARTICLE\s+[A-Z]+", + "ARTICLE", + ), # Catches "ARTICLE SIX", "ARTICLE ONE", etc. + (r"^\s*ARTICLE\s+[IVX0-9]+", "ARTICLE"), # Roman numerals and numbers + # (r'^\s*SECTION\s+\d+(\.\d+)*', 'SECTION'), + (r"^\s*SCHEDULE\s+[A-Z0-9]+", "SCHEDULE"), + (r"^\s*APPENDIX\s+[A-Z0-9]+", "APPENDIX"), + (r"^\s*ATTACHMENT\s+[A-Z0-9]+", "ATTACHMENT"), + # (r'^\s*\d+\.\d+\s+[A-Z]', 'SUBSECTION'), # "6.1 Indemnification" + (r"^\s*[A-Z][A-Z\s]{20,}$", "HEADER"), # All-caps headers (20+ chars) + # (r'^\s*\d+\.\s+[A-Z][A-Z\s]+', 'NUMBERED_SECTION'), # "1. DEFINITIONS" + ] + + lines = page_text.split("\n") + potential_boundaries = [] + + # Track whether we're inside a table region + inside_table = False + + # Find all potential section headers + for line_idx, line in enumerate(lines): + # Check for table boundary markers and update state + if TABLE_START_MARKER in line: + inside_table = True + continue + if TABLE_END_MARKER in line: + inside_table = False + continue + + # Skip any lines inside table boundaries + if inside_table: + continue + + # Skip empty lines + if not line.strip(): + continue + + for pattern, section_type in section_patterns: + if re.match(pattern, line): + # Calculate character position in original text + char_pos = len("\n".join(lines[:line_idx])) + potential_boundaries.append( + { + "line_idx": line_idx, + "char_pos": char_pos, + "type": section_type, + "header_text": line.strip(), + } + ) + break # Only match first pattern per line + + # If no boundaries found, return message + if not potential_boundaries: + return "NO_SECTION_BOUNDARIES_DETECTED" + + # Build output string with context + result_parts = [] + + for i, boundary in enumerate(potential_boundaries): + char_pos = boundary["char_pos"] + + # Extract context + start_pos = max(0, char_pos - context_chars) + end_pos = min( + len(page_text), char_pos + len(boundary["header_text"]) + context_chars + ) + + context = page_text[start_pos:end_pos] + + # Add marker + marker = f"\n\n===BOUNDARY_{i+1}_{boundary['type']}===\n" + result_parts.append(marker + context) + + result_parts.append("\n\n===END_OF_BOUNDARIES===") + + return "".join(result_parts) + + +def get_exhibit_pages_new( + text_dict: Optional[dict[str, str]] = None, + pages_dict: Optional[dict[str, "Page"]] = None, + EXHIBIT_HEADER_MARKERS: Optional[list] = None, + filename: Optional[str] = None, +) -> tuple[list[str], dict[str, list[str]]]: + """ + Identify exhibit pages from the text dictionary or pages dictionary. + + This function processes the dictionary to identify pages that are considered exhibits. + It uses the first page as an exhibit and checks subsequent pages for exhibit headers. + If a page is identified as an exhibit, it is added to the exhibit pages list and its + header is stored in the exhibit header dictionary. + If a page does not have a header or is not an exhibit, it is skipped. + + Args: + text_dict: A dictionary where keys are page numbers (as strings) and values are + the text of those pages (for backward compatibility) + pages_dict: A dictionary where keys are page numbers and values are Page objects + (preferred) + EXHIBIT_HEADER_MARKERS: List of markers to identify exhibit headers + filename: The name of the file being processed, used for logging and LLM invocation. + + Returns: + tuple[list[str], dict[str, list[str]]]: A tuple containing: + - A list of page numbers (as strings) that are identified as exhibit pages. + - A dictionary where keys are page numbers and values are lists of unique + exhibit headers found on that page (excluding N/A values). + + Note: + Either text_dict or pages_dict must be provided. If both are provided, + pages_dict takes precedence. + """ + if EXHIBIT_HEADER_MARKERS is None: + EXHIBIT_HEADER_MARKERS = [] + if filename is None: + filename = "" + + # Determine which dictionary to use + if pages_dict is not None: + page_dict = pages_dict + use_pages = True + elif text_dict is not None: + page_dict = text_dict + use_pages = False + else: + raise ValueError("Either text_dict or pages_dict must be provided") + + sorted_pages = sorted(page_dict.keys(), key=string_utils.page_key_sort) + exhibit_pages = [] + exhibit_header_dict = {} + + is_exhibit = False + for page_num in sorted_pages: + # Get page content - handle both Page objects and strings + if use_pages: + page_obj = page_dict[page_num] + page_content = page_obj.get_text() + else: + page_content = page_dict[page_num] + + if "." not in page_num or ".0" in page_num: + header_response = prompt_calls.prompt_exhibit_header( + page_content, EXHIBIT_HEADER_MARKERS, filename + ) + + # header_response is dict[str, str] - extract unique non-N/A values preserving insertion order + seen = set() + unique_headers = [] + for val in header_response.values(): + if val.upper().strip() != "N/A" and val not in seen: + seen.add(val) + unique_headers.append(val) + if not unique_headers: + is_exhibit = False + else: + exhibit_pages.append(page_num) + exhibit_header_dict[page_num] = unique_headers + is_exhibit = True + elif is_exhibit: + exhibit_pages.append(page_num) + return exhibit_pages, exhibit_header_dict diff --git a/src/prompts/prompt_templates.py b/src/prompts/prompt_templates.py index 1eef9f9..05b49a2 100644 --- a/src/prompts/prompt_templates.py +++ b/src/prompts/prompt_templates.py @@ -203,6 +203,7 @@ def get_cacheable_instructions(): # Preprocessing instructions (exhibit headers, linkage) try: cacheable_instructions["EXHIBIT_HEADER"] = EXHIBIT_HEADER_INSTRUCTION() + cacheable_instructions["EXHIBIT_HEADER_NEW"] = EXHIBIT_HEADER_NEW_INSTRUCTION() cacheable_instructions["EXHIBIT_LINKAGE"] = EXHIBIT_LINKAGE_INSTRUCTION() cacheable_instructions["EXHIBIT_TITLE_MATCH"] = ( EXHIBIT_TITLE_MATCH_INSTRUCTION() @@ -1903,6 +1904,70 @@ Here is the text to analyze: return (prompt, _json_list_parser) +def EXHIBIT_HEADER_NEW_INSTRUCTION() -> str: + """Static instruction for EXHIBIT_HEADER_NEW prompt caching. + Returns formatted instructions for extracting headers from chunked contract pages. + """ + return ( + "[OBJECTIVE]\n" + "Extract section headers from contract page chunks with boundary markers.\n" + "Return a JSON dict in |pipes| with boundary numbers as keys and header as values." + ) + + +def EXHIBIT_HEADER_NEW(context, EXHIBIT_HEADER_MARKERS): + """ + Returns prompt for exhibit header extraction from chunked contract pages. + Chunks are separated by boundary markers like ===BOUNDARY_X_HEADER===. + Call EXHIBIT_HEADER_NEW_INSTRUCTION() separately for the cached instruction. + + Args: + context (str): Contract text with boundary markers separating chunks + EXHIBIT_HEADER_MARKERS (list): List of keywords/phrases that indicate headers + + Returns: + str: Formatted prompt for LLM to extract headers and return as JSON dict + """ + return f"""Analyze the following contract page chunks and extract section headers from each chunk. + +Each chunk is separated by a boundary marker in the format ===BOUNDARY_X_HEADER=== where X is the boundary number. + +The following keywords and phrases should be considered as exhibit/attachment headers: +* {'\n* '.join(EXHIBIT_HEADER_MARKERS)} + +Rules for Inclusion: +* Header can be multi-lined, include all lines that are part of the header +* Include ALL text that appears to be part of a section header +* Full header names and numbers/letters (e.g., "Attachment C: Commercial-Exchange") +* Document titles, agreement names, and section identifiers +* Any words that specify what type of information is contained, such as "Compensation Schedule" or "Definitions" +* Provider/Entity names when listed with exhibit information +* Keep exact capitalization and formatting as shown in document + +Rules for Exclusion: +* Do NOT include body text or paragraph content +* Do NOT include page numbers +* Do NOT include docusign IDs and other metadata +* Do NOT include partial sentences or continuation text +* Do NOT include form field labels like "Provider's Legal Name" or "Authorized Representative's Signature" + +[OUTPUT FORMAT] +Return a JSON dictionary enclosed in |pipes| where: +- Keys are boundary numbers (e.g., "1", "2", "3") +- Values are the header strings found in that chunk +- If no headers are found in a chunk, return "N/A" for that key +- Return {{}} if no headers are found in any chunk + +Example output: +|{{"1": "PROVIDER SERVICES AGREEMENT\\nSIGNATURE PAGE", "2": "N/A", "3": "Attachment A: Fee Schedule"}}| + +Here are the chunks to analyze: +{context.replace('"', "'")} + +Return ONLY the |{{dictionary}}| followed by a brief explanation. +""" + + def EXHIBIT_LINKAGE_INSTRUCTION() -> str: """Static instruction for EXHIBIT_LINKAGE prompt caching. Contains all decision rules for determining meaningful differences. @@ -1941,6 +2006,56 @@ Header 2: return (prompt, _json_list_parser) +def EXHIBIT_HEADER_DEDUP_INSTRUCTION() -> str: + """Static instruction for EXHIBIT_HEADER_DEDUP prompt caching. + Returns formatted instructions for deduplicating exhibit headers across pages. + """ + return ( + "[OBJECTIVE]\n" + "Deduplicate a dictionary of page-number-to-header mappings from a contract.\n" + "Keep only the first page occurrence of each unique section.\n" + "Return the deduplicated JSON dict in |pipes|." + ) + + +def EXHIBIT_HEADER_DEDUP(header_dict_str): + """ + Returns prompt for deduplicating exhibit headers across contract pages. + Call EXHIBIT_HEADER_DEDUP_INSTRUCTION() separately for the cached instruction. + + Args: + header_dict_str (str): JSON string of the exhibit header dictionary + mapping page numbers to lists of headers. + + Returns: + str: Formatted prompt for LLM to deduplicate headers and return as JSON dict + """ + return f"""You are given a JSON dictionary where keys are page numbers (as strings) and values are lists containing header text extracted from each page of a contract document. + +Multiple consecutive pages may belong to the same section and will have similar or identical headers repeated on each page, sometimes with minor variations. + +Your task is to deduplicate this dictionary so that each unique section appears only once, on the first page where it starts. + +[RULES FOR IDENTIFYING DUPLICATES] +* Two headers belong to the same section if they share the same core section identifier AND the same content-type descriptors after case-normalization. +* Minor case variations do NOT make headers different sections. +* Extra sub-lines that are facility names, entity names, or provider names do NOT make it a different section — these are just contextual metadata repeated on continuation pages. +* However, sub-lines that describe the TYPE OF CONTENT or SERVICE CATEGORY DO make it a DIFFERENT section, even if the top-level exhibit/addendum identifier is the same. These descriptors indicate distinct sub-sections with different subject matter under a shared parent exhibit. +* A genuinely NEW section has either a different top-level identifier OR a different content-type/service-category descriptor under the same identifier. +* If a single header on one page bundles multiple distinct section identifiers together (e.g., lists several addenda, exhibits, or schedules within one header block), it is likely a table of contents or index page — NOT an actual section start. When the individual sections referenced in such a bundled header each appear as their own standalone headers on later pages, DROP the bundled/index header entirely and KEEP the individual section headers from later pages. + +[CRITICAL RULE] +* Do NOT modify the header text of the first occurrence of any section. The exact original text MUST be preserved character-for-character because it is used downstream for regex matching to split the document into exhibits. Any modification will break downstream processing. + +[OUTPUT FORMAT] +Return a JSON dictionary enclosed in |pipes| with the same structure as the input: keys are page numbers, values are lists of header strings. The output should contain ONLY the first page of each unique section, with the original header text preserved exactly as-is. + +[INPUT] +{header_dict_str} + +Briefly explain which headers you identified as duplicates and why, then return the deduplicated dictionary in |pipes|.""" + + ##################################################################################### ########################## CHECKS, FIXES, AND VALIDATIONS ########################### ##################################################################################### @@ -1978,7 +2093,7 @@ Must contain a SPECIFIC, CONCRETE payment method you can point to. Valid payment methods: - Dollar amounts (e.g., "$150 per visit") - Percentage rates (e.g., "80% of Medicare allowable") -- Named fee schedule references with specific percentage (e.g., "100% of DMAP fee schedule") +- Named fee schedule references with specific percentage (e.g., "100% of DMAP fee schedule"). Note: a term that establishes or develops a fee schedule at a specific percentage of a named schedule IS valid (e.g., "develop a fee schedule at 110% of Medicaid Fee Schedule"). - Payment formulas or limitations with amounts - "Lesser of" or "greater of" statements with specific rates - Capitation rates (PMPM, PMPY, per member costs) diff --git a/src/tests/test_exhibit_funcs.py b/src/tests/test_exhibit_funcs.py index 84e32ae..9db8349 100644 --- a/src/tests/test_exhibit_funcs.py +++ b/src/tests/test_exhibit_funcs.py @@ -4,13 +4,645 @@ import sys if "write_to_s3=False" not in sys.argv: sys.argv.append("write_to_s3=False") import unittest -from unittest.mock import patch, Mock +from unittest.mock import patch, Mock, MagicMock -from src.pipelines.shared.extraction.exhibit_funcs import Exhibit, get_exhibit_list +from src.pipelines.shared.extraction.exhibit_funcs import ( + Exhibit, + get_exhibit_list, + create_exhibits_from_header_dict, + _find_header_in_text, +) +from src.pipelines.shared.preprocessing.exhibit_smart_chunking_funcs import ( + ExhibitChunk, + SubChunk, + _split_text_at_sentence_boundaries, + _find_table_regions, + _chunk_non_table_text, + _chunk_exhibit_text, + _merge_header_only_chunks, + _merge_small_parent_chunks, + _is_section_header_chunk, + ESC_CONFIG, +) +from src.constants.delimiters import TABLE_START_MARKER, TABLE_END_MARKER +from src.utils.string_utils import keyword_match as _keyword_match +from src.utils.rag_utils import cosine_similarity + +# Constants from config class +DEFAULT_SUBCHUNK_SIZE = ESC_CONFIG.DEFAULT_SUBCHUNK_SIZE +MIN_PARENT_CHUNK_SIZE = ESC_CONFIG.MIN_PARENT_CHUNK_SIZE from src.pipelines.shared.extraction.page_funcs import Page from src.prompts.fieldset import FieldSet +# ============================================================================ +# Tests for SubChunk +# ============================================================================ + + +class TestSubChunk(unittest.TestCase): + + def test_basic_creation(self): + sc = SubChunk(sub_chunk_id="1.1", parent_chunk_id="1", text="Hello world") + self.assertEqual(sc.sub_chunk_id, "1.1") + self.assertEqual(sc.parent_chunk_id, "1") + self.assertEqual(sc.text, "Hello world") + self.assertIsNone(sc.embedding) + + def test_with_embedding(self): + emb = [0.1, 0.2, 0.3] + sc = SubChunk( + sub_chunk_id="2.1", parent_chunk_id="2", text="test", embedding=emb + ) + self.assertEqual(sc.embedding, emb) + + def test_repr_short_text(self): + sc = SubChunk(sub_chunk_id="1.1", parent_chunk_id="1", text="Short") + r = repr(sc) + self.assertIn("1.1", r) + self.assertIn("Short", r) + self.assertIn("no", r) # no embedding + + def test_repr_long_text(self): + sc = SubChunk(sub_chunk_id="1.1", parent_chunk_id="1", text="A" * 100) + r = repr(sc) + self.assertIn("...", r) + + +# ============================================================================ +# Tests for _split_text_at_sentence_boundaries +# ============================================================================ + + +class TestSplitTextAtSentenceBoundaries(unittest.TestCase): + + def test_empty_text(self): + self.assertEqual(_split_text_at_sentence_boundaries(""), []) + + def test_none_text(self): + self.assertEqual(_split_text_at_sentence_boundaries(""), []) + + def test_short_text_returned_as_is(self): + text = "Short sentence." + result = _split_text_at_sentence_boundaries(text, max_size=500) + self.assertEqual(result, [text]) + + def test_splits_at_sentence_boundaries(self): + # Build text that exceeds max_size + sentences = ["This is sentence one. " * 5, "This is sentence two. " * 5] + text = " ".join(sentences) + result = _split_text_at_sentence_boundaries(text, max_size=100) + self.assertGreater(len(result), 1) + # Rejoined text should preserve all content + for chunk in result: + self.assertTrue(len(chunk.strip()) > 0) + + def test_single_long_sentence_no_split(self): + # A single sentence with no sentence boundaries + text = "A" * 600 + result = _split_text_at_sentence_boundaries(text, max_size=500) + self.assertEqual(len(result), 1) + self.assertEqual(result[0], text) + + +# ============================================================================ +# Tests for ExhibitChunk +# ============================================================================ + + +class TestExhibitChunk(unittest.TestCase): + + def test_basic_creation(self): + chunk = ExhibitChunk( + chunk_id="1", + section_header="1.1", + text="Some text", + start_pos=0, + end_pos=9, + ) + self.assertEqual(chunk.chunk_id, "1") + self.assertEqual(chunk.section_header, "1.1") + self.assertFalse(chunk.is_table) + self.assertIsNone(chunk.embedding) + + def test_table_chunk_no_sub_chunks(self): + chunk = ExhibitChunk( + chunk_id="1", + section_header="TABLE", + text="A" * 600, + start_pos=0, + end_pos=600, + is_table=True, + ) + self.assertFalse(chunk.has_sub_chunks()) + self.assertEqual(chunk.sub_chunks, []) + + def test_small_chunk_no_sub_chunks(self): + chunk = ExhibitChunk( + chunk_id="1", + section_header="1.1", + text="Short text", + start_pos=0, + end_pos=10, + ) + self.assertFalse(chunk.has_sub_chunks()) + + def test_large_chunk_creates_sub_chunks(self): + # Build text > 500 chars with sentence boundaries + text = "First sentence. " * 20 + "Second sentence. " * 20 + chunk = ExhibitChunk( + chunk_id="5", + section_header="A.", + text=text, + start_pos=0, + end_pos=len(text), + ) + self.assertTrue(chunk.has_sub_chunks()) + for sc in chunk.sub_chunks: + self.assertEqual(sc.parent_chunk_id, "5") + self.assertTrue(sc.sub_chunk_id.startswith("5.")) + + def test_repr_includes_table_flag(self): + chunk = ExhibitChunk( + chunk_id="1", + section_header="TABLE", + text="data", + start_pos=0, + end_pos=4, + is_table=True, + ) + self.assertIn("TABLE", repr(chunk)) + + def test_sub_chunks_lazy_loaded(self): + """Sub-chunks should only be created on first access.""" + chunk = ExhibitChunk( + chunk_id="1", + section_header="", + text="x" * 600, + start_pos=0, + end_pos=600, + ) + self.assertIsNone(chunk._sub_chunks) + _ = chunk.sub_chunks + self.assertIsNotNone(chunk._sub_chunks) + + +# ============================================================================ +# Tests for _keyword_match +# ============================================================================ + + +class TestKeywordMatch(unittest.TestCase): + + def test_match_dollar_sign(self): + self.assertTrue(_keyword_match("Total cost is $500", ["$"])) + + def test_match_percentage(self): + self.assertTrue(_keyword_match("80% of charges", ["%"])) + + def test_match_case_insensitive(self): + self.assertTrue(_keyword_match("LESSER OF rates", ["lesser of"])) + + def test_no_match(self): + self.assertFalse(_keyword_match("Simple text here", ["$", "%", "lesser of"])) + + def test_empty_keywords(self): + self.assertFalse(_keyword_match("Some text", [])) + + +# ============================================================================ +# Tests for _find_table_regions +# ============================================================================ + + +class TestFindTableRegions(unittest.TestCase): + + def test_no_tables(self): + self.assertEqual(_find_table_regions("No tables here"), []) + + def test_single_table(self): + text = f"Before\n{TABLE_START_MARKER}\ndata\n{TABLE_END_MARKER}\nAfter" + regions = _find_table_regions(text) + self.assertEqual(len(regions), 1) + start, end = regions[0] + self.assertIn(TABLE_START_MARKER, text[start:end]) + self.assertIn(TABLE_END_MARKER, text[start:end]) + + def test_multiple_tables(self): + text = ( + f"Before\n{TABLE_START_MARKER}\ntable1\n{TABLE_END_MARKER}\n" + f"Middle\n{TABLE_START_MARKER}\ntable2\n{TABLE_END_MARKER}\nAfter" + ) + regions = _find_table_regions(text) + self.assertEqual(len(regions), 2) + + def test_unclosed_table(self): + text = f"Before\n{TABLE_START_MARKER}\ndata without end" + regions = _find_table_regions(text) + self.assertEqual(len(regions), 1) + self.assertEqual(regions[0][1], len(text)) + + +# ============================================================================ +# Tests for _chunk_non_table_text +# ============================================================================ + + +class TestChunkNonTableText(unittest.TestCase): + + def test_empty_text(self): + self.assertEqual(_chunk_non_table_text(""), []) + self.assertEqual(_chunk_non_table_text(" "), []) + + def test_text_without_headers_single_chunk(self): + text = "This is a plain paragraph without any section headers." + chunks = _chunk_non_table_text(text) + self.assertEqual(len(chunks), 1) + self.assertEqual(chunks[0].section_header, "") + + def test_numbered_sections(self): + text = "1.1 First section content here.\n1.2 Second section content here." + chunks = _chunk_non_table_text(text) + self.assertGreaterEqual(len(chunks), 2) + headers = [c.section_header for c in chunks] + self.assertIn("1.1", headers) + self.assertIn("1.2", headers) + + def test_article_headers(self): + text = "ARTICLE ONE Some preamble text.\nARTICLE TWO More content." + chunks = _chunk_non_table_text(text) + self.assertGreaterEqual(len(chunks), 1) + has_article = any("ARTICLE" in c.section_header for c in chunks) + self.assertTrue(has_article) + + def test_offset_applied(self): + text = "1.1 Section content." + chunks = _chunk_non_table_text(text, offset=100) + self.assertGreaterEqual(chunks[0].start_pos, 100) + + def test_preamble_before_first_header(self): + text = "This is a long preamble text before any header.\n1.1 First section." + chunks = _chunk_non_table_text(text) + # Should have preamble chunk + section chunk + self.assertGreaterEqual(len(chunks), 2) + + +# ============================================================================ +# Tests for _chunk_exhibit_text +# ============================================================================ + + +class TestChunkExhibitText(unittest.TestCase): + + def test_empty_text(self): + self.assertEqual(_chunk_exhibit_text(""), []) + + def test_text_only_no_tables(self): + text = "1.1 Section one content.\n1.2 Section two content." + chunks = _chunk_exhibit_text(text) + self.assertGreater(len(chunks), 0) + for chunk in chunks: + self.assertNotEqual(chunk.chunk_id, "") # IDs assigned + + def test_table_becomes_separate_chunk(self): + text = ( + "1.1 Some text before table.\n" + f"{TABLE_START_MARKER}\n" + "['Header', 'Value']\n['Row1', '$100']\n" + f"{TABLE_END_MARKER}\n" + "1.2 Text after table." + ) + chunks = _chunk_exhibit_text(text) + table_chunks = [c for c in chunks if c.is_table] + self.assertEqual(len(table_chunks), 1) + self.assertEqual(table_chunks[0].section_header, "TABLE") + + def test_chunk_ids_sequential(self): + text = "1.1 First. 1.2 Second. 1.3 Third." + chunks = _chunk_exhibit_text(text) + ids = [int(c.chunk_id) for c in chunks] + self.assertEqual(ids, list(range(1, len(ids) + 1))) + + def test_merge_header_only_applied(self): + """Chunks where header == text should be merged with next chunk.""" + text = "1.4\nSchedule of Payment. This Agreement shall govern payment terms." + chunks = _chunk_exhibit_text(text) + # "1.4" alone should be merged with the next chunk + for chunk in chunks: + if chunk.section_header == "1.4": + self.assertGreater(len(chunk.text), len("1.4")) + + +# ============================================================================ +# Tests for _merge_header_only_chunks +# ============================================================================ + + +class TestMergeHeaderOnlyChunks(unittest.TestCase): + + def test_empty_list(self): + self.assertEqual(_merge_header_only_chunks([]), []) + + def test_no_header_only_chunks(self): + chunks = [ + ExhibitChunk("1", "1.1", "1.1 Some content", 0, 20), + ExhibitChunk("2", "1.2", "1.2 More content", 20, 40), + ] + result = _merge_header_only_chunks(chunks) + self.assertEqual(len(result), 2) + + def test_merges_header_only_with_next(self): + chunks = [ + ExhibitChunk("1", "1.4", "1.4", 0, 3), + ExhibitChunk("2", "Schedule", "Schedule of Payment. Terms apply.", 4, 40), + ] + result = _merge_header_only_chunks(chunks) + self.assertEqual(len(result), 1) + self.assertIn("1.4", result[0].text) + self.assertIn("Schedule of Payment", result[0].text) + + def test_header_only_at_end_not_merged(self): + chunks = [ + ExhibitChunk("1", "1.1", "1.1 Content", 0, 10), + ExhibitChunk("2", "1.2", "1.2", 10, 13), + ] + result = _merge_header_only_chunks(chunks) + self.assertEqual(len(result), 2) + + +# ============================================================================ +# Tests for _is_section_header_chunk +# ============================================================================ + + +class TestIsSectionHeaderChunk(unittest.TestCase): + + def test_roman_numeral(self): + chunk = ExhibitChunk("1", "", "VI.", 0, 3) + self.assertTrue(_is_section_header_chunk(chunk)) + + def test_roman_numeral_with_title(self): + chunk = ExhibitChunk("1", "", "IV. SERVICES", 0, 12) + self.assertTrue(_is_section_header_chunk(chunk)) + + def test_article_header(self): + chunk = ExhibitChunk("1", "", "ARTICLE ONE", 0, 11) + self.assertTrue(_is_section_header_chunk(chunk)) + + def test_short_all_caps(self): + chunk = ExhibitChunk("1", "", "RELATIONSHIP OF PARTIES", 0, 23) + self.assertTrue(_is_section_header_chunk(chunk)) + + def test_normal_text_not_header(self): + chunk = ExhibitChunk( + "1", "", "This is a regular paragraph with lowercase text.", 0, 48 + ) + self.assertFalse(_is_section_header_chunk(chunk)) + + +# ============================================================================ +# Tests for _merge_small_parent_chunks +# ============================================================================ + + +class TestMergeSmallParentChunks(unittest.TestCase): + + def test_empty_list(self): + self.assertEqual(_merge_small_parent_chunks([]), []) + + def test_single_chunk_returned_as_is(self): + chunk = ExhibitChunk("1", "1.1", "Short", 0, 5) + result = _merge_small_parent_chunks([chunk]) + self.assertEqual(len(result), 1) + + def test_small_chunk_merged_with_next(self): + chunks = [ + ExhibitChunk("1", "VI.", "VI.", 0, 3), # small, header-like + ExhibitChunk("2", "1.1", "A" * 150, 3, 153), # large enough + ] + result = _merge_small_parent_chunks(chunks) + self.assertEqual(len(result), 1) + self.assertIn("VI.", result[0].text) + + def test_table_chunks_never_merged(self): + chunks = [ + ExhibitChunk("1", "TABLE", "table data", 0, 10, is_table=True), + ExhibitChunk("2", "1.1", "A" * 150, 10, 160), + ] + result = _merge_small_parent_chunks(chunks) + table_chunks = [c for c in result if c.is_table] + self.assertEqual(len(table_chunks), 1) + + def test_all_tiny_combined(self): + chunks = [ + ExhibitChunk("1", "A.", "Tiny", 0, 4), + ExhibitChunk("2", "B.", "Also tiny", 5, 14), + ] + result = _merge_small_parent_chunks(chunks) + non_table = [c for c in result if not c.is_table] + self.assertEqual(len(non_table), 1) + + +# ============================================================================ +# Tests for cosine_similarity +# ============================================================================ + + +class TestCosineSimilarity(unittest.TestCase): + + def test_identical_vectors(self): + self.assertAlmostEqual(cosine_similarity([1, 0, 0], [1, 0, 0]), 1.0) + + def test_orthogonal_vectors(self): + self.assertAlmostEqual(cosine_similarity([1, 0], [0, 1]), 0.0) + + def test_opposite_vectors(self): + self.assertAlmostEqual(cosine_similarity([1, 0], [-1, 0]), -1.0) + + +# ============================================================================ +# Tests for _find_header_in_text +# ============================================================================ + + +class TestFindHeaderInText(unittest.TestCase): + + def test_exact_match(self): + text = "Some text\nARTICLE ONE - SERVICES\nMore text" + pos = _find_header_in_text(text, "ARTICLE ONE - SERVICES") + self.assertGreater(pos, -1) + self.assertEqual( + text[pos : pos + len("ARTICLE ONE - SERVICES")], "ARTICLE ONE - SERVICES" + ) + + def test_case_insensitive_match(self): + text = "Some text\narticle one\nMore text" + pos = _find_header_in_text(text, "ARTICLE ONE") + self.assertGreater(pos, -1) + + def test_not_found(self): + pos = _find_header_in_text("Some text", "NONEXISTENT HEADER") + self.assertEqual(pos, -1) + + def test_empty_inputs(self): + self.assertEqual(_find_header_in_text("", "header"), -1) + self.assertEqual(_find_header_in_text("text", ""), -1) + + +# ============================================================================ +# Tests for create_exhibits_from_header_dict +# ============================================================================ + + +class TestCreateExhibitsFromHeaderDict(unittest.TestCase): + + def test_empty_dict(self): + self.assertEqual(create_exhibits_from_header_dict({}), []) + + def test_single_exhibit(self): + text_dict = { + "1": "EXHIBIT A\nContent of exhibit A.", + "2": "More content page 2.", + } + header_dict = {"1": ["EXHIBIT A"]} + exhibits = create_exhibits_from_header_dict(header_dict, text_dict=text_dict) + self.assertEqual(len(exhibits), 1) + self.assertEqual(exhibits[0].exhibit_header, "EXHIBIT A") + self.assertIn("EXHIBIT A", exhibits[0].exhibit_text) + + def test_multiple_exhibits_different_pages(self): + text_dict = { + "1": "EXHIBIT A\nContent A.", + "2": "Still exhibit A content.", + "3": "EXHIBIT B\nContent B.", + } + header_dict = {"1": ["EXHIBIT A"], "3": ["EXHIBIT B"]} + exhibits = create_exhibits_from_header_dict(header_dict, text_dict=text_dict) + self.assertEqual(len(exhibits), 2) + self.assertEqual(exhibits[0].exhibit_header, "EXHIBIT A") + self.assertEqual(exhibits[1].exhibit_header, "EXHIBIT B") + + def test_multiple_headers_on_same_page(self): + text_dict = { + "1": "ARTICLE ONE Terms here.\nARTICLE TWO More terms here.", + "2": "Continued text.", + } + header_dict = {"1": ["ARTICLE ONE", "ARTICLE TWO"]} + exhibits = create_exhibits_from_header_dict(header_dict, text_dict=text_dict) + self.assertEqual(len(exhibits), 2) + self.assertEqual(exhibits[0].exhibit_header, "ARTICLE ONE") + self.assertEqual(exhibits[1].exhibit_header, "ARTICLE TWO") + + def test_linked_list_chain(self): + text_dict = {"1": "EXHIBIT A text.", "2": "EXHIBIT B text."} + header_dict = {"1": ["EXHIBIT A"], "2": ["EXHIBIT B"]} + exhibits = create_exhibits_from_header_dict(header_dict, text_dict=text_dict) + self.assertIsNone(exhibits[0].prev_exhibit) + self.assertIs(exhibits[1].prev_exhibit, exhibits[0]) + + def test_with_pages_dict(self): + page1 = Page("1", "EXHIBIT A\nContent A.") + page2 = Page("2", "EXHIBIT B\nContent B.") + pages_dict = {"1": page1, "2": page2} + header_dict = {"1": ["EXHIBIT A"], "2": ["EXHIBIT B"]} + exhibits = create_exhibits_from_header_dict(header_dict, pages_dict=pages_dict) + self.assertEqual(len(exhibits), 2) + self.assertIsNotNone(exhibits[0].pages) + + def test_raises_without_source(self): + with self.assertRaises(ValueError): + create_exhibits_from_header_dict({"1": ["H"]}) + + def test_page_char_ranges_set(self): + text_dict = {"1": "EXHIBIT A\nContent.", "2": "Page 2 content."} + header_dict = {"1": ["EXHIBIT A"]} + exhibits = create_exhibits_from_header_dict(header_dict, text_dict=text_dict) + self.assertIsNotNone(exhibits[0]._page_char_ranges) + self.assertGreater(len(exhibits[0]._page_char_ranges), 0) + + +# ============================================================================ +# Tests for Exhibit chunk-related methods +# ============================================================================ + + +class TestExhibitChunkMethods(unittest.TestCase): + + def _make_exhibit_with_chunks(self): + """Helper: create an exhibit with pre-assigned chunks (bypass embedding).""" + exhibit = Exhibit( + exhibit_page="1", + exhibit_page_nums=["1"], + exhibit_header="Test Exhibit", + exhibit_text="1.1 Section one content here.\n1.2 Section two content here.", + ) + # Pre-populate chunks to avoid embedding calls + chunks = [ + ExhibitChunk("1", "1.1", "1.1 Section one content here.", 0, 30), + ExhibitChunk("2", "1.2", "1.2 Section two content here.", 31, 60), + ] + exhibit._chunks = chunks + return exhibit + + def test_get_chunk_by_id(self): + exhibit = self._make_exhibit_with_chunks() + chunk = exhibit.get_chunk_by_id("1") + self.assertIsNotNone(chunk) + self.assertEqual(chunk.section_header, "1.1") + + def test_get_chunk_by_id_not_found(self): + exhibit = self._make_exhibit_with_chunks() + self.assertIsNone(exhibit.get_chunk_by_id("999")) + + def test_get_chunks_by_header_pattern(self): + exhibit = self._make_exhibit_with_chunks() + matches = exhibit.get_chunks_by_header_pattern(r"1\.\d") + self.assertEqual(len(matches), 2) + + def test_group_chunks_by_page(self): + exhibit = Exhibit( + exhibit_page="1", + exhibit_page_nums=["1", "2"], + exhibit_header="Test", + exhibit_text="Page 1 content.\nPage 2 content.", + ) + # Set page char ranges manually + exhibit._page_char_ranges = [("1", 0, 15), ("2", 16, 31)] + chunks = [ + ExhibitChunk("1", "A.", "chunk on page 1", 0, 15), + ExhibitChunk("2", "B.", "chunk on page 2", 16, 31), + ] + result = exhibit.group_chunks_by_page(chunks) + self.assertIn("1", result) + self.assertIn("2", result) + self.assertEqual(len(result["1"]), 1) + self.assertEqual(len(result["2"]), 1) + + def test_group_chunks_by_page_empty(self): + exhibit = self._make_exhibit_with_chunks() + result = exhibit.group_chunks_by_page([]) + self.assertEqual(result, {}) + + @patch("src.pipelines.shared.extraction.exhibit_funcs._compute_chunk_embeddings") + def test_chunks_property_calls_chunking(self, mock_embed): + """The chunks property should parse exhibit text and compute embeddings.""" + mock_embed.side_effect = lambda chunks: chunks # passthrough + exhibit = Exhibit( + exhibit_page="1", + exhibit_page_nums=["1"], + exhibit_header="Test", + exhibit_text="1.1 First section. 1.2 Second section.", + ) + chunks = exhibit.chunks + self.assertIsNotNone(chunks) + mock_embed.assert_called_once() + + +# ============================================================================ +# Original TestExhibitFuncs (pre-existing tests) +# ============================================================================ + + class TestExhibitFuncs(unittest.TestCase): def test_add_reimbursement_rows_sets_flag(self): diff --git a/src/tests/test_prompt_calls.py b/src/tests/test_prompt_calls.py index 9d3da09..07381ec 100644 --- a/src/tests/test_prompt_calls.py +++ b/src/tests/test_prompt_calls.py @@ -720,18 +720,18 @@ class TestPromptCalls(unittest.TestCase): self.assertIn(result, ["YES", "NO"]) @patch("src.utils.llm_utils.invoke_claude") - def test_prompt_exhibit_header_returns_string(self, mock_invoke): - """Test that prompt_exhibit_header returns a string.""" + def test_prompt_exhibit_header_returns_parsed_result(self, mock_invoke): + """Test that prompt_exhibit_header returns parsed JSON from LLM output.""" # Simulate LLM returning header as list mock_invoke.return_value = '["Exhibit A - Services"]' result = prompt_calls.prompt_exhibit_header( - "Page content with Exhibit A header", self.filename + "Page content with Exhibit A header", ["Exhibit"], self.filename ) - # Should return string (extracted from list) - self.assertIsInstance(result, str) - self.assertIn("Exhibit", result) + # universal_json_load parses the JSON string into a Python object + self.assertIsInstance(result, list) + self.assertIn("Exhibit A - Services", result) @patch("src.utils.llm_utils.invoke_claude") def test_prompt_exhibit_title_match_returns_string(self, mock_invoke): diff --git a/src/utils/rag_utils.py b/src/utils/rag_utils.py index b8b1506..92d06bb 100644 --- a/src/utils/rag_utils.py +++ b/src/utils/rag_utils.py @@ -10,6 +10,8 @@ import logging from dataclasses import dataclass from typing import Any, List, Dict +import numpy as np + import boto3 from langchain_classic.retrievers import EnsembleRetriever from langchain_core.documents import Document @@ -167,3 +169,10 @@ def rerank_documents( f"Reranking failed: {e} | Selecting {top_k} from base retrieved docs" ) return docs[:top_k] + + +def cosine_similarity(vec1: List[float], vec2: List[float]) -> float: + """Compute cosine similarity between two vectors.""" + a = np.array(vec1) + b = np.array(vec2) + return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))) diff --git a/src/utils/string_utils.py b/src/utils/string_utils.py index 1e422c9..d11fb10 100644 --- a/src/utils/string_utils.py +++ b/src/utils/string_utils.py @@ -1092,3 +1092,21 @@ def token_signature(text: str) -> tuple: # Sort tokens to ensure order-independent comparison return tuple(sorted(tokens)) + + +def keyword_match(text: str, keywords: list) -> bool: + """ + Check if text contains any of the keywords (case-insensitive). + + Args: + text: Text to search in + keywords: List of keywords to look for + + Returns: + True if any keyword is found in text + """ + text_lower = text.lower() + for keyword in keywords: + if keyword.lower() in text_lower: + return True + return False