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
This commit is contained in:
Karan Desai
committed by
Katon Minhas
Katon Minhas
parent
eda5af10f1
commit
786bab6118
@@ -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"
|
||||
|
||||
@@ -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--------"
|
||||
|
||||
@@ -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]+)",
|
||||
]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)))
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user