Merged in feature/document-index (pull request #1005)
Feature/document index
* Add Document Index preprocessing — Layers 1, 2, and 3 wiring
Parse the Textract-emitted Document Index block at the top of each contract
with a single cached LLM call (prompt_document_index) instead of one per-page
call per page. Layer 2 verifies parsed entries via literal string match and
structural regex sweep, escalating suspect pages back to the existing per-page
path. Layer 1+2 failure triggers a full fallback to today's per-page flow.
New symbols:
- preprocessing_funcs.extract_document_index_block — regex slice of index prefix
- preprocessing_funcs.verify_index_against_pages — structural verifier (plain dict return)
- prompt_templates.DOCUMENT_INDEX_INSTRUCTION / DOCUMENT_INDEX — cached prompt pair
- prompt_calls.prompt_document_index — LLM wrapper (usage_label DOCUMENT_INDEX_PARSE)
- config: DOCUMENT_INDEX_PARSE_ENABLED and three threshold flags
- instrumentation: DOCUMENT_INDEX_PARSE mapped to preprocessing segment
one_to_n_exhibit_chunking now accepts optional contract_text; file_processing.py
pas…
* Add exhibit-level relevance filtering and deprecate pipe-delimited output
Step 5: extend the Document Index LLM call to emit reimbursement_relevant per
exhibit (no extra cost), then filter exhibit_header_dict before downstream
one-to-n extraction. Safety net A (keyword override) and safety net B (never
drop all exhibits) prevent false negatives. Controlled by
DOCUMENT_INDEX_RELEVANCE_FILTER_ENABLED flag.
Remove all pipe-delimited output from prompt_templates.py:
EXHIBIT_HEADER_INSTRUCTION, EXHIBIT_HEADER, EXHIBIT_HEADER_DEDUP_INSTRUCTION,
and EXHIBIT_HEADER_DEDUP now use JSON_DICT_FORMAT_INSTRUCTIONS and the
standard "return your final answer as a valid JSON dictionary" closing pattern.
DOCUMENT_INDEX_INSTRUCTION already used the JSON standard; consistent across
all preprocessing prompts.
* Add Tier 1 and Tier 2 unit tests for Document Index feature
Tier 1 — src/tests/test_document_index.py (31 tests):
- TestExtractDocumentIndexBlock: regex extraction (6 tests)
- TestDocumentIndexPromptTemplates: instruction shape, parser, cache
registration (11 tests)
- TestPromptDocumentIndex: mocked LLM wrapper — usage_label, cache=True,
fallback on malformed/non-list/empty responses (6 tests)
- TestVerifyIndexAgainstPages: literal match pass/fail, coverage sweep
escalation, page drift, high-failure-ratio fallback, behavioural pin
that Layer 2 never adds headers from regex (8 tests)
Tier 2 — src/tests/test_preprocessing_funcs.py (6 new tests):
- TestOneToNExhibitChunkingHybrid: Layer 1 success path, escalation routing,
global fallback, feature-flag-off, no contract_text, no index block
* Fix relevance overwrite bug and remove unused config flag
Bug: relevance_lookup was built from all parsed entries, so a landmark or
noise entry on the same page as an exhibit_start could overwrite the
exhibit_start's reimbursement_relevant=True flag. Fixed by filtering the
lookup to exhibit_start entries only.
Remove DOCUMENT_INDEX_COVERAGE_THRESHOLD from config.py — the flag was
defined but never referenced in the verifier or orchestrator. Removing
it avoids misleading operators expecting coverage-threshold fallback.
Docstring correction: extract_document_index_block docstring claimed
BOM-tolerance; removed that claim since Textract-generated .txt files
don't carry BOMs and the implementation doesn't strip .
Add TestRelevanceFilterLogic (3 tests): landmark-overwrite regression,
irrelevant exhibit exclusion, and safety-net-B preservation.
* Add Tier 3 offline study script for Document Index
Runs Layer 1+2 (extract, LLM parse, structural verify) against every .txt
in an input dir without the full pipeline. Emits a per-file CSV with 17
metrics columns (coverage_score, classification distribution, verification
failures, escalation count, fallback flag, runtime) and a summary markdown.
Marked DELETABLE post-rollout — intended for threshold tuning before Tier 4
E2E runs, not for production.
Usage:
uv run python -m src.testbed.document_index_offline
uv run python -m src.testbed.document_index_offline --input-dir data/pacific-source
uv run python -m src.testbed.document_index_offline --max-files 10
* Fix relevance overwrite bug and remove unused config flag
Remove coverage-sweep escalation — high-trust architecture
The Document Index path is now authoritative when it passes quality
thresholds. The per-page escalation of individual uncovered pages is
removed; the only fallback is the whole-document global fallback (empty
parse, page drift, low exhibit density, high failure ratio). This
eliminates the 8–50 false-positive escalations per file observed in the
Tier 3 offline study (continuation pages containing sub-headings were
being sent back to the per-page LLM unnecessarily).
verify_index_against_pages now returns 3 keys instead of 4:
verified_dict, metrics, global_fallback (escalation_pages removed)
one_to_n_exhibit_chunking success path no longer runs get_exhibit_pages_new
or prompt_header_deduplication — it returns verified_dict directly after
the relevance filter.
Add TestVerifyIndexAgainstPagesV2 (40 tests) written contract-first by the
testing agent: return-shape invariants, all 4 fallback tr…
* Trim TestVerifyIndexAgainstPagesV2 to 10 targeted boundary tests
* Remove exhibit-level relevance filter (Step 5)
The filter classified exhibits as reimbursement-relevant or not based
purely on header text from the Document Index table of contents — never
reading the exhibit content. This is unreliable: "AMENDMENT NO. 3" or
"GENERAL TERMS" could contain critical rate changes; the header alone
cannot distinguish them from truly irrelevant exhibits.
Gate 4 (has_reimbursements check after Step 1) already handles this
correctly by reading actual content. The header-based pre-filter added
correctness risk without proportional benefit.
Removed: DOCUMENT_INDEX_RELEVANCE_FILTER_ENABLED config flag, the Step 5
filter block in one_to_n_exhibit_chunking, reimbursement_relevant from
DOCUMENT_INDEX_INSTRUCTION output schema (back to 3 keys: page, header,
classification), TestRelevanceFilterLogic test class. Plan document
updated with reasoning.
Also fix document_index_offline.py to use split_text + lightweight page
wrapper instead of split_text_with_pages, avoiding expensive table
sp…
* Enrich comparison script with page spans and DI parent mapping
* Merged dev into feature/document-index
* Restructure Document Index quality gates: fuzzy match + failure-ratio + tail-coverage
Verifier semantics changed to lead with the index and only fall back on real
quality signals:
- Fuzzy header match: literal-with-whitespace match first, then rapidfuzz
partial_ratio (threshold 85) on the top of the page. Handles dash variants
and minor formatting differences.
- Failure-ratio gate: >10% of exhibit_start entries failing the fuzzy match
triggers fallback (parse is untrustworthy).
- Tail-coverage gate: replaces the old 'first exhibit must start at page 1'
rule, which fired falsely on contracts where page 1 is preamble. New rule
measures the trailing span: tail_pages = total_pages - max_exhibit_start.
For docs >15 pages, fallback when tail > max(15, 30%) of the document — the
signal that the LLM stopped parsing partway through the index.
- Thresholds are hardcoded module constants in preprocessing_funcs.py
(_DI_FUZZY_MATCH_THRESHOLD, _DI_MAX_FAILURE_RATIO, _DI_TAIL_COVERAGE_*) —
not config fla…
* Move Document Index testbed scripts to local_scripts (untracked)
The DI offline study and comparison tools are throwaway local utilities, not
production code. Moving them to local_scripts/ (already gitignored) keeps
src/testbed/ clean and removes them from the repo.
* Document Index: add SECTION as exhibit_start keyword, disambiguate sub-clauses
Empirical signal from a 100-file generic-corpus comparison run showed 37
legacy-only exhibit boundaries on three SECTION-style contracts (Gulf Coast
Division, UTMB Amendment 10, Rehab Designs) that DI was missing because
SECTION wasn't listed alongside EXHIBIT/ATTACHMENT/ARTICLE/etc. Also added
a counter-example to the landmark bullet so numeric sub-clauses like
'2.3 Confidentiality' or '10.4 Billing Procedures' stay out of exhibit_start.
Two unit tests pin the new behavior: SECTION appears in the exhibit_start
paragraph, and the numeric-sub-clause disambiguation appears in landmark.
* Trim Document Index tests: drop low-value prompt-content pinning, consolidate
41 tests for a ~200-line feature was over-investment. Cut to 25 by:
- Dropping prompt-string-content tests that pinned copy rather than behavior
(will-break-on-tuning, value-add zero) — non-empty-string, required-section
headers, noise-example pinning, SECTION/landmark prompt-content checks
- Removing trivial happy-path tests for DOCUMENT_INDEX factory (returns tuple,
embeds block, parser parses valid JSON) that duplicate parser-level coverage
in test_json_parsers.py
- Trimming TestPromptDocumentIndex to the three meaningful contracts:
usage_label, cache=True, malformed-JSON tolerance
- Folding TestVerifyIndexAgainstPagesBoundaries into the main verify class
and consolidating four 'doesn't crash' tests into a single sweep test
Net: 41 → 25 tests. Test density now matches comparable small features in
the codebase (test_clean_header_footer 13 tests, test_chc_extraction 6).
* Document Index: rewrite verified headers with page-text via single LLM call
The Document Index returns thin TOC entries ("EXHIBIT A - FEE SCHEDULE");
downstream code (chunk-boundary matching, dynamic-primary LLM input,
EXHIBIT_TITLE column, active-rates exhibit-title standardization) is
calibrated for the rich page-text version that the legacy per-page path
produced (multi-line headers with provider names, effective dates, etc.).
After DI parse + verification, bundle the full text of each verified page
under per-page boundary markers and run a single LLM call to extract the
page-formatted headers. Two-headers-on-one-page is supported by allowing
list values per boundary — that case is the original reason DI is the
preferred path. On any LLM/parse failure the original verified_dict is
preserved so the pipeline never breaks.
Cost stays bounded by count of header-bearing pages (typically 5-15 per
contract), independent of exhibit length. A 60-page contract with 8
exhibits costs 2 LLM calls (DI parse + bundle …
* Merge remote-tracking branch 'origin/dev' into feature/document-index
* Document Index: recover orphan-anchor pages dropped by TOC
When OCR strips the "ATTACHMENT X" / "Exhibit N" label from a Document Index
line, leaving only the secondary title (e.g. "MEDICARE ADVANTAGE COMPENSATION"
where "ATTACHMENT B" was expected), the DI LLM classifies the orphan secondary
as `landmark`, the page never enters verified_dict, and the exhibit is silently
dropped - even though the actual page begins with the anchor label.
Add a regex-based recovery pass that runs between verify_index_against_pages
and replace_index_headers_with_page_text. Scans every unverified page for a
strict whole-line header anchor, then applies five guards before promoting:
1. Strict whole-line regex (optional " - Secondary Title" suffix)
2. Length cap (anchors are short; <=80 chars)
3. Running-header filter (line on >=2 pages AND >40% of pages)
4. List detection (>=3 distinct anchors on one page -> body-text list, skip)
5. Continuation guard (same anchor as previous page -> multi-page exhibit, skip)
No new …
* Document Index: strip running-header bleed from page-text rewriter
The verified-pages header LLM was concatenating document running headers
(e.g. 'PROVIDER SERVICES AGREEMENT' repeated on every page) into the
rewritten exhibit title, producing entries like
'PROVIDER SERVICES AGREEMENTExhibit 1Services & Compensation'.
Lift the running-header set computation to the DI orchestrator so a single
set is shared between orphan-anchor recovery (Fix A) and the page-text
rewriter (Fix B). The rewriter now strips lines matching the set from each
page before bundling, so the LLM never sees the bleed and can't include it.
Verified on the 6 affected Clover contracts: DAVITA p7, Amedisys p34, HMH
p3, Inspira p4/p5, PHOEBE p17, CedarBridge p6 — all now produce clean
anchor + newline + secondary-title titles.
* Merge remote-tracking branch 'origin/dev' into feature/document-index
Approved-by: Karan Desai
Approved-by: Faizan Mohiuddin
This commit is contained in:
Praneel Panchigar
committed by
Katon Minhas
Katon Minhas
parent
d0edd6b4b0
commit
2306334b40
@@ -375,6 +375,10 @@ MIN_TABLE_ROWS_TO_SPLIT = int(get_arg_value("min_table_rows_to_split", "20"))
|
||||
# Table splitting method (currently only "cells" is supported)
|
||||
TABLE_SPLITTING_METHOD = get_arg_value("table_splitting_method", "cells")
|
||||
|
||||
DOCUMENT_INDEX_PARSE_ENABLED = (
|
||||
get_arg_value("document_index_parse_enabled", "True") == "True"
|
||||
)
|
||||
|
||||
######################################## OUTPUT FILE SPLITTING SETTINGS ########################################
|
||||
# Maximum rows per split file (files are split to keep filenames together)
|
||||
MAX_ROWS_PER_SPLIT = int(get_arg_value("max_rows_per_split", "70000"))
|
||||
|
||||
@@ -97,6 +97,7 @@ def process_file(file_object, constants: Constants, run_timestamp):
|
||||
text_dict=text_dict,
|
||||
EXHIBIT_HEADER_MARKERS=constants.EXHIBIT_HEADER_MARKERS,
|
||||
filename=filename,
|
||||
contract_text=contract_text,
|
||||
)
|
||||
|
||||
logging.info(f"{datetime_str()} Preprocessing Complete - {filename}")
|
||||
|
||||
@@ -850,6 +850,98 @@ def prompt_header_deduplication(
|
||||
return exhibit_header_dict
|
||||
|
||||
|
||||
def prompt_document_index(index_block: str, filename: str) -> list[dict]:
|
||||
"""Parse a Document Index block into structured exhibit entries using LLM.
|
||||
|
||||
Args:
|
||||
index_block: Raw Document Index text from the contract prefix.
|
||||
filename: Source filename for logging and LLM attribution.
|
||||
|
||||
Returns:
|
||||
list[dict]: Parsed entries with keys 'page', 'header', 'classification'.
|
||||
Returns [] on parse failure (caller treats empty as global fallback).
|
||||
"""
|
||||
if not index_block:
|
||||
return []
|
||||
|
||||
prompt, parser = prompt_templates.DOCUMENT_INDEX(index_block)
|
||||
|
||||
llm_answer_raw = llm_utils.invoke_claude(
|
||||
prompt,
|
||||
"sonnet_latest",
|
||||
filename,
|
||||
max_tokens=4096,
|
||||
cache=True,
|
||||
instruction=prompt_templates.DOCUMENT_INDEX_INSTRUCTION(),
|
||||
usage_label="DOCUMENT_INDEX_PARSE",
|
||||
)
|
||||
|
||||
try:
|
||||
result = parser(llm_answer_raw)
|
||||
if not isinstance(result, list):
|
||||
logging.warning(
|
||||
f"DOCUMENT_INDEX_PARSE returned non-list, falling back - {filename}"
|
||||
)
|
||||
return []
|
||||
return result
|
||||
except (json.JSONDecodeError, ValueError) as e:
|
||||
logging.error(
|
||||
f"Failed to parse DOCUMENT_INDEX_PARSE response: {e} - {filename}"
|
||||
)
|
||||
return []
|
||||
|
||||
|
||||
def prompt_exhibit_header_from_verified_pages(
|
||||
bundle: str, filename: str
|
||||
) -> dict[str, list[str]]:
|
||||
"""Extract page-formatted headers from a bundle of DI-verified pages.
|
||||
|
||||
The Document Index gives us thin TOC entries; this call rewrites each one
|
||||
with the rich page-text version (multi-line headers, provider names,
|
||||
effective dates, etc.) so downstream literal matching and dynamic-primary
|
||||
LLM extraction get the same shape they get from the legacy per-page path.
|
||||
|
||||
Args:
|
||||
bundle: Concatenated full-page texts separated by ===BOUNDARY_X_HEADER===
|
||||
markers (one boundary per verified page).
|
||||
filename: Source filename for logging and LLM attribution.
|
||||
|
||||
Returns:
|
||||
dict[str, list[str]]: {boundary_number_str: [header, ...]}. Returns {}
|
||||
on parse failure (caller treats empty as "keep DI index headers").
|
||||
"""
|
||||
if not bundle:
|
||||
return {}
|
||||
|
||||
prompt, parser = prompt_templates.EXHIBIT_HEADER_FROM_VERIFIED_PAGES(bundle)
|
||||
|
||||
llm_answer_raw = llm_utils.invoke_claude(
|
||||
prompt,
|
||||
"sonnet_latest",
|
||||
filename,
|
||||
max_tokens=4096,
|
||||
cache=True,
|
||||
instruction=prompt_templates.EXHIBIT_HEADER_FROM_VERIFIED_PAGES_INSTRUCTION(),
|
||||
usage_label="EXHIBIT_HEADER_FROM_VERIFIED_PAGES",
|
||||
)
|
||||
|
||||
try:
|
||||
result = parser(llm_answer_raw)
|
||||
if not isinstance(result, dict):
|
||||
logging.warning(
|
||||
f"EXHIBIT_HEADER_FROM_VERIFIED_PAGES returned non-dict; "
|
||||
f"falling back - {filename}"
|
||||
)
|
||||
return {}
|
||||
return result
|
||||
except (json.JSONDecodeError, ValueError) as e:
|
||||
logging.error(
|
||||
f"Failed to parse EXHIBIT_HEADER_FROM_VERIFIED_PAGES response: "
|
||||
f"{e} - {filename}"
|
||||
)
|
||||
return {}
|
||||
|
||||
|
||||
def prompt_date_fix(date: str) -> str:
|
||||
"""Processes dates using `prompt_templates.date_fix_prompt`.
|
||||
|
||||
|
||||
@@ -81,6 +81,7 @@ def one_to_n_exhibit_chunking(
|
||||
pages_dict=None,
|
||||
EXHIBIT_HEADER_MARKERS=None,
|
||||
filename=None,
|
||||
contract_text=None,
|
||||
) -> dict[str, list[str]]:
|
||||
"""
|
||||
Process the text dictionary or pages dictionary to identify exhibit headers.
|
||||
@@ -96,6 +97,10 @@ def one_to_n_exhibit_chunking(
|
||||
(preferred)
|
||||
EXHIBIT_HEADER_MARKERS: List of markers to identify exhibit headers
|
||||
filename: The name of the file being processed.
|
||||
contract_text: Optional raw contract text (post-clean_text, pre-split_text).
|
||||
When provided and DOCUMENT_INDEX_PARSE_ENABLED is True, enables
|
||||
the Layer 1 Document Index path (one LLM call instead of per-page).
|
||||
Falls back to per-page path on parse failure. (default: None)
|
||||
|
||||
Returns:
|
||||
dict[str, list[str]]: A dictionary where keys are page numbers and values are
|
||||
@@ -106,35 +111,110 @@ def one_to_n_exhibit_chunking(
|
||||
Either text_dict or pages_dict must be provided. If both are provided,
|
||||
pages_dict takes precedence.
|
||||
"""
|
||||
# Determine which dictionary to use
|
||||
if pages_dict is not None:
|
||||
_, 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:
|
||||
_, 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")
|
||||
from src import config as _config
|
||||
|
||||
# ── Helper: legacy per-page flow (only runs when no Document Index block) ─
|
||||
def _run_per_page_path():
|
||||
if pages_dict is not None:
|
||||
_, result = preprocessing_funcs.get_exhibit_pages_new(
|
||||
pages_dict=pages_dict,
|
||||
EXHIBIT_HEADER_MARKERS=EXHIBIT_HEADER_MARKERS,
|
||||
filename=filename,
|
||||
)
|
||||
elif text_dict is not None:
|
||||
_, result = 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")
|
||||
return result
|
||||
|
||||
# ── Document Index path (primary) ────────────────────────────────────────
|
||||
# Lead with the index. Fall back to the legacy per-page path when:
|
||||
# 1. The contract has no Document Index block, OR
|
||||
# 2. The index parse fails the regex sanity check (>10% headers don't
|
||||
# verify on their claimed page), OR
|
||||
# 3. The parse leaves coverage gaps (first exhibit doesn't start at page 1).
|
||||
if (
|
||||
_config.DOCUMENT_INDEX_PARSE_ENABLED
|
||||
and contract_text
|
||||
and pages_dict is not None
|
||||
):
|
||||
index_block = preprocessing_funcs.extract_document_index_block(contract_text)
|
||||
if index_block:
|
||||
parsed = prompt_calls.prompt_document_index(index_block, filename)
|
||||
verification = preprocessing_funcs.verify_index_against_pages(
|
||||
parsed, pages_dict
|
||||
)
|
||||
logging.info(
|
||||
f"Document Index parse: {verification['metrics']} - {filename}"
|
||||
)
|
||||
|
||||
if not verification["should_fallback"]:
|
||||
# Compute the document's running-header set once and share it
|
||||
# between Fix A (orphan-anchor recovery) and Fix B (page-text
|
||||
# rewriter). Single source of truth — they can't drift apart
|
||||
# on the definition of "running header" for this file.
|
||||
running_headers = preprocessing_funcs.build_running_header_set(
|
||||
pages_dict
|
||||
)
|
||||
# Recover pages whose TOC dropped the leading "ATTACHMENT X" /
|
||||
# "Exhibit N" label, leaving only the secondary title. The DI
|
||||
# LLM classifies the orphan secondary as `landmark`, so the page
|
||||
# never enters verified_dict despite the page itself starting
|
||||
# with the anchor. Strict regex-based promotion; no new LLM call.
|
||||
recovered_dict = preprocessing_funcs.recover_orphan_anchor_pages(
|
||||
verification["verified_dict"],
|
||||
pages_dict,
|
||||
filename,
|
||||
running_headers=running_headers,
|
||||
)
|
||||
# The index gives thin TOC headers; replace with rich
|
||||
# page-text headers so downstream (chunk-boundary matching,
|
||||
# dynamic primary, EXHIBIT_TITLE, active-rates standardization)
|
||||
# gets the same shape it gets from the legacy per-page path.
|
||||
# Strip running-header lines before bundling so the LLM doesn't
|
||||
# concatenate "PROVIDER SERVICES AGREEMENT" (or similar) into
|
||||
# the rewritten exhibit title (Fix B).
|
||||
page_text_dict = (
|
||||
preprocessing_funcs.replace_index_headers_with_page_text(
|
||||
recovered_dict,
|
||||
pages_dict,
|
||||
filename,
|
||||
running_headers=running_headers,
|
||||
)
|
||||
)
|
||||
exhibit_header_dict = dict(
|
||||
sorted(
|
||||
page_text_dict.items(),
|
||||
key=lambda x: int(x[0].split(".")[0]),
|
||||
)
|
||||
)
|
||||
return exhibit_header_dict
|
||||
|
||||
logging.info(
|
||||
f"Document Index quality gate failed "
|
||||
f"({verification['fallback_reason']}); "
|
||||
f"falling back to legacy per-page path - {filename}"
|
||||
)
|
||||
else:
|
||||
logging.info(
|
||||
f"No Document Index block found — falling back to legacy per-page path - {filename}"
|
||||
)
|
||||
|
||||
# ── Legacy per-page fallback ─────────────────────────────────────────────
|
||||
exhibit_header_dict = _run_per_page_path()
|
||||
|
||||
# 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
|
||||
)
|
||||
|
||||
# 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]))
|
||||
)
|
||||
|
||||
# Fallback: if no exhibits were detected, use the first 10 non-empty lines of the
|
||||
# first page as a synthetic header so downstream extraction has something to work with
|
||||
if not exhibit_header_dict:
|
||||
source = pages_dict if pages_dict is not None else text_dict
|
||||
sorted_pages = sorted(source.keys(), key=lambda x: int(x.split(".")[0]))
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
import logging
|
||||
import re
|
||||
from collections import Counter
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from rapidfuzz import fuzz
|
||||
|
||||
import src.utils.string_utils as string_utils
|
||||
from src.pipelines.shared.prompts import prompt_calls
|
||||
from src.pipelines.shared.extraction.exhibit_funcs import _find_header_in_text
|
||||
from src import config
|
||||
from src.constants.delimiters import (
|
||||
TABLE_START_MARKER,
|
||||
@@ -16,6 +20,526 @@ if TYPE_CHECKING:
|
||||
from src.pipelines.shared.extraction.exhibit_funcs import Exhibit
|
||||
|
||||
|
||||
# ── Document Index quality gates (hardcoded — these don't move often) ────────
|
||||
# Two gates trigger fallback to the legacy per-page path:
|
||||
# 1. Failure ratio: more than _DI_MAX_FAILURE_RATIO of exhibit_start entries
|
||||
# fail the fuzzy header-on-page check.
|
||||
# 2. Tail coverage gap: the largest verified exhibit_start sits well before
|
||||
# total_pages, leaving a long tail span that strongly suggests the LLM
|
||||
# stopped parsing partway through the index. Pages BEFORE the first
|
||||
# exhibit are NOT checked — preamble/recitals/cover pages are normal and
|
||||
# page 1 may itself be where the index lives.
|
||||
_DI_FUZZY_MATCH_THRESHOLD = 85 # rapidfuzz partial_ratio score (0–100)
|
||||
_DI_MAX_FAILURE_RATIO = 0.10
|
||||
_DI_TAIL_COVERAGE_MIN_PAGES = 15 # don't apply the tail check on shorter docs
|
||||
_DI_TAIL_COVERAGE_FLOOR_PAGES = 15 # tail must exceed this absolute floor
|
||||
_DI_TAIL_COVERAGE_RATIO = 0.30 # ...and this fraction of the doc
|
||||
_DI_PAGE_HEAD_CHARS = 2000 # only fuzzy-match against the top of the page
|
||||
|
||||
|
||||
def extract_document_index_block(contract_text: str) -> Optional[str]:
|
||||
"""Extract the Document Index block from raw contract text.
|
||||
|
||||
Slices everything before the first "Start of Page No. =" marker.
|
||||
Returns None if the prefix does not begin with "Document Index"
|
||||
(case-insensitive, leading-whitespace-tolerant) or if the marker is absent.
|
||||
"""
|
||||
marker = "Start of Page No. ="
|
||||
idx = contract_text.find(marker)
|
||||
if idx == -1:
|
||||
return None
|
||||
prefix = contract_text[:idx].strip()
|
||||
if not prefix.lower().startswith("document index"):
|
||||
return None
|
||||
return prefix
|
||||
|
||||
|
||||
def _normalize_for_fuzzy(s: str) -> str:
|
||||
"""Normalize a string for fuzzy header-on-page matching."""
|
||||
s = s.lower()
|
||||
# Unify dash variants (em-dash, en-dash, minus)
|
||||
for d in ("—", "–", "−"):
|
||||
s = s.replace(d, "-")
|
||||
# Collapse whitespace
|
||||
return re.sub(r"\s+", " ", s).strip()
|
||||
|
||||
|
||||
def _fuzzy_header_in_page(page_text: str, header: str) -> bool:
|
||||
"""Return True if `header` appears (closely enough) on this page.
|
||||
|
||||
Strategy: try the existing literal-with-whitespace match first (cheap and
|
||||
decisive when it hits). On miss, normalize both sides and run rapidfuzz
|
||||
partial_ratio against just the top of the page (where exhibit headers live).
|
||||
"""
|
||||
if not header or not page_text:
|
||||
return False
|
||||
if _find_header_in_text(page_text, header) >= 0:
|
||||
return True
|
||||
norm_h = _normalize_for_fuzzy(header)
|
||||
if not norm_h:
|
||||
return False
|
||||
norm_page_head = _normalize_for_fuzzy(page_text[:_DI_PAGE_HEAD_CHARS])
|
||||
return fuzz.partial_ratio(norm_h, norm_page_head) >= _DI_FUZZY_MATCH_THRESHOLD
|
||||
|
||||
|
||||
def verify_index_against_pages(
|
||||
parsed_index: list[dict],
|
||||
pages_dict: dict,
|
||||
) -> dict:
|
||||
"""Verify LLM-parsed Document Index entries against the actual page text.
|
||||
|
||||
For each `exhibit_start` entry the LLM produced, run a fuzzy header-on-page
|
||||
check (whitespace-normalized literal match, falling back to rapidfuzz
|
||||
partial_ratio against the top of the page). Pass → entry lands in
|
||||
`verified_dict`; fail → counted toward the failure ratio.
|
||||
|
||||
The path falls back to the legacy per-page system if either quality gate
|
||||
fails:
|
||||
• failure_ratio > _DI_MAX_FAILURE_RATIO (10%) — too many index entries
|
||||
that don't actually appear on the pages they claim.
|
||||
• tail coverage gap — the largest verified exhibit_start sits well before
|
||||
the end of the document, suggesting the LLM stopped parsing partway
|
||||
through. Pages BEFORE the first verified exhibit are NOT a gap; cover
|
||||
pages and recitals are normal.
|
||||
|
||||
Args:
|
||||
parsed_index: List of dicts from prompt_document_index, each with keys
|
||||
'page' (int|None), 'header' (str), 'classification' (str).
|
||||
pages_dict: Dict mapping page-number strings to Page objects.
|
||||
|
||||
Returns:
|
||||
dict with keys 'verified_dict', 'metrics', 'should_fallback',
|
||||
'fallback_reason' (empty string when not falling back).
|
||||
"""
|
||||
exhibit_starts = [
|
||||
e for e in parsed_index if e.get("classification") == "exhibit_start"
|
||||
]
|
||||
|
||||
# Total page count (base pages, ignoring sub-page splits like "3.1")
|
||||
total_pages = len({k.split(".")[0] for k in pages_dict.keys()})
|
||||
|
||||
metrics: dict = {
|
||||
"total_parsed": len(parsed_index),
|
||||
"exhibit_starts": len(exhibit_starts),
|
||||
"verification_passes": 0,
|
||||
"verification_failures": 0,
|
||||
"failure_ratio": 0.0,
|
||||
"max_start_page": None,
|
||||
"tail_pages": None,
|
||||
"total_pages": total_pages,
|
||||
}
|
||||
|
||||
verified_dict: dict[str, list[str]] = {}
|
||||
verified_start_pages: list[int] = []
|
||||
|
||||
for entry in exhibit_starts:
|
||||
raw_page = entry.get("page")
|
||||
header = entry.get("header", "").strip()
|
||||
if not header or raw_page is None:
|
||||
metrics["verification_failures"] += 1
|
||||
continue
|
||||
|
||||
page_key = str(raw_page)
|
||||
# Resolve page_key to an actual key in pages_dict (handle sub-pages)
|
||||
if page_key not in pages_dict:
|
||||
candidate = next(
|
||||
(k for k in pages_dict if k.split(".")[0] == page_key), None
|
||||
)
|
||||
if candidate is None:
|
||||
metrics["verification_failures"] += 1
|
||||
logging.debug(
|
||||
f"verify_index_against_pages: claimed page {page_key} absent from "
|
||||
f"pages_dict for header {header[:60]!r}"
|
||||
)
|
||||
continue
|
||||
page_key = candidate
|
||||
|
||||
page_text = pages_dict[page_key].get_text()
|
||||
if _fuzzy_header_in_page(page_text, header):
|
||||
verified_dict.setdefault(page_key, []).append(header)
|
||||
metrics["verification_passes"] += 1
|
||||
try:
|
||||
verified_start_pages.append(int(page_key.split(".")[0]))
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
metrics["verification_failures"] += 1
|
||||
logging.debug(
|
||||
f"verify_index_against_pages: header not found on page {page_key}: "
|
||||
f"{header[:60]!r}"
|
||||
)
|
||||
|
||||
# ── Quality gates ────────────────────────────────────────────────────────
|
||||
if exhibit_starts:
|
||||
metrics["failure_ratio"] = round(
|
||||
metrics["verification_failures"] / len(exhibit_starts), 4
|
||||
)
|
||||
|
||||
fallback_reason = ""
|
||||
if not exhibit_starts:
|
||||
fallback_reason = "no_exhibit_starts"
|
||||
elif metrics["failure_ratio"] > _DI_MAX_FAILURE_RATIO:
|
||||
fallback_reason = (
|
||||
f"failure_ratio_{metrics['failure_ratio']:.2%}"
|
||||
f"_exceeds_{_DI_MAX_FAILURE_RATIO:.0%}"
|
||||
)
|
||||
elif not verified_start_pages:
|
||||
fallback_reason = "no_verified_exhibits"
|
||||
else:
|
||||
max_start = max(verified_start_pages)
|
||||
tail_pages = max(0, total_pages - max_start)
|
||||
metrics["max_start_page"] = max_start
|
||||
metrics["tail_pages"] = tail_pages
|
||||
# Only check tail coverage on docs large enough for it to matter.
|
||||
if total_pages > _DI_TAIL_COVERAGE_MIN_PAGES:
|
||||
tail_threshold = max(
|
||||
_DI_TAIL_COVERAGE_FLOOR_PAGES,
|
||||
int(total_pages * _DI_TAIL_COVERAGE_RATIO),
|
||||
)
|
||||
if tail_pages > tail_threshold:
|
||||
fallback_reason = (
|
||||
f"tail_coverage_gap_{tail_pages}p_after_p{max_start}"
|
||||
f"_of_{total_pages}p"
|
||||
)
|
||||
|
||||
return {
|
||||
"verified_dict": verified_dict,
|
||||
"metrics": metrics,
|
||||
"should_fallback": bool(fallback_reason),
|
||||
"fallback_reason": fallback_reason,
|
||||
}
|
||||
|
||||
|
||||
# Soft cap on the bundle size sent to the page-text header LLM call. Sonnet's
|
||||
# context window is 200K tokens (~800K chars); we leave plenty of headroom for
|
||||
# prompt + response. Logged as a warning when exceeded — we still send the
|
||||
# bundle (the model can handle it), this is just an early-warning signal.
|
||||
_PAGE_HEADER_BUNDLE_WARN_CHARS = 500_000
|
||||
|
||||
# ── Orphan-anchor recovery (Fix A: TOC drops "ATTACHMENT X" / "Exhibit N") ───
|
||||
# Strict whole-line header anchors. The line must contain ONLY the anchor
|
||||
# (optionally followed by " — Secondary Title"), so body references like
|
||||
# "As required by Exhibit 3 of this Agreement..." don't match.
|
||||
# ATTACHMENT\s* (with optional whitespace) handles OCR-mangled forms like
|
||||
# "ATTACHMENTC" (no space between word and letter).
|
||||
_ANCHOR_LINE_RE = re.compile(
|
||||
r"^\s*(?P<anchor>"
|
||||
r"EXHIBIT\s+[A-Z0-9][-A-Z0-9]*"
|
||||
r"|ARTICLE\s+[IVXLCDM0-9]+"
|
||||
r"|SCHEDULE\s+[A-Z0-9][-A-Z0-9]*"
|
||||
r"|APPENDIX\s+[A-Z0-9][-A-Z0-9]*"
|
||||
r"|ADDENDUM\s+[A-Z0-9][-A-Z0-9]*"
|
||||
r"|ATTACHMENT\s*[A-Z0-9][-A-Z0-9]*"
|
||||
r")"
|
||||
r"(?:\s*[—\-:|]\s*.+)?" # Optional " — Secondary Title"
|
||||
r"[.:]?\s*$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
# Lines appearing on more than this fraction of pages are treated as running
|
||||
# headers/footers and excluded from anchor matching.
|
||||
_RUNNING_HEADER_PAGE_RATIO = 0.4
|
||||
# Headers are short lines. Anything longer is almost certainly a body sentence
|
||||
# that happens to start with an anchor word.
|
||||
_MAX_ANCHOR_LINE_CHARS = 80
|
||||
# Pages with this many or more distinct anchor matches are body-text lists
|
||||
# (e.g. "Section 12.14 Entire Agreement: Attachment A — ... Attachment B — ..."),
|
||||
# not exhibit pages. Skip them.
|
||||
_MAX_DISTINCT_ANCHORS_ON_PAGE = 3
|
||||
|
||||
|
||||
def replace_index_headers_with_page_text(
|
||||
verified_dict: dict[str, list[str]],
|
||||
pages_dict: dict,
|
||||
filename: str,
|
||||
running_headers: Optional[set[str]] = None,
|
||||
) -> dict[str, list[str]]:
|
||||
"""Replace DI-supplied index-formatted headers with page-formatted ones.
|
||||
|
||||
The Document Index gives us thin TOC entries like "EXHIBIT A - FEE SCHEDULE";
|
||||
the actual page header is typically richer ("EXHIBIT A – FEE SCHEDULE\\n
|
||||
State of California Medicaid\\nEffective 07/01/2025"). Downstream code
|
||||
(chunk-boundary matching, dynamic-primary LLM, EXHIBIT_TITLE column,
|
||||
active-rates standardization) is calibrated for the rich page-text version.
|
||||
|
||||
This function bundles the FULL text of each verified page (not just the top —
|
||||
exhibit headers can start mid-page) under a per-page boundary marker, runs
|
||||
one LLM call to extract page-formatted headers, and replaces the verified_dict
|
||||
entries. On any failure, the original verified_dict is preserved so the
|
||||
pipeline never breaks.
|
||||
|
||||
Args:
|
||||
verified_dict: Output of verify_index_against_pages — {page_key: [headers]}.
|
||||
pages_dict: Page objects keyed by page string.
|
||||
filename: For logging and LLM attribution.
|
||||
running_headers: Optional pre-computed set of running-header lines for
|
||||
this document. When provided, each line matching the
|
||||
set is dropped from every page before bundling — this
|
||||
prevents the LLM from concatenating the document
|
||||
running header (e.g. "PROVIDER SERVICES AGREEMENT")
|
||||
into the rewritten exhibit title. Computed on demand
|
||||
when None.
|
||||
|
||||
Returns:
|
||||
New {page_key: [headers]} dict with page-text headers in place of index ones.
|
||||
"""
|
||||
if not verified_dict:
|
||||
return verified_dict
|
||||
|
||||
if running_headers is None:
|
||||
running_headers = build_running_header_set(pages_dict)
|
||||
|
||||
sorted_pages = sorted(verified_dict.keys(), key=lambda x: int(x.split(".")[0]))
|
||||
boundary_to_page = {i + 1: p for i, p in enumerate(sorted_pages)}
|
||||
|
||||
parts: list[str] = []
|
||||
for boundary_num, page_key in enumerate(sorted_pages, start=1):
|
||||
page = pages_dict.get(page_key)
|
||||
if page is None:
|
||||
continue
|
||||
parts.append(f"===BOUNDARY_{boundary_num}_HEADER===")
|
||||
parts.append(_strip_running_header_lines(page.get_text(), running_headers))
|
||||
parts.append("===END_OF_BOUNDARIES===")
|
||||
bundle = "\n".join(parts)
|
||||
|
||||
if len(bundle) > _PAGE_HEADER_BUNDLE_WARN_CHARS:
|
||||
logging.warning(
|
||||
f"Page-text header bundle is large ({len(bundle):,} chars) for {filename}; "
|
||||
f"sending in a single call but watch token usage."
|
||||
)
|
||||
|
||||
bundle_headers = prompt_calls.prompt_exhibit_header_from_verified_pages(
|
||||
bundle, filename
|
||||
)
|
||||
if not bundle_headers:
|
||||
# Empty dict returned on LLM failure or parse failure — keep DI headers
|
||||
return verified_dict
|
||||
|
||||
# Aggregate per-boundary results back to {page_key: [headers]}. The prompt
|
||||
# asks for list values per boundary, but tolerate a single string too.
|
||||
result_by_page: dict[str, list[str]] = {}
|
||||
for boundary_str, val in bundle_headers.items():
|
||||
try:
|
||||
boundary_num = int(boundary_str)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
page_key = boundary_to_page.get(boundary_num)
|
||||
if page_key is None:
|
||||
continue
|
||||
items = val if isinstance(val, list) else [val]
|
||||
cleaned = [
|
||||
s.strip()
|
||||
for s in items
|
||||
if isinstance(s, str) and s.strip() and s.strip().upper() != "N/A"
|
||||
]
|
||||
if cleaned:
|
||||
result_by_page[page_key] = cleaned
|
||||
|
||||
# For each verified page, accept the page-text headers when at least as many
|
||||
# were returned as DI claimed. If page-text returned fewer (likely the LLM
|
||||
# merged two distinct exhibits into one string), keep DI's index entries to
|
||||
# preserve the boundary count. This trades richness for correctness on the
|
||||
# rare merge case.
|
||||
final: dict[str, list[str]] = {}
|
||||
rewritten = 0
|
||||
kept_di = 0
|
||||
for page_key in sorted_pages:
|
||||
di_count = len(verified_dict[page_key])
|
||||
page_text_headers = result_by_page.get(page_key, [])
|
||||
if page_text_headers and len(page_text_headers) >= di_count:
|
||||
final[page_key] = page_text_headers
|
||||
rewritten += 1
|
||||
else:
|
||||
final[page_key] = verified_dict[page_key]
|
||||
kept_di += 1
|
||||
if page_text_headers:
|
||||
logging.info(
|
||||
f"Page-text returned {len(page_text_headers)} header(s) but DI "
|
||||
f"had {di_count} on page {page_key}; keeping DI entries - "
|
||||
f"{filename}"
|
||||
)
|
||||
|
||||
logging.info(
|
||||
f"Page-text header replacement: rewrote {rewritten} page(s), "
|
||||
f"kept DI on {kept_di} page(s) - {filename}"
|
||||
)
|
||||
return final
|
||||
|
||||
|
||||
def build_running_header_set(pages_dict: dict) -> set[str]:
|
||||
"""Identify lines that appear on >40% of pages — document running headers/footers.
|
||||
|
||||
Counts each distinct stripped line once per page (a line that repeats within
|
||||
a single page only counts once for that page). Returns the set of lines that
|
||||
appear on at least 2 pages AND on a >40% fraction of pages. The 2-page floor
|
||||
is essential — a line appearing on a single page can't be a running header
|
||||
by definition, and the fraction check alone would mis-classify it on short
|
||||
documents (e.g. 1/2 pages = 50% > 40%).
|
||||
|
||||
Public so the DI orchestrator can compute the set once per file and pass
|
||||
it into both `recover_orphan_anchor_pages` (Fix A — exclude running lines
|
||||
from anchor matching) and `replace_index_headers_with_page_text` (Fix B —
|
||||
strip running lines from page text before the LLM bundles it).
|
||||
"""
|
||||
line_counts: Counter[str] = Counter()
|
||||
total_pages = 0
|
||||
for page in pages_dict.values():
|
||||
total_pages += 1
|
||||
page_text = page.get_text() if hasattr(page, "get_text") else str(page)
|
||||
seen_on_page: set[str] = set()
|
||||
for line in page_text.split("\n"):
|
||||
stripped = line.strip()
|
||||
if stripped and stripped not in seen_on_page:
|
||||
seen_on_page.add(stripped)
|
||||
line_counts[stripped] += 1
|
||||
if total_pages == 0:
|
||||
return set()
|
||||
threshold = total_pages * _RUNNING_HEADER_PAGE_RATIO
|
||||
return {ln for ln, c in line_counts.items() if c >= 2 and c > threshold}
|
||||
|
||||
|
||||
def _strip_running_header_lines(page_text: str, running_lines: set[str]) -> str:
|
||||
"""Drop lines whose stripped form is in the running-header set.
|
||||
|
||||
Used before bundling page text for the LLM rewriter so the document
|
||||
running header (e.g. "PROVIDER SERVICES AGREEMENT" repeated on every page)
|
||||
doesn't get concatenated into the extracted exhibit title.
|
||||
"""
|
||||
if not running_lines:
|
||||
return page_text
|
||||
return "\n".join(
|
||||
line for line in page_text.split("\n") if line.strip() not in running_lines
|
||||
)
|
||||
|
||||
|
||||
def _find_anchor(page_text: str, running_lines: set[str]) -> Optional[str]:
|
||||
"""Return the first strict whole-line header anchor on the page, or None.
|
||||
|
||||
Applies these guards:
|
||||
1. Length cap (anchors are short)
|
||||
2. Running-header filter (skip lines that repeat across pages)
|
||||
3. Strict whole-line regex (anchor is the entire line)
|
||||
4. List-detection: if 3+ distinct anchors appear on one page, this is
|
||||
a body-text list (e.g. "Entire Agreement: Attachment A — title,
|
||||
Attachment B — title, ..."), not a real exhibit page; return None.
|
||||
"""
|
||||
matches: list[str] = []
|
||||
for line in page_text.split("\n"):
|
||||
stripped = line.strip()
|
||||
if not stripped:
|
||||
continue
|
||||
if len(stripped) > _MAX_ANCHOR_LINE_CHARS:
|
||||
continue
|
||||
if stripped in running_lines:
|
||||
continue
|
||||
m = _ANCHOR_LINE_RE.match(stripped)
|
||||
if m:
|
||||
matches.append(m.group("anchor").strip())
|
||||
if not matches:
|
||||
return None
|
||||
distinct = {m.upper().replace(" ", "") for m in matches}
|
||||
if len(distinct) >= _MAX_DISTINCT_ANCHORS_ON_PAGE:
|
||||
return None
|
||||
return matches[0]
|
||||
|
||||
|
||||
def _anchor_key(s: str) -> str:
|
||||
"""Normalize an anchor for comparison (case- and whitespace-insensitive)."""
|
||||
return s.upper().replace(" ", "").rstrip(".:")
|
||||
|
||||
|
||||
def _previous_page_anchor(
|
||||
page_keys_sorted: list[str], current_index: int, page_headers: dict[str, list[str]]
|
||||
) -> Optional[str]:
|
||||
"""Return the anchor of the previous page's first header, or None.
|
||||
|
||||
Used to detect multi-page continuations: a page whose top anchor matches
|
||||
the previous page's anchor is a continuation (e.g. ATTACHMENT C spans 2
|
||||
pages, the second of which is a table), not a new exhibit boundary.
|
||||
"""
|
||||
if current_index == 0:
|
||||
return None
|
||||
prev_key = page_keys_sorted[current_index - 1]
|
||||
prev_headers = page_headers.get(prev_key, [])
|
||||
if not prev_headers:
|
||||
return None
|
||||
first_line = prev_headers[0].split("\n")[0].strip()
|
||||
m = _ANCHOR_LINE_RE.match(first_line)
|
||||
if m:
|
||||
return _anchor_key(m.group("anchor"))
|
||||
return None
|
||||
|
||||
|
||||
def recover_orphan_anchor_pages(
|
||||
verified_dict: dict[str, list[str]],
|
||||
pages_dict: dict,
|
||||
filename: str,
|
||||
running_headers: Optional[set[str]] = None,
|
||||
) -> dict[str, list[str]]:
|
||||
"""Promote pages whose body contains a strict header anchor that DI missed.
|
||||
|
||||
Some contracts have a Document Index that drops the leading "ATTACHMENT X"
|
||||
or "Exhibit N" label, leaving only a secondary title. The DI LLM correctly
|
||||
classifies the orphan secondary as `landmark`, so the page never enters
|
||||
`verified_dict`, even though the actual page DOES start with the anchor.
|
||||
|
||||
This helper scans every unverified page for a strict whole-line header
|
||||
anchor and adds matching pages to `verified_dict` with the anchor as a
|
||||
seed entry. The downstream rewriter then enriches each recovered page with
|
||||
the rich multi-line page-text header.
|
||||
|
||||
Args:
|
||||
running_headers: Optional pre-computed running-header set. When the
|
||||
orchestrator computes this once for the file and shares
|
||||
it between Fix A (anchor recovery) and Fix B (page-text
|
||||
rewriter), pass it here to skip the duplicate work.
|
||||
Computed on demand when None.
|
||||
|
||||
Returns a new dict containing the original verified entries plus any
|
||||
recovered pages.
|
||||
"""
|
||||
if not pages_dict:
|
||||
return verified_dict
|
||||
running_lines = (
|
||||
running_headers
|
||||
if running_headers is not None
|
||||
else build_running_header_set(pages_dict)
|
||||
)
|
||||
recovered: dict[str, list[str]] = dict(verified_dict)
|
||||
sorted_keys = sorted(pages_dict.keys(), key=lambda k: int(k.split(".")[0]))
|
||||
recovered_count = 0
|
||||
for i, page_key in enumerate(sorted_keys):
|
||||
if page_key in verified_dict:
|
||||
continue
|
||||
# Skip sub-page table chunks (e.g. "3.1"); they aren't independent pages.
|
||||
if "." in page_key and not page_key.endswith(".0"):
|
||||
continue
|
||||
page = pages_dict[page_key]
|
||||
page_text = page.get_text() if hasattr(page, "get_text") else str(page)
|
||||
anchor = _find_anchor(page_text, running_lines)
|
||||
if not anchor:
|
||||
continue
|
||||
# Skip if the previous page already carries the same anchor — that's
|
||||
# a multi-page exhibit continuation (e.g. ATTACHMENT C spans pages
|
||||
# 35-36, the second being a table), not a new boundary.
|
||||
prev_anchor = _previous_page_anchor(sorted_keys, i, recovered)
|
||||
if prev_anchor and prev_anchor == _anchor_key(anchor):
|
||||
logging.debug(
|
||||
f"Skipped p{page_key} continuation (same anchor as previous "
|
||||
f"page): {anchor!r} - {filename}"
|
||||
)
|
||||
continue
|
||||
recovered[page_key] = [anchor]
|
||||
recovered_count += 1
|
||||
logging.info(
|
||||
f"Recovered orphan-anchor page {page_key}: {anchor!r} - {filename}"
|
||||
)
|
||||
if recovered_count:
|
||||
logging.info(
|
||||
f"Orphan-anchor recovery: promoted {recovered_count} page(s) - {filename}"
|
||||
)
|
||||
return recovered
|
||||
|
||||
|
||||
def remove_page_indicators(contract_text: str) -> str:
|
||||
"""Clean textract output by removing page number indicators in the form of "Page X of Y"
|
||||
|
||||
|
||||
+191
-10
@@ -238,6 +238,10 @@ def get_cacheable_instructions():
|
||||
# Preprocessing instructions (exhibit headers, linkage)
|
||||
try:
|
||||
cacheable_instructions["EXHIBIT_HEADER"] = EXHIBIT_HEADER_INSTRUCTION()
|
||||
cacheable_instructions["DOCUMENT_INDEX"] = DOCUMENT_INDEX_INSTRUCTION()
|
||||
cacheable_instructions["EXHIBIT_HEADER_FROM_VERIFIED_PAGES"] = (
|
||||
EXHIBIT_HEADER_FROM_VERIFIED_PAGES_INSTRUCTION()
|
||||
)
|
||||
cacheable_instructions["EXHIBIT_LINKAGE"] = EXHIBIT_LINKAGE_INSTRUCTION()
|
||||
cacheable_instructions["EXHIBIT_TITLE_MATCH"] = (
|
||||
EXHIBIT_TITLE_MATCH_INSTRUCTION()
|
||||
@@ -3026,9 +3030,9 @@ def EXHIBIT_HEADER_INSTRUCTION() -> str:
|
||||
"""Static instruction for EXHIBIT_HEADER prompt caching.
|
||||
Returns formatted instructions for extracting headers from chunked contract pages.
|
||||
"""
|
||||
return """[OBJECTIVE]
|
||||
return f"""[OBJECTIVE]
|
||||
Extract section headers from contract page chunks with boundary markers.
|
||||
Return a JSON dict in |pipes| with boundary numbers as keys and header as values.
|
||||
Return a JSON dict with boundary numbers as keys and header as values.
|
||||
|
||||
Each chunk is separated by a boundary marker in the format ===BOUNDARY_X_HEADER=== where X is the boundary number.
|
||||
|
||||
@@ -3114,14 +3118,16 @@ Output key "4": "N/A"
|
||||
- Form field placeholders ("Provider's Legal Name: ____", "Date Signed: ____") are not headers even when they appear at the top of a chunk.
|
||||
|
||||
[OUTPUT FORMAT]
|
||||
Return a JSON dictionary enclosed in |pipes| where:
|
||||
Return a JSON dictionary 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
|
||||
- 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"}|"""
|
||||
{{"1": "PROVIDER SERVICES AGREEMENT\\nSIGNATURE PAGE", "2": "N/A", "3": "Attachment A: Fee Schedule"}}
|
||||
|
||||
{JSON_DICT_FORMAT_INSTRUCTIONS}"""
|
||||
|
||||
|
||||
def EXHIBIT_HEADER(context) -> Tuple[str, Callable[[str], dict]]:
|
||||
@@ -3139,7 +3145,100 @@ def EXHIBIT_HEADER(context) -> Tuple[str, Callable[[str], dict]]:
|
||||
Here are the chunks to analyze:
|
||||
{context.replace('"', "'")}
|
||||
|
||||
Return ONLY the |{{dictionary}}| followed by a brief explanation."""
|
||||
Briefly explain your reasoning, then return your final answer as a valid JSON dictionary."""
|
||||
|
||||
return (prompt, _json_dict_parser)
|
||||
|
||||
|
||||
def EXHIBIT_HEADER_FROM_VERIFIED_PAGES_INSTRUCTION() -> str:
|
||||
"""Static instruction for EXHIBIT_HEADER_FROM_VERIFIED_PAGES prompt caching.
|
||||
|
||||
Used after Document Index parsing+verification: each chunk is one FULL contract
|
||||
page that the index already flagged as containing one or more exhibit/section
|
||||
starts. The job is to extract the page-formatted header text exactly as it
|
||||
appears (multi-line, with provider names, effective dates, etc.) so it
|
||||
matches verbatim downstream.
|
||||
"""
|
||||
return f"""[OBJECTIVE]
|
||||
Each chunk below is one FULL contract page that the document index has already
|
||||
flagged as containing at least one exhibit/section start. Extract the page-formatted
|
||||
header text(s) for each chunk, exactly as the text appears on the page.
|
||||
|
||||
Each chunk is separated by a boundary marker in the format ===BOUNDARY_X_HEADER===
|
||||
where X is the boundary number (one boundary = one full page).
|
||||
|
||||
A single page may carry MORE THAN ONE distinct exhibit start — that is the whole
|
||||
reason the document index pointed us at this page. When a page carries multiple
|
||||
distinct exhibit starts, return ALL of them (one per list item, in the order they
|
||||
appear on the page).
|
||||
|
||||
The following keywords mark exhibit/section starts:
|
||||
* EXHIBIT
|
||||
* ATTACHMENT
|
||||
* ARTICLE
|
||||
* SECTION
|
||||
* AMENDMENT
|
||||
* SCHEDULE
|
||||
* ADDENDUM
|
||||
* APPENDIX
|
||||
* RIDER
|
||||
* PARTICIPATING PROVIDER AGREEMENT (and similar document-level titles)
|
||||
|
||||
[INCLUSION RULES]
|
||||
* Header can be multi-line — include all lines that are part of the header
|
||||
(e.g., provider/entity names, effective dates, sub-titles immediately under
|
||||
the main exhibit identifier).
|
||||
* Keep exact capitalization, punctuation (em-dashes, en-dashes, colons), and
|
||||
formatting as it appears on the page. Downstream code does literal-string
|
||||
matching against the page text, so any modification will break the match.
|
||||
* When multiple distinct exhibits start on the same page (e.g., short
|
||||
attachments back-to-back), return them as separate items in the list, in
|
||||
the order they appear top-to-bottom on the page.
|
||||
|
||||
[EXCLUSION RULES]
|
||||
* Do NOT include body text, paragraph content, or numbered sub-clauses
|
||||
("2.3 Confidentiality", "10.4 Billing Procedures").
|
||||
* Do NOT include page-footer repeats, docusign IDs, page numbers, contract
|
||||
reference numbers.
|
||||
* Do NOT include form-field labels ("Provider's Legal Name", "Authorized
|
||||
Signature").
|
||||
* Do NOT modify or shorten the header text — preserve it character-for-character.
|
||||
|
||||
[OUTPUT FORMAT]
|
||||
Return a JSON dictionary where:
|
||||
- Keys are boundary numbers as strings ("1", "2", "3", …)
|
||||
- Values are LISTS of header strings found in that chunk
|
||||
- If no headers are found in a chunk (rare, since the index pointed us here),
|
||||
return an empty list [] for that boundary.
|
||||
|
||||
Example output:
|
||||
{{"1": ["PROVIDER SERVICES AGREEMENT\\nSIGNATURE PAGE"],
|
||||
"2": ["EXHIBIT A - FEE SCHEDULE\\nState of California Medicaid\\nEffective 07/01/2025",
|
||||
"EXHIBIT B - DRUG FORMULARY\\nState of California Medicaid"],
|
||||
"3": ["Attachment C: Commercial-Exchange"]}}
|
||||
|
||||
{JSON_DICT_WITH_LISTS_FORMAT_INSTRUCTIONS}"""
|
||||
|
||||
|
||||
def EXHIBIT_HEADER_FROM_VERIFIED_PAGES(
|
||||
bundle: str,
|
||||
) -> Tuple[str, Callable[[str], dict]]:
|
||||
"""Returns ONLY dynamic content for the verified-pages header extraction.
|
||||
Call EXHIBIT_HEADER_FROM_VERIFIED_PAGES_INSTRUCTION() separately for the
|
||||
cached instruction.
|
||||
|
||||
Args:
|
||||
bundle: Concatenated full-page texts separated by ===BOUNDARY_X_HEADER===
|
||||
markers (one boundary per verified page).
|
||||
|
||||
Returns:
|
||||
Tuple of (prompt_string, parser_function) where parser expects a JSON dict.
|
||||
"""
|
||||
prompt = f"""[CONTEXT]
|
||||
Here are the verified-page chunks to analyze:
|
||||
{bundle.replace('"', "'")}
|
||||
|
||||
Briefly explain your reasoning, then return your final answer as a valid JSON dictionary."""
|
||||
|
||||
return (prompt, _json_dict_parser)
|
||||
|
||||
@@ -3189,8 +3288,7 @@ def EXHIBIT_HEADER_DEDUP_INSTRUCTION() -> str:
|
||||
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|.\n\n"
|
||||
"Keep only the first page occurrence of each unique section.\n\n"
|
||||
"[EFFECTIVE DATE RULE]\n"
|
||||
"An exhibit or section that reappears with a different effective date is still a duplicate of the first occurrence — "
|
||||
"a change in effective date alone does NOT create a distinct section. Drop the later-dated repetition and keep only the first."
|
||||
@@ -3228,12 +3326,95 @@ Your task is to deduplicate this dictionary so that each unique section appears
|
||||
* 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.
|
||||
Return a JSON dictionary 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|."""
|
||||
Briefly explain which headers you identified as duplicates and why, then return your final answer as a valid JSON dictionary.
|
||||
|
||||
{JSON_DICT_FORMAT_INSTRUCTIONS}"""
|
||||
|
||||
|
||||
def DOCUMENT_INDEX_INSTRUCTION() -> str:
|
||||
"""Static instruction for DOCUMENT_INDEX prompt caching.
|
||||
Returns formatted instructions for parsing a Document Index block.
|
||||
"""
|
||||
return f"""[OBJECTIVE]
|
||||
Parse a Textract-emitted Document Index block from a healthcare contract and classify each entry.
|
||||
|
||||
The Document Index is a table of contents that appears at the very start of the contract text.
|
||||
Each line typically contains a section title and a page number, e.g.:
|
||||
EXHIBIT A – FEE SCHEDULE 3
|
||||
Attachment B: Drug Formulary 12
|
||||
|
||||
[CLASSIFICATIONS]
|
||||
Assign each entry exactly one classification:
|
||||
- exhibit_start : A major section boundary — anything labelled EXHIBIT, ATTACHMENT, ARTICLE,
|
||||
SECTION, AMENDMENT, SCHEDULE, ADDENDUM, APPENDIX, RIDER, or a document-level title
|
||||
like "PARTICIPATING PROVIDER AGREEMENT". These are the primary chunking signals.
|
||||
Note: SECTION-prefixed entries qualify only when SECTION is the first word
|
||||
(e.g., "SECTION 5 - COMPENSATION"). Bare numeric sub-clauses without the
|
||||
SECTION word ("2.3 Confidentiality", "10.4 Billing") do NOT qualify here.
|
||||
- landmark : A named sub-section or titled chapter that is NOT one of the above but helps
|
||||
orient a reader (e.g., "DEFINITIONS", "GENERAL TERMS", "COMPENSATION METHODOLOGY").
|
||||
Numeric sub-clauses like "2.3 Confidentiality" or "10.4 Billing Procedures"
|
||||
belong here too — they are NOT exhibit boundaries.
|
||||
- noise : Lines that carry NO structural signal:
|
||||
* DocuSign envelope IDs (e.g., "DocuSign Envelope ID: xxxxxxxx-...")
|
||||
* Email thread headers (From:, To:, Sent:, Subject:)
|
||||
* IRS or tax forms (W-9, 1099, etc.)
|
||||
* Bare dates or timestamps
|
||||
* URLs or email addresses
|
||||
* Phone/fax numbers
|
||||
* Page-footer repeat lines (a short-form echo of an exhibit name near a page number)
|
||||
* Any line that is purely numeric or purely punctuation
|
||||
- unknown : Entry format is too ambiguous to classify confidently.
|
||||
|
||||
[RULES]
|
||||
1. Preserve the header text CHARACTER-FOR-CHARACTER as it appears in the index — downstream code
|
||||
does a literal string search for these values in the contract pages.
|
||||
2. Page numbers in the index refer to Textract page numbers (integer strings like "3", "12").
|
||||
Extract the rightmost integer on the line as the page number. If no integer is present, use null.
|
||||
3. A single index line may describe a multi-line header; treat it as one entry.
|
||||
4. Out-of-sequence page numbers are normal (e.g., Attachment A-1 on page 3, Attachment C-1 on page 8
|
||||
with no B-1 in between). Do NOT infer or fill gaps.
|
||||
5. If the index is empty or malformed, return an empty list [].
|
||||
|
||||
[OUTPUT FORMAT]
|
||||
Each element is an object with exactly these keys:
|
||||
"page" : integer or null — Textract page number where this section starts
|
||||
"header" : string — exact header text from the index
|
||||
"classification" : one of "exhibit_start", "landmark", "noise", "unknown"
|
||||
|
||||
Example:
|
||||
[
|
||||
{{"page": 1, "header": "PARTICIPATING PROVIDER AGREEMENT", "classification": "exhibit_start"}},
|
||||
{{"page": 3, "header": "EXHIBIT A – FEE SCHEDULE", "classification": "exhibit_start"}},
|
||||
{{"page": 8, "header": "DEFINITIONS", "classification": "landmark"}},
|
||||
{{"page": null, "header": "DocuSign Envelope ID: abc-123", "classification": "noise"}}
|
||||
]
|
||||
|
||||
{JSON_LIST_OF_DICTS_FORMAT_INSTRUCTIONS}"""
|
||||
|
||||
|
||||
def DOCUMENT_INDEX(index_block: str) -> Tuple[str, Callable[[str], list]]:
|
||||
"""
|
||||
Returns ONLY dynamic content for Document Index parsing.
|
||||
Call DOCUMENT_INDEX_INSTRUCTION() separately for the cached instruction.
|
||||
|
||||
Args:
|
||||
index_block: The raw Document Index text extracted from the contract prefix.
|
||||
|
||||
Returns:
|
||||
Tuple of (prompt_string, parser_function) where parser expects a JSON list.
|
||||
"""
|
||||
prompt = f"""[DOCUMENT INDEX]
|
||||
{index_block.replace('"', "'")}
|
||||
|
||||
Briefly explain your reasoning, then return your final answer as a properly formatted JSON list of dictionaries."""
|
||||
|
||||
return (prompt, _json_list_parser)
|
||||
|
||||
|
||||
#####################################################################################
|
||||
|
||||
@@ -0,0 +1,734 @@
|
||||
"""
|
||||
Unit tests for the Document Index preprocessing feature.
|
||||
|
||||
Covers:
|
||||
- extract_document_index_block: pure regex extraction
|
||||
- DOCUMENT_INDEX_INSTRUCTION / DOCUMENT_INDEX: contract-level invariants only
|
||||
- prompt_document_index: LLM wrapper (mocked)
|
||||
- verify_index_against_pages: structural verifier + quality gates + tolerance
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
import src.prompts.prompt_templates as prompt_templates
|
||||
from src.pipelines.saas.prompts import prompt_calls
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
extract_document_index_block,
|
||||
verify_index_against_pages,
|
||||
)
|
||||
|
||||
|
||||
class TestExtractDocumentIndexBlock(unittest.TestCase):
|
||||
"""Test the pure-regex Document Index block extractor."""
|
||||
|
||||
def test_present_returns_prefix(self):
|
||||
"""Returns the index block when the prefix and marker are both present."""
|
||||
text = (
|
||||
"Document Index\n"
|
||||
"EXHIBIT A - FEE SCHEDULE 3\n"
|
||||
"EXHIBIT B - PROVIDERS 8\n"
|
||||
"Start of Page No. = 1\n"
|
||||
"Page 1 content here..."
|
||||
)
|
||||
result = extract_document_index_block(text)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("Document Index", result)
|
||||
self.assertIn("EXHIBIT A - FEE SCHEDULE", result)
|
||||
self.assertIn("EXHIBIT B - PROVIDERS", result)
|
||||
self.assertNotIn("Start of Page No. =", result)
|
||||
self.assertNotIn("Page 1 content here", result)
|
||||
|
||||
def test_absent_marker_returns_none(self):
|
||||
"""Returns None when the 'Start of Page No. =' marker is absent."""
|
||||
text = "Document Index\nEXHIBIT A - FEE SCHEDULE 3\n"
|
||||
self.assertIsNone(extract_document_index_block(text))
|
||||
|
||||
def test_no_document_index_prefix_returns_none(self):
|
||||
"""Returns None when prefix doesn't start with 'Document Index'."""
|
||||
text = (
|
||||
"Some other random content here\n"
|
||||
"Start of Page No. = 1\n"
|
||||
"Page 1 content..."
|
||||
)
|
||||
self.assertIsNone(extract_document_index_block(text))
|
||||
|
||||
def test_whitespace_tolerant(self):
|
||||
"""Handles leading newlines, spaces, or blank lines before 'Document Index'."""
|
||||
text = (
|
||||
"\n\n Document Index\n"
|
||||
"EXHIBIT A 3\n"
|
||||
"Start of Page No. = 1\n"
|
||||
"Content..."
|
||||
)
|
||||
result = extract_document_index_block(text)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("EXHIBIT A", result)
|
||||
|
||||
def test_case_insensitive(self):
|
||||
"""Matches 'DOCUMENT INDEX' (all caps)."""
|
||||
text = "DOCUMENT INDEX\nEXHIBIT A 3\nStart of Page No. = 1\nContent..."
|
||||
result = extract_document_index_block(text)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("EXHIBIT A", result)
|
||||
|
||||
def test_empty_string_returns_none(self):
|
||||
"""Empty input returns None."""
|
||||
self.assertIsNone(extract_document_index_block(""))
|
||||
|
||||
|
||||
class TestDocumentIndexPromptTemplates(unittest.TestCase):
|
||||
"""Pin only the contract-level invariants of the prompt — not its copy."""
|
||||
|
||||
def test_instruction_contains_all_classification_values(self):
|
||||
"""All four classification enum values must be enumerated."""
|
||||
result = prompt_templates.DOCUMENT_INDEX_INSTRUCTION()
|
||||
self.assertIn("exhibit_start", result)
|
||||
self.assertIn("landmark", result)
|
||||
self.assertIn("noise", result)
|
||||
self.assertIn("unknown", result)
|
||||
|
||||
def test_instruction_has_no_pipe_delimiters(self):
|
||||
"""Project-wide rule: no |pipes| in LLM instruction strings."""
|
||||
result = prompt_templates.DOCUMENT_INDEX_INSTRUCTION()
|
||||
self.assertNotIn("|", result)
|
||||
|
||||
def test_document_index_registered_in_cacheable_instructions(self):
|
||||
"""Bedrock cache wiring contract — DOCUMENT_INDEX must be present."""
|
||||
cacheable = prompt_templates.get_cacheable_instructions()
|
||||
self.assertIn("DOCUMENT_INDEX", cacheable)
|
||||
self.assertIsInstance(cacheable["DOCUMENT_INDEX"], str)
|
||||
self.assertTrue(len(cacheable["DOCUMENT_INDEX"]) > 0)
|
||||
|
||||
|
||||
class TestPromptDocumentIndex(unittest.TestCase):
|
||||
"""Test the prompt_document_index LLM wrapper (with mocked invoke_claude)."""
|
||||
|
||||
def setUp(self):
|
||||
self.filename = "test_contract.txt"
|
||||
self.index_block = "Document Index\nEXHIBIT A - FEE SCHEDULE 3"
|
||||
|
||||
@patch("src.utils.llm_utils.invoke_claude")
|
||||
def test_calls_llm_with_usage_label(self, mock_invoke):
|
||||
"""invoke_claude is called with usage_label='DOCUMENT_INDEX_PARSE' — instrumentation contract."""
|
||||
mock_invoke.return_value = "[]"
|
||||
prompt_calls.prompt_document_index(self.index_block, self.filename)
|
||||
self.assertTrue(mock_invoke.called)
|
||||
self.assertEqual(
|
||||
mock_invoke.call_args.kwargs.get("usage_label"), "DOCUMENT_INDEX_PARSE"
|
||||
)
|
||||
|
||||
@patch("src.utils.llm_utils.invoke_claude")
|
||||
def test_calls_llm_with_cache_true(self, mock_invoke):
|
||||
"""cache=True is passed to invoke_claude — Bedrock prompt cache contract."""
|
||||
mock_invoke.return_value = "[]"
|
||||
prompt_calls.prompt_document_index(self.index_block, self.filename)
|
||||
self.assertTrue(mock_invoke.call_args.kwargs.get("cache"))
|
||||
|
||||
@patch("src.utils.llm_utils.invoke_claude")
|
||||
def test_returns_empty_list_on_malformed_json(self, mock_invoke):
|
||||
"""Malformed LLM output → returns [] without raising."""
|
||||
mock_invoke.return_value = "not valid json ~~~"
|
||||
try:
|
||||
result = prompt_calls.prompt_document_index(self.index_block, self.filename)
|
||||
except Exception as e:
|
||||
self.fail(
|
||||
f"prompt_document_index raised on malformed JSON: {type(e).__name__}: {e}"
|
||||
)
|
||||
self.assertEqual(result, [])
|
||||
|
||||
|
||||
class TestVerifyIndexAgainstPages(unittest.TestCase):
|
||||
"""Cover the verifier's quality gates, fuzzy matching, and tolerance."""
|
||||
|
||||
def _fake_page(self, text):
|
||||
page = MagicMock()
|
||||
page.get_text.return_value = text
|
||||
return page
|
||||
|
||||
# ── Quality gates ────────────────────────────────────────────────────────
|
||||
|
||||
def test_empty_parse_signals_fallback(self):
|
||||
pages_dict = {"1": self._fake_page("Some page text")}
|
||||
result = verify_index_against_pages([], pages_dict)
|
||||
self.assertEqual(result["verified_dict"], {})
|
||||
self.assertTrue(result["should_fallback"])
|
||||
self.assertEqual(result["fallback_reason"], "no_exhibit_starts")
|
||||
|
||||
def test_clean_match_does_not_fall_back(self):
|
||||
pages_dict = {
|
||||
"1": self._fake_page("EXHIBIT A - FEE SCHEDULE\nSome more page text"),
|
||||
}
|
||||
parsed = [
|
||||
{
|
||||
"page": 1,
|
||||
"header": "EXHIBIT A - FEE SCHEDULE",
|
||||
"classification": "exhibit_start",
|
||||
}
|
||||
]
|
||||
result = verify_index_against_pages(parsed, pages_dict)
|
||||
self.assertIn("1", result["verified_dict"])
|
||||
self.assertFalse(result["should_fallback"])
|
||||
self.assertEqual(result["metrics"]["verification_passes"], 1)
|
||||
|
||||
def test_fuzzy_match_handles_em_dash_variation(self):
|
||||
pages_dict = {"1": self._fake_page("EXHIBIT A - FEE SCHEDULE\nbody text here")}
|
||||
parsed = [
|
||||
{
|
||||
"page": 1,
|
||||
"header": "EXHIBIT A — FEE SCHEDULE", # em-dash
|
||||
"classification": "exhibit_start",
|
||||
}
|
||||
]
|
||||
result = verify_index_against_pages(parsed, pages_dict)
|
||||
self.assertIn("1", result["verified_dict"])
|
||||
self.assertFalse(result["should_fallback"])
|
||||
|
||||
def test_high_failure_ratio_triggers_fallback(self):
|
||||
"""4 of 5 entries fail → 80% > 10% → fallback."""
|
||||
pages_dict = {
|
||||
"1": self._fake_page("EXHIBIT A - FEE SCHEDULE\ncontent"),
|
||||
"2": self._fake_page("Random unrelated content zzz"),
|
||||
"3": self._fake_page("Random unrelated content yyy"),
|
||||
"4": self._fake_page("Random unrelated content xxx"),
|
||||
"5": self._fake_page("Random unrelated content www"),
|
||||
}
|
||||
parsed = [
|
||||
{
|
||||
"page": 1,
|
||||
"header": "EXHIBIT A - FEE SCHEDULE",
|
||||
"classification": "exhibit_start",
|
||||
},
|
||||
{"page": 2, "header": "QQQQ TWO", "classification": "exhibit_start"},
|
||||
{"page": 3, "header": "QQQQ THREE", "classification": "exhibit_start"},
|
||||
{"page": 4, "header": "QQQQ FOUR", "classification": "exhibit_start"},
|
||||
{"page": 5, "header": "QQQQ FIVE", "classification": "exhibit_start"},
|
||||
]
|
||||
result = verify_index_against_pages(parsed, pages_dict)
|
||||
self.assertTrue(result["should_fallback"])
|
||||
self.assertTrue(result["fallback_reason"].startswith("failure_ratio_"))
|
||||
|
||||
def test_low_failure_ratio_does_not_trigger_fallback(self):
|
||||
"""1 of 10 = 10% — not strictly greater, so no fallback."""
|
||||
pages_dict = {
|
||||
str(i): self._fake_page(f"EXHIBIT {chr(64+i)} HEADER\ncontent")
|
||||
for i in range(1, 11)
|
||||
}
|
||||
parsed = [
|
||||
{
|
||||
"page": i,
|
||||
"header": f"EXHIBIT {chr(64+i)} HEADER",
|
||||
"classification": "exhibit_start",
|
||||
}
|
||||
for i in range(1, 10)
|
||||
] + [
|
||||
{
|
||||
"page": 10,
|
||||
"header": "EXHIBIT ZNONEXISTENT XYZ",
|
||||
"classification": "exhibit_start",
|
||||
}
|
||||
]
|
||||
result = verify_index_against_pages(parsed, pages_dict)
|
||||
self.assertFalse(result["should_fallback"])
|
||||
|
||||
def test_first_exhibit_on_page_2_does_not_fall_back(self):
|
||||
"""Pre-first-exhibit pages are normal (cover page / DI block); no fallback."""
|
||||
pages_dict = {
|
||||
"1": self._fake_page("Cover page intro content"),
|
||||
"2": self._fake_page("EXHIBIT A - FEE SCHEDULE\ncontent"),
|
||||
}
|
||||
parsed = [
|
||||
{
|
||||
"page": 2,
|
||||
"header": "EXHIBIT A - FEE SCHEDULE",
|
||||
"classification": "exhibit_start",
|
||||
}
|
||||
]
|
||||
result = verify_index_against_pages(parsed, pages_dict)
|
||||
self.assertFalse(result["should_fallback"])
|
||||
|
||||
def test_tail_coverage_gap_triggers_fallback(self):
|
||||
"""Last exhibit_start at p10 of a 50-page doc → 40-page tail → fallback."""
|
||||
pages_dict = {
|
||||
str(i): self._fake_page(f"page {i} content") for i in range(1, 51)
|
||||
}
|
||||
pages_dict["1"] = self._fake_page("EXHIBIT A - FEE SCHEDULE\ncontent")
|
||||
pages_dict["10"] = self._fake_page("EXHIBIT B - HOSPITAL\ncontent")
|
||||
parsed = [
|
||||
{
|
||||
"page": 1,
|
||||
"header": "EXHIBIT A - FEE SCHEDULE",
|
||||
"classification": "exhibit_start",
|
||||
},
|
||||
{
|
||||
"page": 10,
|
||||
"header": "EXHIBIT B - HOSPITAL",
|
||||
"classification": "exhibit_start",
|
||||
},
|
||||
]
|
||||
result = verify_index_against_pages(parsed, pages_dict)
|
||||
self.assertTrue(result["should_fallback"])
|
||||
self.assertIn("tail_coverage_gap", result["fallback_reason"])
|
||||
|
||||
def test_small_tail_does_not_trigger_fallback(self):
|
||||
"""Last exhibit at p47 of 50p doc → 3-page tail → no fallback."""
|
||||
pages_dict = {
|
||||
str(i): self._fake_page(f"page {i} content") for i in range(1, 51)
|
||||
}
|
||||
pages_dict["1"] = self._fake_page("EXHIBIT A\ncontent")
|
||||
pages_dict["47"] = self._fake_page("EXHIBIT B\ncontent")
|
||||
parsed = [
|
||||
{"page": 1, "header": "EXHIBIT A", "classification": "exhibit_start"},
|
||||
{"page": 47, "header": "EXHIBIT B", "classification": "exhibit_start"},
|
||||
]
|
||||
result = verify_index_against_pages(parsed, pages_dict)
|
||||
self.assertFalse(result["should_fallback"])
|
||||
|
||||
def test_short_doc_skips_tail_check(self):
|
||||
"""Doc with total_pages <= 15 skips the tail-coverage check."""
|
||||
pages_dict = {str(i): self._fake_page(f"page {i}") for i in range(1, 11)}
|
||||
pages_dict["1"] = self._fake_page("EXHIBIT A\ncontent")
|
||||
parsed = [{"page": 1, "header": "EXHIBIT A", "classification": "exhibit_start"}]
|
||||
result = verify_index_against_pages(parsed, pages_dict)
|
||||
self.assertFalse(result["should_fallback"])
|
||||
|
||||
# ── Classification handling ──────────────────────────────────────────────
|
||||
|
||||
def test_only_exhibit_start_entries_land_in_verified_dict(self):
|
||||
"""Landmark and other classifications must NOT appear in verified_dict,
|
||||
even when the text matches on the claimed page."""
|
||||
pages = {
|
||||
"1": self._fake_page("EXHIBIT A\ncontent"),
|
||||
"2": self._fake_page("SECTION 1 GENERAL TERMS\ncontent"),
|
||||
"3": self._fake_page("RANDOM HEADER\ncontent"),
|
||||
}
|
||||
parsed = [
|
||||
{"page": 1, "header": "EXHIBIT A", "classification": "exhibit_start"},
|
||||
{
|
||||
"page": 2,
|
||||
"header": "SECTION 1 GENERAL TERMS",
|
||||
"classification": "landmark",
|
||||
},
|
||||
{
|
||||
"page": 3,
|
||||
"header": "RANDOM HEADER",
|
||||
"classification": "other_thing", # any non-exhibit_start
|
||||
},
|
||||
]
|
||||
result = verify_index_against_pages(parsed, pages)
|
||||
all_headers = [h for hs in result["verified_dict"].values() for h in hs]
|
||||
self.assertEqual(all_headers, ["EXHIBIT A"])
|
||||
|
||||
def test_subpage_key_resolves_to_base_page(self):
|
||||
"""When the index says page=3 but the dict has '3.1', the verifier resolves it."""
|
||||
pages = {"3.1": self._fake_page("EXHIBIT A – PHYSICIAN FEES\ntext")}
|
||||
parsed = [
|
||||
{
|
||||
"page": 3,
|
||||
"header": "EXHIBIT A – PHYSICIAN FEES",
|
||||
"classification": "exhibit_start",
|
||||
}
|
||||
]
|
||||
result = verify_index_against_pages(parsed, pages)
|
||||
all_headers = [h for hs in result["verified_dict"].values() for h in hs]
|
||||
self.assertIn("EXHIBIT A – PHYSICIAN FEES", all_headers)
|
||||
|
||||
# ── Tolerance: never raise on garbage entries ────────────────────────────
|
||||
|
||||
def test_malformed_entries_do_not_crash_or_pollute_output(self):
|
||||
"""Single sweep test: missing-header, empty-header, missing-page, missing-classification.
|
||||
Verifier must not raise, and no garbage entry may end up in verified_dict."""
|
||||
pages = {"1": self._fake_page("EXHIBIT A\ncontent")}
|
||||
parsed = [
|
||||
{"page": 1, "classification": "exhibit_start"}, # missing header
|
||||
{"page": 1, "header": "", "classification": "exhibit_start"}, # empty
|
||||
{"header": "EXHIBIT A", "classification": "exhibit_start"}, # missing page
|
||||
{"page": 1, "header": "EXHIBIT A"}, # missing classification
|
||||
]
|
||||
try:
|
||||
result = verify_index_against_pages(parsed, pages)
|
||||
except Exception as exc:
|
||||
self.fail(f"verifier raised on malformed entries: {exc}")
|
||||
self.assertIsInstance(result, dict)
|
||||
all_headers = [h for hs in result["verified_dict"].values() for h in hs]
|
||||
# Empty strings and missing-classification entries must NOT appear
|
||||
self.assertNotIn("", all_headers)
|
||||
# The missing-classification entry has a valid header but no class —
|
||||
# treating it as exhibit_start would be unsafe, so it must be excluded
|
||||
self.assertEqual(all_headers, [])
|
||||
|
||||
def test_returns_dict_with_expected_keys(self):
|
||||
"""API contract: response shape is stable."""
|
||||
pages_dict = {"1": self._fake_page("Some text")}
|
||||
result = verify_index_against_pages([], pages_dict)
|
||||
self.assertEqual(
|
||||
set(result.keys()),
|
||||
{"verified_dict", "metrics", "should_fallback", "fallback_reason"},
|
||||
)
|
||||
|
||||
|
||||
class TestReplaceIndexHeadersWithPageText(unittest.TestCase):
|
||||
"""Cover the page-text header replacement helper.
|
||||
|
||||
The helper takes a thin DI verified_dict and rewrites each header with the
|
||||
rich page-text version returned by the LLM. On any failure it must preserve
|
||||
the original verified_dict so the pipeline keeps working.
|
||||
"""
|
||||
|
||||
def _fake_page(self, text):
|
||||
page = MagicMock()
|
||||
page.get_text.return_value = text
|
||||
return page
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.prompt_calls.prompt_exhibit_header_from_verified_pages"
|
||||
)
|
||||
def test_two_headers_on_one_page_both_get_rewritten(self, mock_prompt):
|
||||
"""When DI flagged two exhibits on a single page and the LLM returns a
|
||||
list of rich headers for that boundary, both get rewritten and order
|
||||
preserved — this is the core DI use case (multi-exhibit page)."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
replace_index_headers_with_page_text,
|
||||
)
|
||||
|
||||
verified_dict = {"1": ["Exhibit A", "Exhibit B"]}
|
||||
pages_dict = {
|
||||
"1": self._fake_page(
|
||||
"EXHIBIT A — FEE SCHEDULE\nState of CA Medicaid\n"
|
||||
"[body content]\n"
|
||||
"EXHIBIT B — DRUG FORMULARY\nState of CA Medicaid"
|
||||
)
|
||||
}
|
||||
mock_prompt.return_value = {
|
||||
"1": [
|
||||
"EXHIBIT A — FEE SCHEDULE\nState of CA Medicaid",
|
||||
"EXHIBIT B — DRUG FORMULARY\nState of CA Medicaid",
|
||||
]
|
||||
}
|
||||
|
||||
result = replace_index_headers_with_page_text(
|
||||
verified_dict, pages_dict, "test.txt"
|
||||
)
|
||||
self.assertEqual(len(result["1"]), 2)
|
||||
self.assertIn("EXHIBIT A — FEE SCHEDULE", result["1"][0])
|
||||
self.assertIn("EXHIBIT B — DRUG FORMULARY", result["1"][1])
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.prompt_calls.prompt_exhibit_header_from_verified_pages"
|
||||
)
|
||||
def test_llm_failure_preserves_original_verified_dict(self, mock_prompt):
|
||||
"""Empty wrapper response (LLM/parse failure) → verified_dict returned unchanged."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
replace_index_headers_with_page_text,
|
||||
)
|
||||
|
||||
mock_prompt.return_value = {}
|
||||
|
||||
verified_dict = {"1": ["EXHIBIT A"], "5": ["EXHIBIT B"]}
|
||||
pages_dict = {
|
||||
"1": self._fake_page("EXHIBIT A content"),
|
||||
"5": self._fake_page("EXHIBIT B content"),
|
||||
}
|
||||
result = replace_index_headers_with_page_text(
|
||||
verified_dict, pages_dict, "test.txt"
|
||||
)
|
||||
self.assertEqual(result, verified_dict)
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.prompt_calls.prompt_exhibit_header_from_verified_pages"
|
||||
)
|
||||
def test_thin_index_header_replaced_with_rich_page_header(self, mock_prompt):
|
||||
"""Happy path: thin TOC entry rewritten to multi-line page-text version."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
replace_index_headers_with_page_text,
|
||||
)
|
||||
|
||||
verified_dict = {"3": ["EXHIBIT A - FEE SCHEDULE"]}
|
||||
pages_dict = {
|
||||
"3": self._fake_page(
|
||||
"EXHIBIT A — FEE SCHEDULE\nPHYSICIAN COMPENSATION\n"
|
||||
"State of California Medicaid\nEffective 07/01/2025"
|
||||
)
|
||||
}
|
||||
mock_prompt.return_value = {
|
||||
"1": [
|
||||
"EXHIBIT A — FEE SCHEDULE\nPHYSICIAN COMPENSATION\n"
|
||||
"State of California Medicaid\nEffective 07/01/2025"
|
||||
]
|
||||
}
|
||||
|
||||
result = replace_index_headers_with_page_text(
|
||||
verified_dict, pages_dict, "test.txt"
|
||||
)
|
||||
self.assertEqual(len(result["3"]), 1)
|
||||
self.assertIn("PHYSICIAN COMPENSATION", result["3"][0])
|
||||
self.assertIn("State of California Medicaid", result["3"][0])
|
||||
# Original thin header replaced
|
||||
self.assertNotEqual(result["3"][0], "EXHIBIT A - FEE SCHEDULE")
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.prompt_calls.prompt_exhibit_header_from_verified_pages"
|
||||
)
|
||||
def test_running_header_lines_stripped_before_bundling(self, mock_prompt):
|
||||
"""Fix B: a line that repeats on every page (e.g. "PROVIDER SERVICES
|
||||
AGREEMENT") must not reach the LLM bundle — otherwise the LLM concatenates
|
||||
it into the rewritten exhibit title."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
replace_index_headers_with_page_text,
|
||||
)
|
||||
|
||||
# 5 pages, every one carries "PROVIDER SERVICES AGREEMENT" at the top.
|
||||
# That makes it a document running header (>40% of pages).
|
||||
running = "PROVIDER SERVICES AGREEMENT"
|
||||
pages_dict = {
|
||||
"1": self._fake_page(f"{running}\nCover page body\n"),
|
||||
"2": self._fake_page(f"{running}\nRecitals body\n"),
|
||||
"3": self._fake_page(f"{running}\nMore recitals\n"),
|
||||
"4": self._fake_page(f"{running}\nDefinitions body\n"),
|
||||
"5": self._fake_page(
|
||||
f"{running}\nExhibit 1\nServices & Compensation\n[body]"
|
||||
),
|
||||
}
|
||||
verified_dict = {"5": ["Exhibit 1"]}
|
||||
mock_prompt.return_value = {"1": ["Exhibit 1\nServices & Compensation"]}
|
||||
|
||||
replace_index_headers_with_page_text(verified_dict, pages_dict, "test.txt")
|
||||
|
||||
bundle_arg = mock_prompt.call_args[0][0]
|
||||
self.assertNotIn(
|
||||
running,
|
||||
bundle_arg,
|
||||
"Running-header line must be stripped from the LLM bundle",
|
||||
)
|
||||
# Real header content must still survive
|
||||
self.assertIn("Exhibit 1", bundle_arg)
|
||||
self.assertIn("Services & Compensation", bundle_arg)
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.prompt_calls.prompt_exhibit_header_from_verified_pages"
|
||||
)
|
||||
def test_explicit_running_headers_used_when_passed(self, mock_prompt):
|
||||
"""When the orchestrator passes a pre-computed running_headers set,
|
||||
the rewriter uses it directly without recomputing on pages_dict —
|
||||
this is the lifted-set path used in production."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
replace_index_headers_with_page_text,
|
||||
)
|
||||
|
||||
pages_dict = {
|
||||
"1": self._fake_page("Hackensack Meridian Health, Inc.\nEXHIBIT A\nFEES")
|
||||
}
|
||||
verified_dict = {"1": ["EXHIBIT A"]}
|
||||
mock_prompt.return_value = {"1": ["EXHIBIT A\nFEES"]}
|
||||
|
||||
replace_index_headers_with_page_text(
|
||||
verified_dict,
|
||||
pages_dict,
|
||||
"test.txt",
|
||||
running_headers={"Hackensack Meridian Health, Inc."},
|
||||
)
|
||||
|
||||
bundle_arg = mock_prompt.call_args[0][0]
|
||||
self.assertNotIn("Hackensack Meridian Health, Inc.", bundle_arg)
|
||||
self.assertIn("EXHIBIT A", bundle_arg)
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.prompt_calls.prompt_exhibit_header_from_verified_pages"
|
||||
)
|
||||
def test_no_running_headers_page_text_passes_through(self, mock_prompt):
|
||||
"""When the document has no running headers, every page line reaches
|
||||
the LLM unchanged — the strip is a no-op."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
replace_index_headers_with_page_text,
|
||||
)
|
||||
|
||||
pages_dict = {
|
||||
"1": self._fake_page("EXHIBIT A\nState of California Medicaid\nbody")
|
||||
}
|
||||
verified_dict = {"1": ["EXHIBIT A"]}
|
||||
mock_prompt.return_value = {"1": ["EXHIBIT A\nState of California Medicaid"]}
|
||||
|
||||
replace_index_headers_with_page_text(
|
||||
verified_dict, pages_dict, "test.txt", running_headers=set()
|
||||
)
|
||||
|
||||
bundle_arg = mock_prompt.call_args[0][0]
|
||||
self.assertIn("State of California Medicaid", bundle_arg)
|
||||
self.assertIn("body", bundle_arg)
|
||||
|
||||
|
||||
class TestRecoverOrphanAnchorPages(unittest.TestCase):
|
||||
"""Cover the regex-based recovery of pages where DI missed the anchor.
|
||||
|
||||
Some TOCs drop the leading "ATTACHMENT X" / "Exhibit N" label, leaving
|
||||
only the secondary title. The DI LLM correctly classifies the orphan
|
||||
secondary as `landmark`, so the page never enters verified_dict. This
|
||||
helper scans every unverified page for a strict whole-line anchor and
|
||||
promotes it. No new LLM call.
|
||||
"""
|
||||
|
||||
def _fake_page(self, text):
|
||||
page = MagicMock()
|
||||
page.get_text.return_value = text
|
||||
return page
|
||||
|
||||
def test_anchor_at_top_of_page_recovered(self):
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
recover_orphan_anchor_pages,
|
||||
)
|
||||
|
||||
verified_dict = {"1": ["EXHIBIT A"]}
|
||||
pages_dict = {
|
||||
"1": self._fake_page("EXHIBIT A\nFEE SCHEDULE\nbody text"),
|
||||
"5": self._fake_page("ATTACHMENT B\nMEDICARE ADVANTAGE COMPENSATION\nbody"),
|
||||
}
|
||||
result = recover_orphan_anchor_pages(verified_dict, pages_dict, "test.txt")
|
||||
self.assertIn("5", result)
|
||||
self.assertEqual(result["5"], ["ATTACHMENT B"])
|
||||
# Original entries untouched
|
||||
self.assertEqual(result["1"], ["EXHIBIT A"])
|
||||
|
||||
def test_anchor_mid_page_recovered(self):
|
||||
"""Header may appear after several body lines on the same page."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
recover_orphan_anchor_pages,
|
||||
)
|
||||
|
||||
verified_dict = {}
|
||||
page_text = (
|
||||
"continuation of previous exhibit\n"
|
||||
"some body sentence here\n"
|
||||
"another body sentence\n"
|
||||
"ATTACHMENT C\n"
|
||||
"PROVIDER INFORMATION\n"
|
||||
)
|
||||
pages_dict = {"7": self._fake_page(page_text)}
|
||||
result = recover_orphan_anchor_pages(verified_dict, pages_dict, "test.txt")
|
||||
self.assertEqual(result["7"], ["ATTACHMENT C"])
|
||||
|
||||
def test_body_sentence_starting_with_exhibit_not_recovered(self):
|
||||
"""A sentence like 'As required by Exhibit 3 of this Agreement...' must NOT match."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
recover_orphan_anchor_pages,
|
||||
)
|
||||
|
||||
page_text = (
|
||||
"PROVIDER SERVICES AGREEMENT\n"
|
||||
"Payor's affiliates are those entities controlling, controlled by, or under common control with Payor who will\n"
|
||||
"be bound by the terms and conditions of this Agreement.\n"
|
||||
"Exhibit 3 to this Agreement. Payor shall notify DaVita of any newly created affiliated Payor entities.\n"
|
||||
)
|
||||
pages_dict = {"4": self._fake_page(page_text)}
|
||||
result = recover_orphan_anchor_pages({}, pages_dict, "test.txt")
|
||||
self.assertNotIn("4", result)
|
||||
|
||||
def test_running_header_excluded(self):
|
||||
"""A line appearing on >40% of pages is a running header — never promote it."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
recover_orphan_anchor_pages,
|
||||
)
|
||||
|
||||
# 5 pages all starting with "EXHIBIT A" — that's a document running
|
||||
# header, not 5 exhibits. The repetition guard must reject all of them.
|
||||
pages_dict = {
|
||||
str(i): self._fake_page("EXHIBIT A\nbody content\nmore body\n")
|
||||
for i in range(1, 6)
|
||||
}
|
||||
result = recover_orphan_anchor_pages({}, pages_dict, "test.txt")
|
||||
self.assertEqual(result, {})
|
||||
|
||||
def test_long_line_with_anchor_at_start_rejected(self):
|
||||
"""Lines >80 chars are body sentences, not headers; reject even if anchor leads."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
recover_orphan_anchor_pages,
|
||||
)
|
||||
|
||||
long_line = (
|
||||
"EXHIBIT A - A really long body sentence that goes on and on with body details "
|
||||
"that should clearly NOT be classified as a header"
|
||||
)
|
||||
pages_dict = {"3": self._fake_page(long_line + "\n\nother content")}
|
||||
result = recover_orphan_anchor_pages({}, pages_dict, "test.txt")
|
||||
self.assertNotIn("3", result)
|
||||
|
||||
def test_no_space_attachment_form_recovered(self):
|
||||
"""OCR-mangled ATTACHMENTC (no space) is recovered."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
recover_orphan_anchor_pages,
|
||||
)
|
||||
|
||||
pages_dict = {"6": self._fake_page("ATTACHMENTC\nPROVIDERINFORMATION")}
|
||||
result = recover_orphan_anchor_pages({}, pages_dict, "test.txt")
|
||||
self.assertEqual(result["6"], ["ATTACHMENTC"])
|
||||
|
||||
def test_sub_page_chunk_skipped(self):
|
||||
"""Sub-page table chunks (e.g. '3.1') are not independent pages — skip them."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
recover_orphan_anchor_pages,
|
||||
)
|
||||
|
||||
pages_dict = {
|
||||
"3.1": self._fake_page("ATTACHMENT B\nbody"),
|
||||
}
|
||||
result = recover_orphan_anchor_pages({}, pages_dict, "test.txt")
|
||||
self.assertEqual(result, {})
|
||||
|
||||
def test_already_verified_page_not_duplicated(self):
|
||||
"""A page already in verified_dict must be left unchanged."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
recover_orphan_anchor_pages,
|
||||
)
|
||||
|
||||
verified_dict = {"1": ["EXHIBIT A - FEE SCHEDULE"]}
|
||||
pages_dict = {"1": self._fake_page("EXHIBIT A\nFEE SCHEDULE\nbody")}
|
||||
result = recover_orphan_anchor_pages(verified_dict, pages_dict, "test.txt")
|
||||
self.assertEqual(result, verified_dict)
|
||||
|
||||
def test_list_of_attachments_in_body_not_recovered(self):
|
||||
"""A page enumerating multiple attachments (Section 12.14-style 'Entire
|
||||
Agreement' clause) is NOT a real exhibit — skip it."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
recover_orphan_anchor_pages,
|
||||
)
|
||||
|
||||
page_text = (
|
||||
"12.14 Entire Agreement. This Agreement including all Attachments\n"
|
||||
"listed below, is the entire agreement between the Parties.\n"
|
||||
"Attachment A\n"
|
||||
"- Medicare Advantage Program Attachment\n"
|
||||
"Attachment B\n"
|
||||
"- Compensation\n"
|
||||
"Attachment C\n"
|
||||
"- Contracted Provider and Health Care Practitioner Locations\n"
|
||||
"Attachment D\n"
|
||||
"- Certification Regarding Lobbying\n"
|
||||
)
|
||||
pages_dict = {"26": self._fake_page(page_text)}
|
||||
result = recover_orphan_anchor_pages({}, pages_dict, "test.txt")
|
||||
self.assertNotIn("26", result)
|
||||
|
||||
def test_continuation_page_same_anchor_as_previous_not_recovered(self):
|
||||
"""A page whose top anchor matches the previous verified page's anchor
|
||||
is a multi-page exhibit continuation (e.g. table page), not a new boundary."""
|
||||
from src.pipelines.shared.preprocessing.preprocessing_funcs import (
|
||||
recover_orphan_anchor_pages,
|
||||
)
|
||||
|
||||
verified_dict = {
|
||||
"35": ["ATTACHMENT C"],
|
||||
}
|
||||
pages_dict = {
|
||||
"35": self._fake_page("ATTACHMENT C\nPROVIDER LOCATIONS\n[content]"),
|
||||
# p36 starts with ATTACHMENT C too — but this is the continuation
|
||||
# table of the same exhibit, not a new exhibit.
|
||||
"36": self._fake_page("ATTACHMENT C\nAMEDISYS NEW JERSEY, LLC\n[table]"),
|
||||
"37": self._fake_page("ATTACHMENT D\nCERTIFICATION REGARDING LOBBYING"),
|
||||
}
|
||||
result = recover_orphan_anchor_pages(verified_dict, pages_dict, "test.txt")
|
||||
# p36 must NOT be added (continuation), but p37 must (new exhibit)
|
||||
self.assertNotIn("36", result)
|
||||
self.assertEqual(result["37"], ["ATTACHMENT D"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -257,6 +257,235 @@ class TestPreprocessingFuncs:
|
||||
assert "Table 1" in result or "-------Table Start--------" in result
|
||||
|
||||
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
|
||||
class TestOneToNExhibitChunkingHybrid(unittest.TestCase):
|
||||
"""Integration tests for the Document Index path in one_to_n_exhibit_chunking.
|
||||
|
||||
The DI path is the primary system: if the contract has a Document Index block,
|
||||
we always trust the LLM parse. We only fall back to the legacy per-page path
|
||||
when the block is missing entirely.
|
||||
"""
|
||||
|
||||
def setUp(self):
|
||||
page_1 = MagicMock()
|
||||
page_1.get_text.return_value = "EXHIBIT A - FEE SCHEDULE\nPage 1 body text"
|
||||
page_2 = MagicMock()
|
||||
page_2.get_text.return_value = "EXHIBIT B - PROVIDERS\nPage 2 body text"
|
||||
self.fake_pages = {"1": page_1, "2": page_2}
|
||||
self.filename = "test.txt"
|
||||
self.contract_text = (
|
||||
"Document Index\n"
|
||||
"EXHIBIT A - FEE SCHEDULE 1\n"
|
||||
"EXHIBIT B - PROVIDERS 2\n"
|
||||
"Start of Page No. = 1\n"
|
||||
"Page 1 body..."
|
||||
)
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.get_exhibit_pages_new"
|
||||
)
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.replace_index_headers_with_page_text"
|
||||
)
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.recover_orphan_anchor_pages"
|
||||
)
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.verify_index_against_pages"
|
||||
)
|
||||
@patch("src.pipelines.saas.prompts.prompt_calls.prompt_document_index")
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.extract_document_index_block"
|
||||
)
|
||||
def test_di_block_present_skips_per_page_llm(
|
||||
self,
|
||||
mock_extract,
|
||||
mock_prompt_doc_idx,
|
||||
mock_verify,
|
||||
mock_recover,
|
||||
mock_replace,
|
||||
mock_per_page,
|
||||
):
|
||||
"""DI block exists → DI path returns verified_dict; per-page LLM never runs."""
|
||||
from src.pipelines.shared.preprocessing.preprocess import (
|
||||
one_to_n_exhibit_chunking,
|
||||
)
|
||||
|
||||
mock_extract.return_value = "Document Index\nEXHIBIT A - FEE SCHEDULE 1"
|
||||
mock_prompt_doc_idx.return_value = [
|
||||
{"page": 1, "header": "EXHIBIT A", "classification": "exhibit_start"}
|
||||
]
|
||||
mock_verify.return_value = {
|
||||
"verified_dict": {"1": ["EXHIBIT A"]},
|
||||
"metrics": {},
|
||||
"should_fallback": False,
|
||||
"fallback_reason": "",
|
||||
}
|
||||
# Passthrough — this test verifies the per-page LLM doesn't run, not
|
||||
# the orphan-anchor recovery or header-rewriting behavior. Return the
|
||||
# verified_dict unchanged through both downstream helpers.
|
||||
mock_recover.side_effect = lambda vd, pd, fn, **kwargs: vd
|
||||
mock_replace.side_effect = lambda vd, pd, fn, **kwargs: vd
|
||||
|
||||
with patch("src.config.DOCUMENT_INDEX_PARSE_ENABLED", True):
|
||||
result = one_to_n_exhibit_chunking(
|
||||
pages_dict=self.fake_pages,
|
||||
EXHIBIT_HEADER_MARKERS=[],
|
||||
filename=self.filename,
|
||||
contract_text=self.contract_text,
|
||||
)
|
||||
|
||||
self.assertFalse(mock_per_page.called)
|
||||
self.assertEqual(result, {"1": ["EXHIBIT A"]})
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.get_exhibit_pages_new"
|
||||
)
|
||||
@patch("src.pipelines.saas.prompts.prompt_calls.prompt_header_deduplication")
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.verify_index_against_pages"
|
||||
)
|
||||
@patch("src.pipelines.saas.prompts.prompt_calls.prompt_document_index")
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.extract_document_index_block"
|
||||
)
|
||||
def test_di_quality_gate_failure_runs_legacy_per_page_path(
|
||||
self,
|
||||
mock_extract,
|
||||
mock_prompt_doc_idx,
|
||||
mock_verify,
|
||||
mock_dedup,
|
||||
mock_per_page,
|
||||
):
|
||||
"""DI block exists but quality gate fails → legacy per-page path runs."""
|
||||
from src.pipelines.shared.preprocessing.preprocess import (
|
||||
one_to_n_exhibit_chunking,
|
||||
)
|
||||
|
||||
mock_extract.return_value = "Document Index\n(garbage)"
|
||||
mock_prompt_doc_idx.return_value = []
|
||||
mock_verify.return_value = {
|
||||
"verified_dict": {},
|
||||
"metrics": {},
|
||||
"should_fallback": True,
|
||||
"fallback_reason": "no_exhibit_starts",
|
||||
}
|
||||
mock_per_page.return_value = ({}, {"1": ["EXHIBIT A"]})
|
||||
mock_dedup.side_effect = lambda d, _filename: d
|
||||
|
||||
with patch("src.config.DOCUMENT_INDEX_PARSE_ENABLED", True):
|
||||
one_to_n_exhibit_chunking(
|
||||
pages_dict=self.fake_pages,
|
||||
EXHIBIT_HEADER_MARKERS=[],
|
||||
filename=self.filename,
|
||||
contract_text=self.contract_text,
|
||||
)
|
||||
|
||||
self.assertTrue(mock_per_page.called)
|
||||
self.assertTrue(mock_dedup.called)
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.get_exhibit_pages_new"
|
||||
)
|
||||
@patch("src.pipelines.saas.prompts.prompt_calls.prompt_header_deduplication")
|
||||
@patch("src.pipelines.saas.prompts.prompt_calls.prompt_document_index")
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.extract_document_index_block"
|
||||
)
|
||||
def test_no_di_block_runs_legacy_per_page_path(
|
||||
self,
|
||||
mock_extract,
|
||||
mock_prompt_doc_idx,
|
||||
mock_dedup,
|
||||
mock_per_page,
|
||||
):
|
||||
"""No DI block → per-page LLM runs (with dedup)."""
|
||||
from src.pipelines.shared.preprocessing.preprocess import (
|
||||
one_to_n_exhibit_chunking,
|
||||
)
|
||||
|
||||
mock_extract.return_value = None
|
||||
mock_per_page.return_value = ({}, {"1": ["EXHIBIT A"]})
|
||||
mock_dedup.side_effect = lambda d, _filename: d
|
||||
|
||||
with patch("src.config.DOCUMENT_INDEX_PARSE_ENABLED", True):
|
||||
one_to_n_exhibit_chunking(
|
||||
pages_dict=self.fake_pages,
|
||||
EXHIBIT_HEADER_MARKERS=[],
|
||||
filename=self.filename,
|
||||
contract_text=self.contract_text,
|
||||
)
|
||||
|
||||
self.assertFalse(mock_prompt_doc_idx.called)
|
||||
self.assertTrue(mock_per_page.called)
|
||||
self.assertTrue(mock_dedup.called)
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.get_exhibit_pages_new"
|
||||
)
|
||||
@patch("src.pipelines.saas.prompts.prompt_calls.prompt_header_deduplication")
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.extract_document_index_block"
|
||||
)
|
||||
def test_feature_flag_off_skips_di_path(
|
||||
self,
|
||||
mock_extract,
|
||||
mock_dedup,
|
||||
mock_per_page,
|
||||
):
|
||||
from src.pipelines.shared.preprocessing.preprocess import (
|
||||
one_to_n_exhibit_chunking,
|
||||
)
|
||||
|
||||
mock_per_page.return_value = ({}, {"1": ["EXHIBIT A"]})
|
||||
mock_dedup.side_effect = lambda d, _filename: d
|
||||
|
||||
with patch("src.config.DOCUMENT_INDEX_PARSE_ENABLED", False):
|
||||
one_to_n_exhibit_chunking(
|
||||
pages_dict=self.fake_pages,
|
||||
EXHIBIT_HEADER_MARKERS=[],
|
||||
filename=self.filename,
|
||||
contract_text=self.contract_text,
|
||||
)
|
||||
|
||||
self.assertFalse(mock_extract.called)
|
||||
self.assertTrue(mock_per_page.called)
|
||||
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.get_exhibit_pages_new"
|
||||
)
|
||||
@patch("src.pipelines.saas.prompts.prompt_calls.prompt_header_deduplication")
|
||||
@patch(
|
||||
"src.pipelines.shared.preprocessing.preprocessing_funcs.extract_document_index_block"
|
||||
)
|
||||
def test_no_contract_text_skips_di_path(
|
||||
self,
|
||||
mock_extract,
|
||||
mock_dedup,
|
||||
mock_per_page,
|
||||
):
|
||||
from src.pipelines.shared.preprocessing.preprocess import (
|
||||
one_to_n_exhibit_chunking,
|
||||
)
|
||||
|
||||
mock_per_page.return_value = ({}, {"1": ["EXHIBIT A"]})
|
||||
mock_dedup.side_effect = lambda d, _filename: d
|
||||
|
||||
with patch("src.config.DOCUMENT_INDEX_PARSE_ENABLED", True):
|
||||
one_to_n_exhibit_chunking(
|
||||
pages_dict=self.fake_pages,
|
||||
EXHIBIT_HEADER_MARKERS=[],
|
||||
filename=self.filename,
|
||||
contract_text=None,
|
||||
)
|
||||
|
||||
self.assertFalse(mock_extract.called)
|
||||
self.assertTrue(mock_per_page.called)
|
||||
|
||||
|
||||
class TestIdentifySectionBoundaries:
|
||||
TABLE_START = "-------Table Start--------"
|
||||
TABLE_END = "-------Table End--------"
|
||||
|
||||
@@ -63,6 +63,8 @@ USAGE_LABEL_TO_SEGMENT: dict[str, str] = {
|
||||
# --- preprocessing (exhibit boundary detection, runs before one_to_n) ---
|
||||
"EXHIBIT_HEADER": "preprocessing", # src/pipelines/saas/prompts/prompt_calls.py:605
|
||||
"EXHIBIT_HEADER_DEDUP": "preprocessing", # prompt_calls.py:649
|
||||
"DOCUMENT_INDEX_PARSE": "preprocessing", # prompt_calls.py:prompt_document_index
|
||||
"EXHIBIT_HEADER_FROM_VERIFIED_PAGES": "preprocessing", # prompt_calls.py:prompt_exhibit_header_from_verified_pages
|
||||
"EXHIBIT_LINKAGE": "preprocessing", # preprocessing_funcs.py:198 (cross-page exhibit match)
|
||||
# --- one_to_n primary (exhibit_level + reimbursement-primary extraction) ---
|
||||
"REIMBURSEMENT_PRIMARY": "one_to_n_primary", # prompt_calls.py:175
|
||||
|
||||
Reference in New Issue
Block a user