Alex Galarce
f84023d742
Bugfix/table preprocessing * restructure get_exhibit_dict * exhibit header * update exhibit header handling * update docstring * bugfix * remove test - update docstring * Merged main into bugfix/table-preprocessing * refactor: update page_key_sort to return float for numeric keys * refactor: enhance get_exhibit_dict to improve chunking logic for exhibits * refactor: update page_key_sort to return float for numeric keys * test: add unit tests for get_exhibit_dict function * test: fix test case 7 * Add final, difficult test * refactor: enhance get_exhibit_dict to handle single decimal page exhibits and improve chunking logic * pulled `one_to_n_test.py` from dynamic primary branch * fix: update common_file_names logic to remove specific entry and log changes * docs: update get_exhibit_dict docstring to clarify input and output formats * test: add test for exhibit handling with decimal pages in get_exhibit_dict * feat: add debug logging for exhibit chunk mapping and exhibit dictionary keys in one_to_n_exhibit_chunking * Merge remote-tracking branch 'origin/main' into bugfix/table-preprocessing * feat: enhance spatial processing rules and consistency checks for reimbursement terms * feat: update global scan methodology to enhance context extraction for reimbursement constraints * Merge remote-tracking branch 'origin/main' into bugfix/table-preprocessing Approved-by: Katon Minhas
460 lines
17 KiB
Python
460 lines
17 KiB
Python
import pytest
|
|
from src.preprocessing_funcs import split_text
|
|
from src.preprocessing_funcs import (
|
|
remove_page_indicators,
|
|
split_text,
|
|
clean_law_symbols,
|
|
filter_quick_review,
|
|
get_exhibit_pages,
|
|
)
|
|
|
|
from src.client.smart_chunking_funcs import smart_chunk_ac
|
|
|
|
class TestPreprocessingFuncs:
|
|
# Test cases for clean_newlines
|
|
@pytest.mark.parametrize("input_text, expected_output", [
|
|
# Removes page number indicators
|
|
("Page 1 of 10\n\nThis is a test.", " This is a test."),
|
|
# Handles multiple newlines
|
|
("Line 1\n\nLine 2\n\nLine 3", "Line 1\n\nLine 2\n\nLine 3"),
|
|
# Handles no newlines
|
|
("This is a single line.", "This is a single line."),
|
|
])
|
|
def test_clean_newlines(self, input_text, expected_output):
|
|
assert remove_page_indicators(input_text) == expected_output
|
|
|
|
# Test cases for clean_law_symbols
|
|
@pytest.mark.parametrize("input_text, expected_output", [
|
|
# Replaces double dollar signs
|
|
("This is a $$ test.", "This is a $ test."),
|
|
# Replaces U.S.C. symbol
|
|
("U.S.C. $1234", "U.S.C.§1234"),
|
|
# Replaces C.F.R. symbol
|
|
("C.F.R. $1234", "C.F.R.§1234"),
|
|
# Replaces $ followed by three decimal places
|
|
("$123.456", "§123.456"),
|
|
# Handles no replacements
|
|
("This is a normal text.", "This is a normal text."),
|
|
# Handles mixed replacements
|
|
("U.S.C. $1234 and C.F.R. $5678 and $123.456", "U.S.C.§1234 and C.F.R.§5678 and §123.456"),
|
|
])
|
|
def test_clean_law_symbols(self, input_text, expected_output):
|
|
assert clean_law_symbols(input_text) == expected_output
|
|
|
|
# Test cases for smart_chunk_ac
|
|
@pytest.mark.parametrize("methodology, mock_return, chunking_fn_name, expected_chunk, expected_pages", [
|
|
# OR methodology
|
|
("or", [1, 2], "or", "Content1\nContent2", {"test": [1, 2]}),
|
|
# Hierarchy methodology
|
|
("hierarchy", [1, 2], "hierarchical", "Content1\nContent2", {"test": [1, 2]}),
|
|
# AND methodology
|
|
("and", [1], "and", "Content1", {"test": [1]}),
|
|
# Regex methodology
|
|
("regex", [1, 2], "regex", "Content1\nContent2", {"test": [1, 2]}),
|
|
# Include/Exclude methodology
|
|
("include_exclude", [1], "with_include_exclude_keywords", "Content1", {"test": [1]}),
|
|
])
|
|
|
|
def test_smart_chunk_ac(self, mocker, methodology, mock_return, chunking_fn_name, expected_chunk, expected_pages):
|
|
# Mock the chunking functions
|
|
mock_chunk_func = mocker.patch(f"src.client.smart_chunking_funcs.chunk_{chunking_fn_name}", return_value=mock_return)
|
|
mock_keyword_search = mocker.patch("src.client.smart_chunking_funcs.keyword_search", return_value={"test": mock_return})
|
|
|
|
text_dict = {"1": "Content1", "2": "Content2"}
|
|
keyword_mappings = {"group1": {"methodology": methodology, "keywords": ["test"], "case_sensitive": False}}
|
|
|
|
if methodology == "regex": keyword_mappings["group1"]["regex"] = "Content"
|
|
|
|
if methodology == "include_exclude":
|
|
del keyword_mappings["group1"]["keywords"]
|
|
keyword_mappings["group1"]["included_keywords"] = ["Content"]
|
|
keyword_mappings["group1"]["excluded_keywords"] = []
|
|
|
|
result = smart_chunk_ac(text_dict, "Content1\nContent2", keyword_mappings)
|
|
|
|
assert result["group1"] == expected_chunk
|
|
assert result["group1_pages"] == expected_pages
|
|
mock_chunk_func.assert_called_once()
|
|
mock_keyword_search.assert_called_once()
|
|
|
|
# Test cases for filter_quick_review
|
|
@pytest.mark.parametrize("input_dict, expected_contract, expected_quick_review", [
|
|
# Filters quick review pages
|
|
(
|
|
{"1": "QUICK REVIEW", "2": "Normal content", "3": "COVER SHEET"},
|
|
{"2": "Normal content"},
|
|
{"1": "QUICK REVIEW", "3": "COVER SHEET"}
|
|
),
|
|
# No quick review pages
|
|
(
|
|
{"1": "Normal content", "2": "More content"},
|
|
{"1": "Normal content", "2": "More content"},
|
|
{}
|
|
),
|
|
# All quick review pages
|
|
(
|
|
{"1": "QUICK REVIEW", "2": "TOP SHEET"},
|
|
{},
|
|
{"1": "QUICK REVIEW", "2": "TOP SHEET"}
|
|
),
|
|
# Case-insensitive matching
|
|
(
|
|
{"1": "quick review", "2": "Normal content"},
|
|
{"2": "Normal content"},
|
|
{"1": "quick review"}
|
|
),
|
|
])
|
|
def test_filter_quick_review(self, input_dict, expected_contract, expected_quick_review):
|
|
contract, quick_review = filter_quick_review(input_dict)
|
|
assert contract == expected_contract
|
|
assert quick_review == expected_quick_review
|
|
|
|
# Test cases for get_exhibit_pages
|
|
@pytest.mark.parametrize("mock_response, expected_pages, text_dict", [
|
|
# First page included by default, second page not an exhibit
|
|
(["|N|"], ["1"], {"1": "Exhibit content", "2": "Non-exhibit content"}),
|
|
# Multiple pages with mixed exhibit status
|
|
(["|N|", "|Y|"], ["1", "3"], {"1": "First page", "2": "Regular content", "3": "Exhibit B"}),
|
|
# No exhibit page detected
|
|
([], [], {}),
|
|
# Multiple exhibit pages
|
|
(["|Y|","|Y|"], ["1", "2"], {"1": "Exhibit content", "2": "Attachment"}),
|
|
# Invalid response; still include the first page
|
|
(["Invalid"], ["1"], {"1": "Non-exhibit content"}),
|
|
])
|
|
def test_get_exhibit_pages(self, mocker, mock_response, expected_pages, text_dict):
|
|
# Mock the LLM response
|
|
mocker_invoke_claude = mocker.patch("src.utils.llm_utils.invoke_claude", side_effect=mock_response)
|
|
filename = "test.pdf"
|
|
assert get_exhibit_pages(text_dict, filename) == expected_pages
|
|
|
|
# Verify calls made - should be number of pages minus 1 (first page), but never negative
|
|
expected_calls = max(0, len(text_dict) - 1)
|
|
assert mocker_invoke_claude.call_count == expected_calls
|
|
|
|
|
|
from src.preprocessing_funcs import get_exhibit_dict
|
|
|
|
class TestGetExhibitDict:
|
|
|
|
def test_case_1_short_table_first_page(self):
|
|
"""Case 1: Short Table on First Page"""
|
|
text_dict = {
|
|
'5': 'Exhibit 5 header text',
|
|
'5.0': 'Short table content'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'5': '5',
|
|
'5.0': '5'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'5.0': 'Exhibit 5 header text\nShort table content'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_case_2_continued_short_table(self):
|
|
"""Case 2: Continued Short Table on Subsequent Pages"""
|
|
text_dict = {
|
|
'5': 'Exhibit 5 header text',
|
|
'5.0': 'Table start content',
|
|
'5.1': 'Table continuation content'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'5': '5',
|
|
'5.0': '5',
|
|
'5.1': '5'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'5.0': 'Exhibit 5 header text\nTable start content',
|
|
'5.1': 'Exhibit 5 header text\nTable continuation content'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_case_3_long_table_first_page(self):
|
|
"""Case 3: Long Table on First Page (split by row limit)"""
|
|
text_dict = {
|
|
'5': 'Exhibit 5 header text',
|
|
'5.0': 'Long table part 1',
|
|
'5.1': 'Long table part 2',
|
|
'5.2': 'Long table part 3'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'5': '5',
|
|
'5.0': '5',
|
|
'5.1': '5',
|
|
'5.2': '5'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'5.0': 'Exhibit 5 header text\nLong table part 1',
|
|
'5.1': 'Exhibit 5 header text\nLong table part 2',
|
|
'5.2': 'Exhibit 5 header text\nLong table part 3'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_case_4_continued_long_table(self):
|
|
"""Case 4: Continued Long Table on Subsequent Pages"""
|
|
text_dict = {
|
|
'5': 'Exhibit 5 header text',
|
|
'6': 'Page 6 text',
|
|
'6.0': 'Continued table part 1',
|
|
'6.1': 'Continued table part 2',
|
|
'7': 'Page 7 text'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'5': '5',
|
|
'6': '5',
|
|
'6.0': '5',
|
|
'6.1': '5',
|
|
'7': '5'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'5.0': 'Exhibit 5 header text\nPage 6 text\nContinued table part 1\nPage 7 text',
|
|
'5.1': 'Exhibit 5 header text\nPage 6 text\nContinued table part 2\nPage 7 text',
|
|
'5.2': 'Exhibit 5 header text\nPage 6 text\nPage 7 text'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_case_5_multiple_tables_first_page(self):
|
|
"""Case 5: Multiple Tables on First page only"""
|
|
text_dict = {
|
|
'5': 'Exhibit 5 header text',
|
|
'5.0': 'First table content',
|
|
'5.1': 'Second table content',
|
|
'5.2': 'Third table content'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'5': '5',
|
|
'5.0': '5',
|
|
'5.1': '5',
|
|
'5.2': '5'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'5.0': 'Exhibit 5 header text\nFirst table content',
|
|
'5.1': 'Exhibit 5 header text\nSecond table content',
|
|
'5.2': 'Exhibit 5 header text\nThird table content'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_case_6_multiple_tables_with_continuation(self):
|
|
"""Case 6: Multiple Tables on Subsequent Page, with first being continuation"""
|
|
text_dict = {
|
|
'5': 'Exhibit 5 header text',
|
|
'6': 'Page 6 text',
|
|
'6.0': 'Continuation of previous table',
|
|
'6.1': 'New table on same page',
|
|
'6.2': 'Another new table on same page'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'5': '5',
|
|
'6': '5',
|
|
'6.0': '5',
|
|
'6.1': '5',
|
|
'6.2': '5'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'5.0': 'Exhibit 5 header text\nPage 6 text\nContinuation of previous table',
|
|
'5.1': 'Exhibit 5 header text\nPage 6 text\nNew table on same page',
|
|
'5.2': 'Exhibit 5 header text\nPage 6 text\nAnother new table on same page'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_case_7_multiple_tables_across_pages(self):
|
|
"""Case 7: Multiple Tables on First and Subsequent Pages"""
|
|
text_dict = {
|
|
'5': 'Exhibit 5 header text',
|
|
'5.0': 'First table on page 5',
|
|
'5.1': 'Second table on page 5',
|
|
'6': 'Page 6 text',
|
|
'6.0': 'First table on page 6',
|
|
'6.1': 'Second table on page 6',
|
|
'7': 'Page 7 text'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'5': '5',
|
|
'5.0': '5',
|
|
'5.1': '5',
|
|
'6': '5',
|
|
'6.0': '5',
|
|
'6.1': '5',
|
|
'7': '5'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'5.0': 'Exhibit 5 header text\nFirst table on page 5\nPage 6 text\nPage 7 text',
|
|
'5.1': 'Exhibit 5 header text\nSecond table on page 5\nPage 6 text\nPage 7 text',
|
|
'5.2': 'Exhibit 5 header text\nPage 6 text\nFirst table on page 6\nPage 7 text',
|
|
'5.3': 'Exhibit 5 header text\nPage 6 text\nSecond table on page 6\nPage 7 text',
|
|
'5.4': 'Exhibit 5 header text\nPage 6 text\nPage 7 text'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_edge_case_exhibit_is_decimal_page(self):
|
|
"""Edge Case: Exhibit itself is a decimal page"""
|
|
text_dict = {
|
|
'5.1': 'Exhibit starting at decimal page'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'5.1': '5.1'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'5.1': 'Exhibit starting at decimal page'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_edge_case_no_decimal_pages(self):
|
|
"""Edge Case: No decimal pages in exhibit (pure text)"""
|
|
text_dict = {
|
|
'5': 'Exhibit 5 header text',
|
|
'6': 'Page 6 text',
|
|
'7': 'Page 7 text'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'5': '5',
|
|
'6': '5',
|
|
'7': '5'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'5.0': 'Exhibit 5 header text\nPage 6 text\nPage 7 text'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_example_from_docstring_1(self):
|
|
"""Test first example from docstring"""
|
|
text_dict = {
|
|
'1.0': 'Page 1.0 text',
|
|
'2': 'Page 2 text',
|
|
'3.0': 'Page 3.0 text',
|
|
'3.1': 'Page 3.1 text',
|
|
'3.2': 'Page 3.2 text',
|
|
'4.0': 'Page 4.0 text'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'1.0': '1.0',
|
|
'2': '2',
|
|
'3.0': '2',
|
|
'3.1': '2',
|
|
'3.2': '2',
|
|
'4.0': '2'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'1.0': 'Page 1.0 text',
|
|
'2.0': 'Page 2 text\nPage 3.0 text',
|
|
'2.1': 'Page 2 text\nPage 3.1 text',
|
|
'2.2': 'Page 2 text\nPage 3.2 text',
|
|
'2.3': 'Page 2 text\nPage 4.0 text'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_example_from_docstring_2(self):
|
|
"""Test second example from docstring"""
|
|
text_dict = {
|
|
'1.0': 'Page 1.0 text',
|
|
'2': 'Page 2 text',
|
|
'3.0': 'Page 3.0 text',
|
|
'3.1': 'Page 3.1 text',
|
|
'3.2': 'Page 3.2 text',
|
|
'4': 'Page 4 text',
|
|
'5': 'Page 5 text'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'1.0': '1.0',
|
|
'2': '2',
|
|
'3.0': '2',
|
|
'3.1': '2',
|
|
'3.2': '2',
|
|
'4': '2',
|
|
'5': '2'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
'1.0': 'Page 1.0 text',
|
|
'2.0': 'Page 2 text\nPage 3.0 text\nPage 4 text\nPage 5 text',
|
|
'2.1': 'Page 2 text\nPage 3.1 text\nPage 4 text\nPage 5 text',
|
|
'2.2': 'Page 2 text\nPage 3.2 text\nPage 4 text\nPage 5 text',
|
|
'2.3': 'Page 2 text\nPage 4 text\nPage 5 text'
|
|
}
|
|
assert result == expected
|
|
|
|
def test_exhibit_with_decimal_pages(self):
|
|
"""Test exhibit with decimal pages"""
|
|
text_dict = {
|
|
"23.0": "Exhibit page 23 header",
|
|
"24": "Page 24 text"
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
"23.0": "23.0",
|
|
"24": "23.0"
|
|
}
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
"23.0.0": "Exhibit page 23 header\nPage 24 text",
|
|
"23.0.1": "Page 24 text"
|
|
}
|
|
assert result == expected
|
|
|
|
def test_multiple_exhibits_with_tables(self):
|
|
"""Test multiple different exhibits, each with their own base pages and tables"""
|
|
text_dict = {
|
|
'10': 'Exhibit A header',
|
|
'10.0': 'Exhibit A table 1',
|
|
'10.1': 'Exhibit A table 2',
|
|
'11': 'Exhibit A page 2',
|
|
|
|
'20': 'Exhibit B header',
|
|
'21': 'Exhibit B page 2',
|
|
'21.0': 'Exhibit B table 1',
|
|
'22': 'Exhibit B page 3',
|
|
|
|
'30': 'Exhibit C header',
|
|
'31': 'Exhibit C page 2'
|
|
}
|
|
exhibit_chunk_mapping = {
|
|
'10': 'A',
|
|
'10.0': 'A',
|
|
'10.1': 'A',
|
|
'11': 'A',
|
|
|
|
'20': 'B',
|
|
'21': 'B',
|
|
'21.0': 'B',
|
|
'22': 'B',
|
|
|
|
'30': 'C',
|
|
'31': 'C'
|
|
}
|
|
|
|
result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
|
|
expected = {
|
|
# Exhibit A
|
|
'A.0': 'Exhibit A header\nExhibit A table 1\nExhibit A page 2',
|
|
'A.1': 'Exhibit A header\nExhibit A table 2\nExhibit A page 2',
|
|
'A.2': 'Exhibit A header\nExhibit A page 2',
|
|
|
|
# Exhibit B
|
|
'B.0': 'Exhibit B header\nExhibit B page 2\nExhibit B table 1\nExhibit B page 3',
|
|
'B.1': 'Exhibit B header\nExhibit B page 2\nExhibit B page 3',
|
|
|
|
# Exhibit C (no decimal pages)
|
|
'C.0': 'Exhibit C header\nExhibit C page 2'
|
|
}
|
|
assert result == expected |