From 524968f09f1a7fcb8ef4a9483c8fa501ca129037 Mon Sep 17 00:00:00 2001 From: Alex Galarce Date: Mon, 10 Feb 2025 22:59:10 +0000 Subject: [PATCH] Merged in bugfix/exhibit-page-fix-pipeline (pull request #390) Bugfix/exhibit page fix pipeline * Fix broken get_exhibit_pages tests * Add __repr__ method to FieldSet for better string representation Approved-by: Katon Minhas --- fieldExtraction/src/preprocessing_funcs.py | 32 ++++++++++++------- .../src/prompts/investment_prompts.py | 6 ++++ .../tests/preprocessing_funcs_test.py | 15 +++++---- 3 files changed, 36 insertions(+), 17 deletions(-) diff --git a/fieldExtraction/src/preprocessing_funcs.py b/fieldExtraction/src/preprocessing_funcs.py index 1704fce..1b05009 100644 --- a/fieldExtraction/src/preprocessing_funcs.py +++ b/fieldExtraction/src/preprocessing_funcs.py @@ -203,23 +203,33 @@ def filter_quick_review(text_dict): or "COVER SHEET" in page_text.upper() } -def get_exhibit_pages(text_dict, filename): +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. + + Args: + text_dict (dict[str, str]): Dictionary valued by string-formatted page numbers and valued by page text + filename (str): Filename of the contract (for tracking purposes) + + Returns: + list[str]: A list of exhibit page numbers extracted from the contract. + """ exhibit_pages = [] first_page = True for page_num, page in text_dict.items(): if first_page: # Include the first page in the exhibit pages exhibit_pages.append(page_num) first_page = False - continue - prompt = preprocessing_prompts.EXHIBIT_CHECK(page[0:100]) - claude_answer_raw = llm_utils.invoke_claude( - prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=10 # 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) + else: + prompt = preprocessing_prompts.EXHIBIT_CHECK(page[0:100]) + claude_answer_raw = llm_utils.invoke_claude( + prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=10 # 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) return exhibit_pages def chunk_by_exhibit(text_dict: dict, diff --git a/fieldExtraction/src/prompts/investment_prompts.py b/fieldExtraction/src/prompts/investment_prompts.py index 3dd6267..a31d57c 100644 --- a/fieldExtraction/src/prompts/investment_prompts.py +++ b/fieldExtraction/src/prompts/investment_prompts.py @@ -194,6 +194,12 @@ class FieldSet: for field in self.fields: field.print() + def __repr__(self): + """Returns a string representation of all fields in the FieldSet.""" + if not self.fields: + return "FieldSet: No fields loaded" + return str([field.field_name for field in self.fields]) + def print_fields(self): if not self.fields: print("No fields loaded.") diff --git a/fieldExtraction/tests/preprocessing_funcs_test.py b/fieldExtraction/tests/preprocessing_funcs_test.py index a65353c..f86f177 100644 --- a/fieldExtraction/tests/preprocessing_funcs_test.py +++ b/fieldExtraction/tests/preprocessing_funcs_test.py @@ -111,14 +111,16 @@ class TestPreprocessingFuncs: # Test cases for get_exhibit_pages @pytest.mark.parametrize("mock_response, expected_pages, text_dict", [ - # Exhibit page detected - (["|Y|","|N|"], ["1"], {"1": "Exhibit content", "2": "Non-exhibit content"}), + # 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 - (["Invalid"], [], {"1": "Non-exhibit content"}), + # 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 @@ -126,5 +128,6 @@ class TestPreprocessingFuncs: filename = "test.pdf" assert get_exhibit_pages(text_dict, filename) == expected_pages - # # verify calls made - # assert mocker_invoke_claude.call_count == len(text_dict) + # 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