Merged in bugfix/identify-tables (pull request #544)
Don't consider table-subpages each individually for reimbursement exhibit identification * Refine EXHIBIT_CHECK prompt * test: rework `get_exhibit_pages` for table subpages * Fix indentation * Experimental table de-chunking * Increase max_tokens for Claude model in exhibit page processing * fix get_pages_to_process docstring * refactor: improved docstring for get_exhibit_pages * move `get_pages_to_process` to `preprocess.py` * test: add unit tests for get_pages_to_process function * Comment out debugging lines * Changed prompt that was breading universal_json_load parsing * fix: handle "no results" case in get_reimbursement_primary function * fixed docstring Approved-by: Katon Minhas
This commit is contained in:
Alex Galarce
@@ -57,6 +57,8 @@ def process_file(file_object, all_dataset, run_timestamp):
|
||||
text_dict, top_sheet_dict = preprocess.split_text(contract_text)
|
||||
text_dict = preprocess.clean_tables(text_dict, filename)
|
||||
exhibit_pages, exhibit_chunk_mapping = preprocess.one_to_n_exhibit_chunking(text_dict, filename)
|
||||
# print(f"Exhibit Pages: {exhibit_pages}") # Debugging line
|
||||
# print(f"Exhibit Chunk Mapping: {exhibit_chunk_mapping}") # Debugging line
|
||||
print(f"{datetime_str()} Preprocessing Complete - {filename}")
|
||||
|
||||
################## ONE TO N ##################
|
||||
@@ -147,27 +149,41 @@ def run_one_to_n_prompts(filename, text_dict, exhibit_pages, exhibit_chunk_mappi
|
||||
# Then identify the compensation exhibits
|
||||
reimbursement_exhibits = one_to_n_funcs.identify_reimbursement_exhibits(all_exhibit_headers, filename)
|
||||
|
||||
pages_to_process = preprocess.get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
# print(f"Pages to process: {pages_to_process}") # Debugging line
|
||||
|
||||
################## RUN PROMPTS ##################
|
||||
one_to_n_results = []
|
||||
for exhibit_page in reimbursement_exhibits:
|
||||
exhibit_chunk = string_utils.get_exhibit_chunk(text_dict, exhibit_chunk_mapping, exhibit_page)
|
||||
for page in pages_to_process:
|
||||
# Use individual page text for subpages, full chunk for regular pages
|
||||
if '.' in page:
|
||||
# Subpage - get the text for this specific page
|
||||
page_text = text_dict[page]
|
||||
processing_id = page
|
||||
else:
|
||||
# Regular page - get the full chunk of text
|
||||
page_text = string_utils.get_exhibit_chunk(text_dict, exhibit_chunk_mapping, page)
|
||||
processing_id = page
|
||||
|
||||
# Find which exhibit this page belongs to (for tracking)
|
||||
parent_exhibit = exhibit_chunk_mapping.get(page, page)
|
||||
|
||||
################## INITIALIZE FIELDS ##################
|
||||
reimbursement_level_fields = FieldSet(relationship="one_to_n", field_type="reimbursement_level", file_path=FIELD_JSON_PATH)
|
||||
|
||||
################## GET EXHIBIT HEADER ##################
|
||||
exhibit_header = one_to_n_funcs.get_exhibit_header(exhibit_chunk, filename)
|
||||
exhibit_header = one_to_n_funcs.get_exhibit_header(page_text, filename)
|
||||
|
||||
################## GET REIMBURSEMENT TIN/NPI ##################
|
||||
tin_npi_answers, reimbursement_level_fields = reimbursement_tin_npi(exhibit_chunk, reimbursement_level_fields, filename)
|
||||
tin_npi_answers, reimbursement_level_fields = reimbursement_tin_npi(page_text, reimbursement_level_fields, filename)
|
||||
|
||||
################## GET EXHIBIT-LEVEL ANSWERS ##################
|
||||
exhibit_level_answers = one_to_n_funcs.get_exhibit_level_answers(exhibit_chunk, filename)
|
||||
exhibit_level_answers['EXHIBIT_PAGE'] = exhibit_page
|
||||
exhibit_level_answers = one_to_n_funcs.get_exhibit_level_answers(page_text, filename)
|
||||
exhibit_level_answers['EXHIBIT_PAGE'] = parent_exhibit
|
||||
exhibit_level_answers['EXHIBIT_TITLE'] = exhibit_header
|
||||
|
||||
################## GET DYNAMIC-PRIMARY ANSWERS ##################
|
||||
dynamic_to_exhibit_level_answers, dynamic_to_reimbursement_level_fields = dynamic_funcs.get_dynamic_answers(exhibit_chunk,
|
||||
dynamic_to_exhibit_level_answers, dynamic_to_reimbursement_level_fields = dynamic_funcs.get_dynamic_answers(page_text,
|
||||
exhibit_header,
|
||||
filename,
|
||||
FieldSet(relationship="one_to_n", field_type="dynamic", file_path=config.FIELD_JSON_PATH))
|
||||
@@ -175,7 +191,7 @@ def run_one_to_n_prompts(filename, text_dict, exhibit_pages, exhibit_chunk_mappi
|
||||
reimbursement_level_fields.combine(dynamic_to_reimbursement_level_fields, inplace=True)
|
||||
|
||||
################## GET DYNAMIC-CODE ANSWERS ##################
|
||||
code_to_exhibit_level_answers, code_to_reimbursement_level_fields = dynamic_funcs.get_dynamic_answers(exhibit_chunk,
|
||||
code_to_exhibit_level_answers, code_to_reimbursement_level_fields = dynamic_funcs.get_dynamic_answers(page_text,
|
||||
exhibit_header,
|
||||
filename,
|
||||
FieldSet(relationship="one_to_n", field_type="dynamic_code", file_path=config.FIELD_JSON_PATH))
|
||||
@@ -183,7 +199,7 @@ def run_one_to_n_prompts(filename, text_dict, exhibit_pages, exhibit_chunk_mappi
|
||||
reimbursement_level_fields.combine(code_to_reimbursement_level_fields, inplace=True)
|
||||
|
||||
################## GET REIMBURSEMENT-LEVEL ANSWERS (INCLUDING DYNAMIC) ##################
|
||||
reimbursement_level_answers = reimbursement_level(exhibit_chunk, filename, reimbursement_level_fields, all_dataset, exhibit_page) # Return list of dictionaries
|
||||
reimbursement_level_answers = reimbursement_level(page_text, filename, reimbursement_level_fields, all_dataset, parent_exhibit) # Return list of dictionaries
|
||||
|
||||
################# COMBINE ANSWERS ##################
|
||||
full_answer_dict = combine_one_to_n_answers(exhibit_level_answers, reimbursement_level_answers, tin_npi_answers)
|
||||
|
||||
@@ -104,11 +104,22 @@ def identify_reimbursement_exhibits(exhibit_headers: dict, filename: str) -> lis
|
||||
def get_reimbursement_primary(reimbursement_level_fields, exhibit_text, filename):
|
||||
field_prompts = reimbursement_level_fields.get_prompt_dict()
|
||||
prompt = investment_prompts.REIMBURSEMENT_LEVEL_PRIMARY(exhibit_text, field_prompts)
|
||||
# print(prompt) # Debugging line
|
||||
llm_answer_raw = llm_utils.invoke_claude(
|
||||
prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, max_tokens=25000
|
||||
)
|
||||
|
||||
llm_answer_final = string_utils.universal_json_load(llm_answer_raw)
|
||||
# print(llm_answer_raw) # Debugging line
|
||||
|
||||
# Check for the special "no results" case
|
||||
if "NO_REIMBURSEMENT_TERMS_FOUND" in llm_answer_raw:
|
||||
return [] # Return empty list if no reimbursement terms found
|
||||
|
||||
try:
|
||||
llm_answer_final = string_utils.universal_json_load(llm_answer_raw)
|
||||
except ValueError as e:
|
||||
print(f"Error parsing LLM response: {e}")
|
||||
print(f"Raw LLM output: {llm_answer_raw}")
|
||||
raise
|
||||
return llm_answer_final
|
||||
|
||||
|
||||
|
||||
@@ -65,9 +65,55 @@ def one_to_n_exhibit_chunking(text_dict, filename):
|
||||
exhibit_pages = preprocessing_funcs.get_exhibit_pages(
|
||||
text_dict, filename
|
||||
) # All pages with exhibit headers
|
||||
# print(f"Exhibit Pages (inside one_to_n_exhibit_chunking): {exhibit_pages}") # Debugging line
|
||||
exhibit_chunk_mapping = preprocessing_funcs.chunk_by_exhibit(text_dict, exhibit_pages)
|
||||
return exhibit_pages, exhibit_chunk_mapping
|
||||
|
||||
def get_pages_to_process(reimbursement_exhibits: list, exhibit_chunk_mapping: dict) -> list:
|
||||
"""Determine how to process reimbursement exhibit pages based on table splitting.
|
||||
|
||||
This function decides whether to process exhibit pages individually (for subpages)
|
||||
or as chunks (for regular pages). When tables are split into subpages (e.g., 5.1, 5.2),
|
||||
each subpage needs individual processing. When pages are regular (e.g., 1, 2, 3),
|
||||
They can be processed as a single chunk for better context.
|
||||
|
||||
Processing Logic:
|
||||
- If ANY page in an exhibit chunk contains a decimal (subpage), process ALL pages individually
|
||||
- If NO pages contain decimals, process as a single chunk using the exhibit page number
|
||||
|
||||
Args:
|
||||
reimbursement_exhibits (list): List of exhibit page numbers that contain reimbursement information.
|
||||
exhibit_chunk_mapping (dict): Maps each page to its parent exhibit page.
|
||||
|
||||
Returns:
|
||||
list: List of page numbers to process. For exhibits with subpages, returns all subpages individually.
|
||||
For regular exhibits, returns the exhibit page number for chunk processing.
|
||||
|
||||
Example:
|
||||
>>> reimbursement_exhibits = ["1", "5"]
|
||||
>>> exhibit_chunk_mapping = {"1": "1", "2": "1", "5.0": "5.0", "5.1": "5.0", "5.2": "5.0"}
|
||||
Returns: ["1", "5.0", "5.1", "5.2"]
|
||||
"""
|
||||
pages_to_process = []
|
||||
|
||||
for exhibit_page in reimbursement_exhibits:
|
||||
# Find all pages mapped to this exhibit
|
||||
exhibit_pages_in_chunk = [page for page, mapped_exhibit in exhibit_chunk_mapping.items()
|
||||
if mapped_exhibit == exhibit_page]
|
||||
|
||||
# Check if any pages are subpages (contain '.')
|
||||
has_subpages = any('.' in page for page in exhibit_pages_in_chunk)
|
||||
|
||||
if has_subpages:
|
||||
# Split into individual subpages
|
||||
print(f"Exhibit {exhibit_page} has subpages - processing individually: {exhibit_pages_in_chunk}")
|
||||
pages_to_process.extend(exhibit_pages_in_chunk)
|
||||
else:
|
||||
# Keep as one chunk (use the exhibit_page as representative)
|
||||
print(f"Exhibit {exhibit_page} is regular - processing as chunk")
|
||||
pages_to_process.append(exhibit_page)
|
||||
|
||||
return pages_to_process
|
||||
|
||||
def one_to_one_smart_chunking(text_dict, contract_text, one_to_one_fields):
|
||||
"""
|
||||
|
||||
@@ -211,32 +211,84 @@ def filter_quick_review(text_dict):
|
||||
}
|
||||
|
||||
def get_exhibit_pages(text_dict: dict[str, str], filename: str) -> list[str]:
|
||||
"""Extract beginning-of-exhibit pages from a contract. The first page is also
|
||||
always included in the exhibit pages.
|
||||
"""Extract beginning-of-exhibit pages from a contract using LLM-based detection.
|
||||
|
||||
This function identifies pages that start new exhibits by checking each page (or
|
||||
subpage) for exhibit headers using an LLM prompt. It handles both regular pages
|
||||
(e.g., "5") and subpages created by table splitting (e.g., "5.1", "5.2").
|
||||
|
||||
Key behaviors:
|
||||
- The first page of the contract is always included in the exhibit pages.
|
||||
- For base pages with subpages, if ANY subpage contains an exhibit header,
|
||||
the FIRST subpage is added to the exhibit list
|
||||
- Each base page number is only processed once to avoid duplicates.
|
||||
|
||||
Args:
|
||||
text_dict (dict[str, str]): Dictionary valued by string-formatted page numbers and valued by page text
|
||||
text_dict (dict[str, str]): Dictionary keyed by string-formatted page numbers and valued by page text
|
||||
Page numbers may include subpages (e.g., "5.1", "5.2")
|
||||
filename (str): Filename of the contract (for tracking purposes)
|
||||
|
||||
Returns:
|
||||
list[str]: A list of exhibit page numbers extracted from the contract.
|
||||
list[str]: Ordered list of page numbers that start new exhibits. May include
|
||||
subpage numbers (e.g., ["1", "5.1", "5.2"]) depending on table splitting.
|
||||
|
||||
Example:
|
||||
Given pages ["1", "2", "5.0", "5.1", "5.2", "7"] where "5.1" contains an exhibit header:
|
||||
Returns ["1", "5.1", "7"] (assuming pages 1 and 7 also have exhibit headers)
|
||||
"""
|
||||
exhibit_pages = []
|
||||
first_page = True
|
||||
for page_num, page in text_dict.items():
|
||||
if first_page: # Include the first page in the exhibit pages
|
||||
processed_base_pages = set() # Track base pages to avoid duplicate processing
|
||||
|
||||
# Sort pages to handle subpages properly (5.0, 5.1, 5.2, etc.)
|
||||
# TODO: augment `string_utils.page_key_sort` to handle subpages
|
||||
sorted_pages = sorted(text_dict.keys(), key=lambda x: (
|
||||
int(x.split('.')[0]),
|
||||
float(x.split('.')[1]) if '.' in x else 0)
|
||||
)
|
||||
|
||||
for page_num in sorted_pages:
|
||||
base_page = page_num.split(".")[0] # Extract base page number (e.g., "5" from "5.1")
|
||||
|
||||
# Always include the first page as an exhibit
|
||||
if first_page:
|
||||
exhibit_pages.append(page_num)
|
||||
processed_base_pages.add(base_page)
|
||||
first_page = False
|
||||
# Skip if page base group has already been processed
|
||||
elif base_page in processed_base_pages:
|
||||
continue
|
||||
|
||||
# For new base pages, check all subpages for exhibit headers
|
||||
# This handles cases where tables are split across subpages
|
||||
else:
|
||||
prompt = investment_prompts.EXHIBIT_CHECK(page[0:200])
|
||||
claude_answer_raw = llm_utils.invoke_claude(
|
||||
prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=100 # TODO: low priority, try increasing max_tokens and maybe pass multiple pages in to reduce overall calls
|
||||
)
|
||||
claude_answer_extracted = string_utils.extract_text_from_delimiters(
|
||||
claude_answer_raw, Delimiter.PIPE
|
||||
)
|
||||
if "Y" in claude_answer_extracted:
|
||||
exhibit_pages.append(page_num)
|
||||
# Check if ANY subpage of the base page is an exhibit
|
||||
subpages = [p for p in text_dict.keys() if p.startswith(base_page + ".") or p == base_page]
|
||||
|
||||
is_exhibit = False
|
||||
for subpage in subpages:
|
||||
page_content = text_dict[subpage]
|
||||
prompt = investment_prompts.EXHIBIT_CHECK(page_content[0:200])
|
||||
# print(f"Prompting LLM for exhibit check on page {subpage}...") # Debugging line
|
||||
# print("Prompt:", prompt) # Debugging line
|
||||
claude_answer_raw = llm_utils.invoke_claude(
|
||||
prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=150
|
||||
)
|
||||
# print("Claude answer:", claude_answer_raw) # Debugging line
|
||||
claude_answer_extracted = string_utils.extract_text_from_delimiters(
|
||||
claude_answer_raw, Delimiter.PIPE
|
||||
)
|
||||
if "Y" in claude_answer_extracted:
|
||||
is_exhibit = True
|
||||
break # We found an exhibit, no need to check further subpages
|
||||
|
||||
# If any subpage is an exhibit, add the first subpage as the exhibit start
|
||||
if is_exhibit:
|
||||
# Find the numerically first subpage (handles both "5" and "5.1" formats)
|
||||
first_subpage = min(subpages, key=lambda x: float(x.split('.')[-1]) if '.' in x else 0)
|
||||
exhibit_pages.append(first_subpage)
|
||||
processed_base_pages.add(base_page)
|
||||
|
||||
return exhibit_pages
|
||||
|
||||
def chunk_by_exhibit(text_dict: dict,
|
||||
|
||||
@@ -257,32 +257,29 @@ class FieldSet:
|
||||
#####################################################################################
|
||||
|
||||
def EXHIBIT_CHECK(page_start):
|
||||
return f"""Analyze the following contract excerpt and determine if it contains an exhibit/section header.
|
||||
|
||||
return f"""Analyze the following contract excerpt and determine if it indicates the start of a section. Follow the following instructions:
|
||||
Look for these patterns at the START of the text:
|
||||
- EXHIBIT [letter/number]
|
||||
- ATTACHMENT [letter/number]
|
||||
- SCHEDULE [letter/number]
|
||||
- ARTICLE [letter/number]
|
||||
- AMENDMENT [letter/number]
|
||||
- ADDENDUM [letter/number]
|
||||
|
||||
1. Look specifically for the words 'Exhibit', 'Article', 'Amendment', 'Schedule', 'Attachment', or 'Addendum'. These words indicate the start of a new section and are usually followed by a number or letter (e.g., 'Attachment A').
|
||||
The header should be the FIRST substantive content, regardless of what follows.
|
||||
|
||||
2. These words should appear in a header-style format. This means:
|
||||
- On their own line
|
||||
- Surrounded by whitespace or formatting like a title
|
||||
- Capitalized and not part of a longer sentence
|
||||
IGNORE:
|
||||
- References within sentences (e.g., "see Exhibit A")
|
||||
- Legal citations (e.g., "Section 85.113")
|
||||
- Numbered clauses (e.g., "1.20", "7.1")
|
||||
|
||||
3. Do not treat the following as a section start:
|
||||
- Legal references such as 'Section 85.113' or 'Article IX'
|
||||
- Numbered clauses like '1.20', '7.1', or '11.1'
|
||||
- Any sentence that simply mentions one of the keywords above
|
||||
Return |Y| if there's a header at the start, |N| if not.
|
||||
|
||||
4. The phrase 'Additional Provisions' does not indicate the start of a section.
|
||||
|
||||
5. Return Y if the excerpt indicates the start of a section, N if not. Please justify your answer, but enclose your final answer in |pipes|.
|
||||
|
||||
Here is the text to analyze:
|
||||
|
||||
## START CONTRACT TEXT ##
|
||||
Text to analyze:
|
||||
{page_start}
|
||||
## END CONTRACT TEXT ##
|
||||
|
||||
State your answer according to the instructions above."""
|
||||
Answer: |Your decision|"""
|
||||
|
||||
def EXHIBIT_HEADER(context):
|
||||
return f"""Analyze the following contract excerpt and extract the full header found.
|
||||
@@ -367,7 +364,7 @@ The output from this prompt will be used to populate downstream fields in the sy
|
||||
[A. OUTPUT FORMAT]
|
||||
Return a JSON array of objects, each containing these attributes:
|
||||
{fields}
|
||||
Return an empty list if no reimbursement terms are found.
|
||||
If no reimbursement terms are found, return the text "NO_REIMBURSEMENT_TERMS_FOUND" instead of an empty JSON array.
|
||||
|
||||
[B. KEY EXTRACTION RULES]
|
||||
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
import pytest
|
||||
from src.investment.preprocess import get_pages_to_process
|
||||
|
||||
|
||||
class TestGetPagesToProcess:
|
||||
"""Test cases for get_pages_to_process function."""
|
||||
|
||||
def test_regular_pages_only(self):
|
||||
"""Test processing exhibits with no subpages - should return exhibit pages for chunk processing."""
|
||||
reimbursement_exhibits = ["1", "5"]
|
||||
exhibit_chunk_mapping = {
|
||||
"1": "1",
|
||||
"2": "1",
|
||||
"3": "1",
|
||||
"5": "5",
|
||||
"6": "5",
|
||||
"7": "7"
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
assert result == ["1", "5"]
|
||||
|
||||
def test_subpages_only(self):
|
||||
"""Test processing exhibits with subpages - should return all subpages individually."""
|
||||
reimbursement_exhibits = ["5"]
|
||||
exhibit_chunk_mapping = {
|
||||
"1": "1",
|
||||
"5.1": "5",
|
||||
"5.2": "5",
|
||||
"5.3": "5",
|
||||
"7": "7"
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
assert set(result) == {"5.1", "5.2", "5.3"}
|
||||
|
||||
def test_mixed_regular_and_subpages(self):
|
||||
"""Test processing mix of regular exhibits and exhibits with subpages."""
|
||||
reimbursement_exhibits = ["1", "5", "8"]
|
||||
exhibit_chunk_mapping = {
|
||||
"1": "1", # Regular exhibit (no subpages)
|
||||
"2": "1",
|
||||
"5.1": "5", # Exhibit with subpages
|
||||
"5.2": "5",
|
||||
"5.3": "5",
|
||||
"8": "8", # Regular exhibit
|
||||
"9": "8",
|
||||
"10": "8"
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
expected = ["1", "5.1", "5.2", "5.3", "8"]
|
||||
assert set(result) == set(expected)
|
||||
|
||||
def test_exhibit_with_base_page_and_subpages(self):
|
||||
"""Test exhibit that has both base page and subpages (e.g., '5', '5.1', '5.2')."""
|
||||
reimbursement_exhibits = ["5"]
|
||||
exhibit_chunk_mapping = {
|
||||
"1": "1",
|
||||
"5": "5", # Base page
|
||||
"5.1": "5", # Subpages
|
||||
"5.2": "5",
|
||||
"7": "7"
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
# Should return all pages in the exhibit since it has subpages
|
||||
assert set(result) == {"5", "5.1", "5.2"}
|
||||
|
||||
def test_empty_reimbursement_exhibits(self):
|
||||
"""Test handling of empty reimbursement exhibits list."""
|
||||
reimbursement_exhibits = []
|
||||
exhibit_chunk_mapping = {
|
||||
"1": "1",
|
||||
"2": "1",
|
||||
"5.1": "5",
|
||||
"5.2": "5"
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
assert result == []
|
||||
|
||||
def test_complex_subpage_numbering(self):
|
||||
"""Test complex subpage numbering patterns."""
|
||||
reimbursement_exhibits = ["5", "10"]
|
||||
exhibit_chunk_mapping = {
|
||||
"1": "1",
|
||||
"5.0": "5", # Different subpage patterns
|
||||
"5.1": "5",
|
||||
"5.10": "5", # Double digit subpage
|
||||
"5.2": "5",
|
||||
"10.1": "10",
|
||||
"10.11": "10", # Double digit subpage
|
||||
"10.2": "10"
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
expected_5 = {"5.0", "5.1", "5.10", "5.2"}
|
||||
expected_10 = {"10.1", "10.11", "10.2"}
|
||||
|
||||
assert set(result) == expected_5.union(expected_10)
|
||||
|
||||
def test_single_exhibit_single_page(self):
|
||||
"""Test single exhibit with single page."""
|
||||
reimbursement_exhibits = ["1"]
|
||||
exhibit_chunk_mapping = {
|
||||
"1": "1"
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
assert result == ["1"]
|
||||
|
||||
def test_decimal_in_exhibit_name_not_subpage(self):
|
||||
"""Test edge case where exhibit name contains decimal but isn't a subpage."""
|
||||
# This might be an edge case depending on your data format
|
||||
reimbursement_exhibits = ["1.5"] # Exhibit named "1.5"
|
||||
exhibit_chunk_mapping = {
|
||||
"1.5": "1.5", # This is the exhibit itself, not a subpage
|
||||
"2": "1.5" # Regular page belonging to exhibit "1.5"
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
# Since "1.5" contains a decimal, it should process pages individually
|
||||
assert set(result) == {"1.5", "2"}
|
||||
|
||||
def test_order_preservation(self):
|
||||
"""Test that the order of processing is preserved/predictable."""
|
||||
reimbursement_exhibits = ["1", "5", "8"]
|
||||
exhibit_chunk_mapping = {
|
||||
"1": "1",
|
||||
"5.2": "5",
|
||||
"5.1": "5", # Intentionally out of order
|
||||
"8": "8"
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
# Should maintain order of exhibits in reimbursement_exhibits
|
||||
assert result[0] == "1" # First exhibit
|
||||
assert "5.2" in result and "5.1" in result # Subpages from second exhibit
|
||||
assert result[-1] == "8" # Last exhibit
|
||||
|
||||
def test_multiple_exhibits_with_subpages(self):
|
||||
"""Test multiple exhibits that all have subpages."""
|
||||
reimbursement_exhibits = ["3", "7", "12"]
|
||||
exhibit_chunk_mapping = {
|
||||
"1": "1",
|
||||
"3.1": "3",
|
||||
"3.2": "3",
|
||||
"7.1": "7",
|
||||
"7.2": "7",
|
||||
"7.3": "7",
|
||||
"12.1": "12",
|
||||
"12.2": "12"
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
expected = ["3.1", "3.2", "7.1", "7.2", "7.3", "12.1", "12.2"]
|
||||
assert set(result) == set(expected)
|
||||
|
||||
|
||||
# Integration test with actual preprocessing flow
|
||||
class TestGetPagesToProcessIntegration:
|
||||
"""Integration tests that simulate real preprocessing scenarios."""
|
||||
|
||||
def test_table_splitting_scenario(self):
|
||||
"""Test scenario where tables are split creating subpages."""
|
||||
# Simulate a scenario where exhibit 5 had a large table that got split
|
||||
reimbursement_exhibits = ["1", "5"]
|
||||
exhibit_chunk_mapping = {
|
||||
"1": "1", # Regular exhibit
|
||||
"2": "1",
|
||||
"5": "5", # Original exhibit page
|
||||
"5.1": "5", # Table continuation pages
|
||||
"5.2": "5",
|
||||
"5.3": "5",
|
||||
"6": "6" # Next exhibit (not reimbursement)
|
||||
}
|
||||
|
||||
result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
|
||||
|
||||
# Exhibit 1: regular processing (chunk)
|
||||
# Exhibit 5: individual processing (has subpages from table splitting)
|
||||
expected = ["1", "5", "5.1", "5.2", "5.3"]
|
||||
assert set(result) == set(expected)
|
||||
Reference in New Issue
Block a user