diff --git a/src/config.py b/src/config.py index 1398e20..687a518 100644 --- a/src/config.py +++ b/src/config.py @@ -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")) diff --git a/src/pipelines/saas/file_processing.py b/src/pipelines/saas/file_processing.py index f82d12e..b1de540 100644 --- a/src/pipelines/saas/file_processing.py +++ b/src/pipelines/saas/file_processing.py @@ -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}") diff --git a/src/pipelines/saas/prompts/prompt_calls.py b/src/pipelines/saas/prompts/prompt_calls.py index 0e2879b..af53177 100644 --- a/src/pipelines/saas/prompts/prompt_calls.py +++ b/src/pipelines/saas/prompts/prompt_calls.py @@ -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`. diff --git a/src/pipelines/shared/preprocessing/preprocess.py b/src/pipelines/shared/preprocessing/preprocess.py index e0322e7..cc25aee 100644 --- a/src/pipelines/shared/preprocessing/preprocess.py +++ b/src/pipelines/shared/preprocessing/preprocess.py @@ -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])) diff --git a/src/pipelines/shared/preprocessing/preprocessing_funcs.py b/src/pipelines/shared/preprocessing/preprocessing_funcs.py index 5ec4164..1c7fdb7 100644 --- a/src/pipelines/shared/preprocessing/preprocessing_funcs.py +++ b/src/pipelines/shared/preprocessing/preprocessing_funcs.py @@ -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" + 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" diff --git a/src/prompts/prompt_templates.py b/src/prompts/prompt_templates.py index 81a84d9..2f1c162 100644 --- a/src/prompts/prompt_templates.py +++ b/src/prompts/prompt_templates.py @@ -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) ##################################################################################### diff --git a/src/tests/test_document_index.py b/src/tests/test_document_index.py new file mode 100644 index 0000000..857e06c --- /dev/null +++ b/src/tests/test_document_index.py @@ -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() diff --git a/src/tests/test_preprocessing_funcs.py b/src/tests/test_preprocessing_funcs.py index b5f588d..72e1abc 100644 --- a/src/tests/test_preprocessing_funcs.py +++ b/src/tests/test_preprocessing_funcs.py @@ -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--------" diff --git a/src/utils/instrumentation.py b/src/utils/instrumentation.py index b0d9768..4564947 100644 --- a/src/utils/instrumentation.py +++ b/src/utils/instrumentation.py @@ -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