Merged in feature/exhibit-linking (pull request #564)

Feature/exhibit linking

* Initialize new exhibit header prompt

* Save

* Merge branch 'main' into feature/exhibit-linking

* update prompt

* cleanup return type

* remove test

* Test fixes

* Fix test

* Add tests

* Merge branch 'main' into feature/exhibit-linking

* Merge branch 'main' into feature/exhibit-linking

* Remove test notebooks

* Delete test notebook

* Save

* Merge branch 'main' into feature/exhibit-linking

* try except for exhbit page finding

* Prep for merge

* update test

* Remove test

* Merge branch 'main' into feature/exhibit-linking

* update unit test

* Fix unit test

* replace try except

* Update page sort

* Add docstring


Approved-by: Alex Galarce
This commit is contained in:
Katon Minhas
2025-06-10 19:50:15 +00:00
parent 0c07167094
commit baf5002b2c
11 changed files with 179 additions and 2336 deletions
+81 -40
View File
@@ -211,67 +211,108 @@ def filter_quick_review(text_dict):
or "COVER SHEET" in page_text.upper()
}
def get_exhibit_pages(text_dict: dict[str, str], filename: str) -> list[str]:
"""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 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]: 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)
def get_exhibit_pages(text_dict: dict[str, str], filename: str) -> tuple[list[str], dict[str, str]]:
"""
Identify exhibit pages from the text dictionary.
This function processes the text 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 (dict[str, str]): A dictionary where keys are page numbers (as strings) and values are the text of those pages.
filename (str): The name of the file being processed, used for logging and LLM invocation.
Returns:
tuple[list[str], dict[str, 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 the headers of those exhibits.
"""
sorted_pages = sorted(text_dict.keys(), key=string_utils.page_key_sort)
exhibit_pages = []
exhibit_header_dict = {}
first_page = True
# 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:
page_content = text_dict[page_num]
# Always include the first page as an exhibit
if first_page:
exhibit_pages.append(page_num)
exhibit_header_dict[page_num] = ""
is_exhibit = True
first_page = False
else:
if "." not in page_num or ".0" in page_num:
page_content = text_dict[page_num]
prompt = investment_prompts.EXHIBIT_CHECK(page_content[0:200])
prompt = investment_prompts.EXHIBIT_HEADER(page_content[0:400])
claude_answer_raw = llm_utils.invoke_claude(
prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=150
)
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)
is_exhibit = True
else:
if "N/A" in claude_answer_extracted:
is_exhibit = False
else:
exhibit_pages.append(page_num)
exhibit_header_dict[page_num] = claude_answer_extracted
is_exhibit = True
elif is_exhibit==True:
exhibit_pages.append(page_num)
return exhibit_pages
return exhibit_pages, exhibit_header_dict
def link_exhibit_pages(all_exhibit_headers, filename):
"""
Filters exhibit pages to include only those with meaningfully different headers.
This function processes exhibit headers sequentially and uses literal comparison and
LLM-based semantic comparison to determine when exhibits represent continuation pages
vs. distinct exhibits. Continuation pages are filtered out of the headers list and
mapping, and when used downstream are included in their parent exhibit's text chunk.
Links exhibit pages based on their headers. If the header of a page is identical (either literally or semantically) to the previous page's header, it is not added to the final list of exhibit pages.
This function will link contiguous header pages together, so that the final list of exhibit pages will only contain pages that are meaningfully different from the previous page. E.g. "1" with header {"Exhibit 1"} and "2" with header {"Continued exhibit 1"} will be linked together, but "3" with header {"Exhibit 2"} will not be linked to either of them.
Args:
all_exhibit_headers (dict[str, str]): A dictionary mapping page numbers (as strings) to their exhibit headers. Expected to be in chronological order
filename (str): The name of the file being processed, used for logging and LLM invocation.
Returns:
tuple[list[str], dict[str, 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 the headers of those exhibits.
Example:
Input: {"1": "Exhibit A", "2": "Exhibit A (continued)", "3": "Exhibit B"}
Output (["1", "3"], {"1": "Exhibit A", "3": "Exhibit B"})
Notes:
- First exhibit page is always included.
- Literal header comparison is performed first to optimize LLM usage and latency.
- Uses `EXHIBIT_LINKAGE` prompt for semantic similarity assessment
- Expects LLM to return "Y" or "N" enclosed in pipe delimiters
- Pages with non-distinct headers are excluded from the final results.
"""
first_exhibit = True
final_exhibit_pages = []
final_exhibit_headers = {}
for page_num, page_header in all_exhibit_headers.items():
if first_exhibit:
exhibit_is_different = "N"
previous_header = page_header
final_exhibit_pages.append(page_num)
final_exhibit_headers[page_num] = page_header
first_exhibit = False
else:
# If literally identical, skip prompt
if page_header.upper().strip() == previous_header.upper().strip():
exhibit_is_different = "N"
else:
prompt = investment_prompts.EXHIBIT_LINKAGE(page_header, previous_header)
claude_answer_raw = llm_utils.invoke_claude(prompt, config.MODEL_ID_CLAUDE35_SONNET, filename)
exhibit_is_different = string_utils.extract_text_from_delimiters(claude_answer_raw, Delimiter.PIPE)
# If the prompt has identified that the current exhibit is meaningfully different than the previous
if "Y" in exhibit_is_different:
final_exhibit_pages.append(page_num)
final_exhibit_headers[page_num] = page_header
previous_header = page_header
return final_exhibit_pages, final_exhibit_headers
def chunk_by_exhibit(text_dict: dict,
exhibit_pages: list