diff --git a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py index 28697c45..5815a412 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -1,9 +1,9 @@ """Locate hierarchy titles on PDF pages and resolve page ranges. -This module is intentionally deterministic: it performs strict title anchoring, -candidate collection, and range assembly. The page-memory residual agent calls -into these primitives for grep-like tools and adds VLM verification outside this -module. +Deterministic title anchoring and range assembly. Leaf starts come from +offset-guided ``match_overrides`` or per-line strict exact. Null-page parents +are located upstream via compact-strict (cross-line) + optional VLM, then +resolved here including parent self-only spans for interstitial pages. """ from __future__ import annotations @@ -19,10 +19,6 @@ TitleMatchSource = Literal[ "anchored", - "page_compact", - "normalized", - "token", - "printed_prior", "h1_result", "agent_vlm", "agent_heuristic", @@ -80,47 +76,6 @@ class _LineHit: score: float -_STOPWORDS = { - "a", - "an", - "and", - "are", - "as", - "for", - "in", - "of", - "on", - "or", - "the", - "to", - "with", -} - - -def locate_title_start_page( - title: str, - *, - scope_pages: list[int], - page_texts: dict[int, str], - printed_page: int | None = None, - page_offset_hint: int | None = None, -) -> TitleMatch | None: - """Locate *title* using deterministic weak evidence. - - This is a candidate-gathering primitive. The C4 page-memory path should only - directly accept :func:`locate_title_strict_exact`; weak results from this - function are meant for residual agent/VLM arbitration. - """ - matches = collect_title_candidate_matches( - title, - scope_pages=scope_pages, - page_texts=page_texts, - printed_page=printed_page, - page_offset_hint=page_offset_hint, - ) - return matches[0] if matches else None - - def locate_title_strict_exact( title: str, *, @@ -135,84 +90,82 @@ def locate_title_strict_exact( return _choose_best_hit( hits, source="anchored", - printed_page=None, - page_offset_hint=None, extra_evidence={"accept": "strict_exact_unique"}, ) -def collect_title_candidate_matches( +def locate_title_compact_strict( title: str, *, scope_pages: list[int], page_texts: dict[int, str], - printed_page: int | None = None, - page_offset_hint: int | None = None, - limit: int | None = None, -) -> list[TitleMatch]: - """Collect grep-style candidate pages for a title without final arbitration.""" - normalized_title = normalize_heading_text(title) - if not normalized_title or not scope_pages: - return [] +) -> TitleMatch | None: + """Locate *title* after cross-line compact cleanup; accept only a unique page. - hits: list[_LineHit] = [] - for _source, finder in ( - ("anchored", _find_anchored_hits), - ("page_compact", _find_page_compact_hits), - ("normalized", _find_normalized_hits), - ("token", _find_token_hits), - ): - hits.extend(finder(normalized_title, scope_pages, page_texts)) - - by_page: dict[int, list[_LineHit]] = {} - for hit in hits: - by_page.setdefault(hit.page, []).append(hit) - - matches = [ - _choose_best_hit( - page_hits, - source=_preferred_source(page_hits), - printed_page=printed_page, - page_offset_hint=page_offset_hint, - ) - for page_hits in by_page.values() - ] + Pipeline: compact(page text) → contiguous strict match of compact(title) → + accept iff exactly one page in ``scope_pages`` hits. Handles PyMuPDF line + splits; does not use token/normalized weak matching. + """ + needle = _compact_match_text(clean_toc_title(title) or title) + if not needle or not scope_pages: + return None - prior_page = _resolve_printed_prior( - printed_page=printed_page, - page_offset_hint=page_offset_hint, - scope_pages=scope_pages, - ) - if prior_page is not None and prior_page not in by_page: - matches.append( - TitleMatch( - page=prior_page, - confidence=0.35, - source="printed_prior", - matched_line="", - score=0.35, - candidates=[prior_page], - evidence={ - "printed_page": printed_page, - "page_offset_hint": page_offset_hint, - }, - ) - ) + hit_pages: list[int] = [] + matched_preview = "" + for page in scope_pages: + haystack = _compact_match_text(page_texts.get(page, "")) + if not haystack or needle not in haystack: + continue + hit_pages.append(page) + if not matched_preview: + matched_preview = needle[:160] - matches.sort( - key=lambda match: ( - match.score, - match.confidence, - -abs(match.page - (printed_page + page_offset_hint)) - if printed_page is not None and page_offset_hint is not None - else 0, - -match.page, - ), - reverse=True, + unique_pages = sorted(set(hit_pages)) + if len(unique_pages) != 1: + return None + + page = unique_pages[0] + return TitleMatch( + page=page, + confidence=0.92, + source="anchored", + matched_line=matched_preview, + score=0.96, + candidates=[page], + evidence={"accept": "compact_strict_unique"}, ) - if limit is not None: - return matches[: max(int(limit), 0)] - return matches + + +def last_leaf_start_under( + node: TitleNode, + parent_titles: tuple[str, ...], + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> int | None: + """Max start page among located leaves under *node*; None if none located.""" + max_page: int | None = None + for leaf_path, _leaf in iter_leaf_title_nodes([node], parent_titles=parent_titles): + match = match_overrides.get(leaf_path) + if match is None: + continue + if max_page is None or match.page > max_page: + max_page = match.page + return max_page + + +def first_leaf_start_under( + node: TitleNode, + parent_titles: tuple[str, ...], + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> int | None: + """Min start page among located leaves under *node*; None if none located.""" + min_page: int | None = None + for leaf_path, _leaf in iter_leaf_title_nodes([node], parent_titles=parent_titles): + match = match_overrides.get(leaf_path) + if match is None: + continue + if min_page is None or match.page < min_page: + min_page = match.page + return min_page def resolve_hierarchy_page_ranges( @@ -221,15 +174,12 @@ def resolve_hierarchy_page_ranges( page_count: int, page_texts: dict[int, str], body_pages: list[int] | None = None, - page_offset_hint: int | None = None, match_overrides: dict[tuple[str, ...], TitleMatch] | None = None, - use_weak_fallback: bool = False, ) -> list[ResolvedHierarchyRange]: - """Resolve leaf hierarchy nodes into closed page ranges. + """Resolve hierarchy nodes into closed page ranges. - The emitted ranges are leaf-first and intentionally closed-closed: if the - next leaf starts on page N, the previous leaf may also include page N. This - preserves page-to-section many-to-many mapping for dense documents. + Emits leaf ranges and parent self-only spans when a parent start is strictly + before its first located descendant leaf. Ranges are closed-closed. """ if page_count <= 0 or not nodes: return [] @@ -248,9 +198,7 @@ def resolve_hierarchy_page_ranges( allowed_pages=allowed_pages, parent_titles=(), page_texts=page_texts, - page_offset_hint=page_offset_hint, match_overrides=match_overrides or {}, - use_weak_fallback=use_weak_fallback, resolved=resolved, ) return resolved @@ -274,9 +222,7 @@ def _resolve_siblings( allowed_pages: set[int], parent_titles: tuple[str, ...], page_texts: dict[int, str], - page_offset_hint: int | None, match_overrides: dict[tuple[str, ...], TitleMatch], - use_weak_fallback: bool, resolved: list[ResolvedHierarchyRange], ) -> None: located: list[tuple[TitleNode, int, TitleMatch | None]] = [] @@ -285,27 +231,13 @@ def _resolve_siblings( for index, node in enumerate(nodes): path_titles = (*parent_titles, node.title) pages = _allowed_pages_between(lower_bound, parent_scope.end, allowed_pages) - match = _match_override(path_titles, match_overrides, pages) - if match is None: - match = _match_physical_hint(node=node, scope_pages=pages) - if match is None: - match = locate_title_strict_exact( - node.title, - scope_pages=pages, - page_texts=page_texts, - ) - if match is None and use_weak_fallback: - match = locate_title_start_page( - node.title, - scope_pages=pages, - page_texts=page_texts, - printed_page=node.printed_page, - page_offset_hint=page_offset_hint, - ) - if match is None and node.children and match_overrides: - match = _infer_start_from_descendant_overrides( - node, parent_titles, match_overrides, pages, - ) + match = _locate_match_for_node( + node, + path_titles=path_titles, + scope_pages=pages, + page_texts=page_texts, + match_overrides=match_overrides, + ) if match is None: start_page = lower_bound else: @@ -321,9 +253,7 @@ def _resolve_siblings( parent_end=parent_scope.end, allowed_pages=allowed_pages, page_texts=page_texts, - page_offset_hint=page_offset_hint, match_overrides=match_overrides, - use_weak_fallback=use_weak_fallback, parent_titles=parent_titles, ) if next_match is not None: @@ -349,15 +279,32 @@ def _resolve_siblings( ) if node.children: + first_child_start = first_leaf_start_under( + node, parent_titles, match_overrides + ) + if ( + match is not None + and first_child_start is not None + and start_page < first_child_start + ): + resolved.append( + ResolvedHierarchyRange( + title=node.title, + level=node.level, + start_page=start_page, + end_page=first_child_start, + path_titles=path_titles, + match=match, + evidence={**evidence, "skeleton_kind": "parent_self_only"}, + ) + ) _resolve_siblings( node.children, parent_scope=PageRange(start_page, end_page), allowed_pages=allowed_pages, parent_titles=path_titles, page_texts=page_texts, - page_offset_hint=page_offset_hint, match_overrides=match_overrides, - use_weak_fallback=use_weak_fallback, resolved=resolved, ) continue @@ -375,6 +322,34 @@ def _resolve_siblings( ) +def _locate_match_for_node( + node: TitleNode, + *, + path_titles: tuple[str, ...], + scope_pages: list[int], + page_texts: dict[int, str], + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> TitleMatch | None: + match = _match_override(path_titles, match_overrides, scope_pages) + if match is not None: + return match + match = _match_physical_hint(node=node, scope_pages=scope_pages) + if match is not None: + return match + if node.children: + # Parent active locate is upstream (compact-strict / visual). Wide-window + # strict_exact is intentionally not used here. + return _infer_start_from_descendant_overrides( + node, parent_titles=path_titles[:-1], match_overrides=match_overrides, + scope_pages=scope_pages, + ) + return locate_title_strict_exact( + node.title, + scope_pages=scope_pages, + page_texts=page_texts, + ) + + def _find_next_located_sibling( *, nodes: list[TitleNode], @@ -383,36 +358,19 @@ def _find_next_located_sibling( parent_end: int, allowed_pages: set[int], page_texts: dict[int, str], - page_offset_hint: int | None, match_overrides: dict[tuple[str, ...], TitleMatch], - use_weak_fallback: bool, parent_titles: tuple[str, ...], ) -> TitleMatch | None: pages = _allowed_pages_between(lower_bound, parent_end, allowed_pages) for sibling in nodes[start_index:]: path_titles = (*parent_titles, sibling.title) - match = _match_override(path_titles, match_overrides, pages) - if match is None: - match = _match_physical_hint(node=sibling, scope_pages=pages) - if match is not None: - return match - match = locate_title_strict_exact( - sibling.title, + match = _locate_match_for_node( + sibling, + path_titles=path_titles, scope_pages=pages, page_texts=page_texts, + match_overrides=match_overrides, ) - if match is None and use_weak_fallback: - match = locate_title_start_page( - sibling.title, - scope_pages=pages, - page_texts=page_texts, - printed_page=sibling.printed_page, - page_offset_hint=page_offset_hint, - ) - if match is None and sibling.children and match_overrides: - match = _infer_start_from_descendant_overrides( - sibling, parent_titles, match_overrides, pages, - ) if match is not None: return match return None @@ -424,13 +382,7 @@ def _infer_start_from_descendant_overrides( match_overrides: dict[tuple[str, ...], TitleMatch], scope_pages: list[int], ) -> TitleMatch | None: - """Infer a parent node's start page from its earliest located descendant leaf. - - When a non-leaf node cannot be directly located (no printed_page, no grep - match), its descendant leaves may already be in match_overrides from - offset-guided bulk anchoring. Use the minimum page among those descendants - as a synthetic match so the resolver can cap the previous sibling's end_page. - """ + """Final fallback: parent start = earliest located descendant leaf page.""" if not node.children or not match_overrides: return None leaves = iter_leaf_title_nodes([node], parent_titles=parent_titles) @@ -457,6 +409,7 @@ def _infer_start_from_descendant_overrides( evidence={ "inferred_from": "descendant_leaf_override", "original_confidence": min_match.confidence, + "status": "degraded", }, ) @@ -487,50 +440,6 @@ def iter_leaf_title_nodes( return leaves -def max_title_depth(nodes: list[TitleNode]) -> int: - if not nodes: - return 0 - return max( - max(node.level, max_title_depth(node.children)) - if node.children - else node.level - for node in nodes - ) - - -def prune_title_nodes_for_emit_depth( - nodes: list[TitleNode], - *, - emit_depth: int, -) -> list[TitleNode]: - pruned: list[TitleNode] = [] - for node in nodes: - if node.level >= emit_depth or not node.children: - pruned.append( - TitleNode( - title=node.title, - level=node.level, - printed_page=node.printed_page, - physical_page_hint=node.physical_page_hint, - children=[], - ) - ) - continue - pruned.append( - TitleNode( - title=node.title, - level=node.level, - printed_page=node.printed_page, - physical_page_hint=node.physical_page_hint, - children=prune_title_nodes_for_emit_depth( - node.children, - emit_depth=emit_depth, - ), - ) - ) - return pruned - - def _next_located_start( located: list[tuple[TitleNode, int, TitleMatch | None]], start_index: int, @@ -604,6 +513,10 @@ def _allowed_pages_between(start: int, end: int, allowed_pages: set[int]) -> lis return [page for page in range(start, end + 1) if page in allowed_pages] +def _compact_match_text(text: str) -> str: + return re.sub(r"\s+", "", normalize_heading_text(text)).casefold() + + def _find_anchored_hits( title: str, scope_pages: list[int], @@ -628,123 +541,24 @@ def _find_anchored_hits( return hits -def _find_page_compact_hits( - title: str, - scope_pages: list[int], - page_texts: dict[int, str], -) -> list[_LineHit]: - hits: list[_LineHit] = [] - needles = _compact_title_variants(title) - if not needles: - return hits - - for page in scope_pages: - raw_text = page_texts.get(page, "") - compact_text = _compact_match_text(raw_text) - if not compact_text: - continue - matched_needle = next((needle for needle in needles if needle in compact_text), None) - if matched_needle is None: - continue - line_index, evidence = _compact_match_evidence(raw_text, matched_needle) - hits.append( - _LineHit( - page=page, - line_index=line_index, - line=evidence, - source="page_compact", - score=_line_score(line=evidence, line_index=line_index, base=0.94), - ) - ) - return hits - - -def _find_normalized_hits( - title: str, - scope_pages: list[int], - page_texts: dict[int, str], -) -> list[_LineHit]: - hits: list[_LineHit] = [] - needle = normalize_heading_text(clean_toc_title(title)).casefold() - if len(needle) < 2: - return hits - for page, line_index, line in _iter_lines(scope_pages, page_texts): - cleaned_line = normalize_heading_text(clean_toc_title(line)).casefold() - if not cleaned_line: - continue - if needle in cleaned_line or _is_strong_reverse_match(cleaned_line, needle): - hits.append( - _LineHit( - page=page, - line_index=line_index, - line=line.strip(), - source="normalized", - score=_line_score(line=line, line_index=line_index, base=0.9), - ) - ) - return hits - - -def _find_token_hits( - title: str, - scope_pages: list[int], - page_texts: dict[int, str], -) -> list[_LineHit]: - title_tokens = _significant_tokens(clean_toc_title(title) or title) - if not title_tokens: - return [] - - hits: list[_LineHit] = [] - for page, line_index, line in _iter_lines(scope_pages, page_texts): - line_tokens = _significant_tokens(line) - if not line_tokens: - continue - coverage = len(title_tokens & line_tokens) / len(title_tokens) - if coverage < 0.8: - continue - hits.append( - _LineHit( - page=page, - line_index=line_index, - line=line.strip(), - source="token", - score=_line_score(line=line, line_index=line_index, base=0.78) - + coverage, - ) - ) - return hits - - def _choose_best_hit( hits: list[_LineHit], *, source: TitleMatchSource, - printed_page: int | None, - page_offset_hint: int | None, extra_evidence: dict[str, Any] | None = None, ) -> TitleMatch: - expected_page = ( - printed_page + page_offset_hint - if printed_page is not None and page_offset_hint is not None - else None + ordered = sorted( + hits, + key=lambda hit: (hit.score, -hit.line_index, -hit.page), + reverse=True, ) - - def sort_key(hit: _LineHit) -> tuple[float, int, int, int]: - printed_bonus = 0 - if expected_page is not None: - printed_bonus = -abs(hit.page - expected_page) - return (hit.score, printed_bonus, -hit.line_index, -hit.page) - - ordered = sorted(hits, key=sort_key, reverse=True) best = ordered[0] pages = sorted({hit.page for hit in ordered}) confidence_by_source = { "anchored": 0.92, - "page_compact": 0.9, - "normalized": 0.84, - "token": 0.72, - "printed_prior": 0.35, "h1_result": 0.88, + "agent_vlm": 0.75, + "agent_heuristic": 0.5, } return TitleMatch( page=best.page, @@ -756,41 +570,11 @@ def sort_key(hit: _LineHit) -> tuple[float, int, int, int]: evidence={ "line_index": best.line_index, "candidate_count": len(pages), - "printed_page": printed_page, - "page_offset_hint": page_offset_hint, **(extra_evidence or {}), }, ) -def _preferred_source(hits: list[_LineHit]) -> TitleMatchSource: - priority = { - "anchored": 50, - "page_compact": 40, - "normalized": 30, - "token": 20, - "printed_prior": 10, - "h1_result": 60, - "agent_vlm": 70, - "agent_heuristic": 15, - } - return max(hits, key=lambda hit: (priority.get(hit.source, 0), hit.score)).source - - -def _resolve_printed_prior( - *, - printed_page: int | None, - page_offset_hint: int | None, - scope_pages: list[int], -) -> int | None: - if printed_page is None or page_offset_hint is None or not scope_pages: - return None - page = printed_page + page_offset_hint - if page in scope_pages: - return page - return None - - def _line_score(*, line: str, line_index: int, base: float) -> float: stripped = normalize_heading_text(line) short_line_bonus = max(0.0, 1.0 - (len(stripped) / 140.0)) @@ -810,55 +594,6 @@ def _iter_lines( return rows -def _compact_title_variants(title: str) -> list[str]: - normalized = normalize_heading_text(clean_toc_title(title) or title).casefold() - compacted: list[str] = [] - compact = _compact_match_text(normalized) - if compact: - compacted.append(compact) - return compacted - - -def _compact_match_text(text: str) -> str: - return re.sub(r"\s+", "", normalize_heading_text(text)).casefold() - - -def _compact_match_evidence(raw_text: str, compact_needle: str) -> tuple[int, str]: - lines = [line.strip() for line in raw_text.splitlines() if line.strip()] - if not lines: - return 0, "" - - compact_so_far = "" - start_index = 0 - for index, line in enumerate(lines): - line_compact = _compact_match_text(line) - if not compact_so_far: - start_index = index - compact_so_far += line_compact - if compact_needle in compact_so_far: - return start_index, " ".join(lines[start_index : index + 1])[:160] - if len(compact_so_far) > len(compact_needle) * 3: - compact_so_far = line_compact - start_index = index - - return 0, " ".join(lines[:3])[:160] - - -def _is_strong_reverse_match(fragment: str, title: str) -> bool: - return len(fragment) >= 6 and fragment in title - - -def _significant_tokens(text: str) -> set[str]: - normalized = normalize_heading_text(clean_toc_title(text) or text).casefold() - latin = { - token - for token in re.findall(r"[a-z0-9][a-z0-9_-]+", normalized) - if token not in _STOPWORDS - } - cjk = set(re.findall(r"[\u4e00-\u9fff]", normalized)) - return latin | cjk - - def _extract_flat_entries(payload: Any) -> list[dict[str, Any]]: if isinstance(payload, list): return [ diff --git a/apps/worker/app/services/document_agent/structure/page_locate_agent.py b/apps/worker/app/services/document_agent/structure/page_locate_agent.py index 7a68f362..f6aa312d 100644 --- a/apps/worker/app/services/document_agent/structure/page_locate_agent.py +++ b/apps/worker/app/services/document_agent/structure/page_locate_agent.py @@ -1,4 +1,10 @@ -"""Residual page-location agent for hierarchy titles.""" +"""VLM page verification for page-memory offset calibration. + +This module provides the deterministic-input, VLM-arbitrated helper used by the +page-memory skeleton calibration to confirm which candidate page starts a given +section title. The former residual ReAct sub-agent has been removed; calibration +now drives offset-guided bulk anchoring directly and only needs this verifier. +""" from __future__ import annotations @@ -6,26 +12,12 @@ import json import os import time -from dataclasses import dataclass, field from typing import Any, cast from loguru import logger from app.services.document_agent.manifest import ToolContext -from app.services.document_agent.structure.hierarchy_locator import ( - TitleMatch, - TitleNode, - collect_title_candidate_matches, - iter_leaf_title_nodes, - locate_title_strict_exact, - max_title_depth, - prune_title_nodes_for_emit_depth, -) -from app.services.document_agent.structure.page_locate_subagent import ( - PageLocateSubAgent, - SubAgentConfig, -) -from shared.models.schemas.page_memory_config import PageMemoryConfig +from app.services.document_agent.structure.hierarchy_locator import TitleMatch VLM_CONFIRMED_DEFAULT_CONFIDENCE = 0.75 GREP_ONLY_CONFIDENCE_CAP = 0.62 @@ -34,224 +26,6 @@ VLM_FAILED_GREP_CONFIDENCE_CAP = 0.54 -@dataclass(frozen=True) -class PageLocateConfig: - residual_agent_limit: int = 50 - max_emit_depth: int = 5 - min_emit_depth: int = 2 - vlm_candidate_page_cap: int = 4 - full_leaf_sections: bool = False - - @classmethod - def from_page_memory_config( - cls, - page_memory_config: PageMemoryConfig, - ) -> "PageLocateConfig": - return cls( - residual_agent_limit=page_memory_config.page_locate_residual_agent_limit, - max_emit_depth=page_memory_config.page_locate_max_emit_depth, - min_emit_depth=page_memory_config.page_locate_min_emit_depth, - vlm_candidate_page_cap=( - page_memory_config.page_locate_vlm_candidate_page_cap - ), - full_leaf_sections=page_memory_config.page_locate_full_leaf_sections, - ) - - -@dataclass(frozen=True) -class ResidualRequest: - path_titles: tuple[str, ...] - node: TitleNode - - -@dataclass(frozen=True) -class PageLocatePrepareResult: - nodes: list[TitleNode] - match_overrides: dict[tuple[str, ...], TitleMatch] - summary: dict[str, Any] = field(default_factory=dict) - - -class PageLocateResidualAgent: - """Batch residual resolver for C4 title-to-page anchoring.""" - - def __init__( - self, - *, - ctx: ToolContext | None, - page_texts: dict[int, str], - body_pages: list[int], - page_count: int, - page_offset_hint: int | None, - config: PageLocateConfig | None = None, - ) -> None: - self.ctx = ctx - self.page_texts = page_texts - self.body_pages = sorted({page for page in body_pages if 1 <= page <= page_count}) - self.page_count = page_count - self.page_offset_hint = page_offset_hint - self.config = config or PageLocateConfig() - - def prepare(self, nodes: list[TitleNode]) -> PageLocatePrepareResult: - if not nodes: - return PageLocatePrepareResult(nodes=[], match_overrides={}, summary={}) - - actual_max_depth = max_title_depth(nodes) - emit_depth = min(actual_max_depth, self.config.max_emit_depth) - emit_depth = max(emit_depth, self.config.min_emit_depth) - direct_matches: dict[tuple[str, ...], TitleMatch] = {} - residuals: list[ResidualRequest] = [] - - while True: - selected_nodes = prune_title_nodes_for_emit_depth(nodes, emit_depth=emit_depth) - direct_matches, residuals = self._classify_residuals(selected_nodes) - if ( - self.config.full_leaf_sections - or len(residuals) <= self.config.residual_agent_limit - or emit_depth <= self.config.min_emit_depth - ): - break - emit_depth -= 1 - - residual_matches = self._resolve_residuals(residuals) - match_overrides = {**direct_matches, **residual_matches} - summary = { - "agent": "page_locate_residual", - "emit_depth": emit_depth, - "actual_max_depth": actual_max_depth, - "direct_exact_count": len(direct_matches), - "residual_count": len(residuals), - "resolved_residual_count": len(residual_matches), - "residual_limit": self.config.residual_agent_limit, - "vlm_candidate_page_cap": self.config.vlm_candidate_page_cap, - "full_leaf_sections": self.config.full_leaf_sections, - } - logger.info("[page_locate.agent] summary={}", summary) - return PageLocatePrepareResult( - nodes=selected_nodes, - match_overrides=match_overrides, - summary=summary, - ) - - def _entry_scope(self, node: TitleNode) -> list[int]: - """Narrow search scope for a node using its printed page + offset hint. - - When a node carries a printed_page and we have a page_offset_hint, - restrict grep to a small window around the expected physical page. - The window spans [printed_page, printed_page + offset] (inclusive, - order-independent) to cover the uncertainty between printed numbering - and physical page positions. - """ - if node.printed_page is None or self.page_offset_hint is None: - return self.body_pages - expected = node.printed_page + self.page_offset_hint - lo = min(node.printed_page, expected) - hi = max(node.printed_page, expected) - return [p for p in self.body_pages if lo <= p <= hi] - - def _classify_residuals( - self, - nodes: list[TitleNode], - ) -> tuple[dict[tuple[str, ...], TitleMatch], list[ResidualRequest]]: - direct: dict[tuple[str, ...], TitleMatch] = {} - residuals: list[ResidualRequest] = [] - for path_titles, node in iter_leaf_title_nodes(nodes): - scope = self._entry_scope(node) - match = locate_title_strict_exact( - node.title, - scope_pages=scope, - page_texts=self.page_texts, - ) - if match is not None: - direct[path_titles] = TitleMatch( - page=match.page, - confidence=match.confidence, - source=match.source, - matched_line=match.matched_line, - score=match.score, - candidates=match.candidates, - evidence={ - **match.evidence, - "page_locate_agent": { - "decision": "direct_strict_exact", - "path_titles": list(path_titles), - }, - }, - ) - else: - residuals.append(ResidualRequest(path_titles=path_titles, node=node)) - return direct, residuals - - def _resolve_residuals( - self, - residuals: list[ResidualRequest], - ) -> dict[tuple[str, ...], TitleMatch]: - matches: dict[tuple[str, ...], TitleMatch] = {} - if not residuals or self.ctx is None: - # No agent runtime (ctx/budget) → leave residuals for physical-hint / - # neighbor-boundary fallback in the resolver. The residual sub-agent - # only runs when a real ToolContext is available. - return matches - - sub_config = SubAgentConfig( - candidate_cap=max(self.config.vlm_candidate_page_cap * 2, 4), - verify_page_cap=max(self.config.vlm_candidate_page_cap, 1), - ) - for residual in residuals: - scope = self._entry_scope(residual.node) - agent = PageLocateSubAgent( - ctx=self.ctx, - scope_pages=scope, - page_count=self.page_count, - config=sub_config, - page_offset_hint=self.page_offset_hint, - ) - result = agent.locate( - title=residual.node.title, - printed_page=residual.node.printed_page, - ) - if result.match is None: - logger.warning( - "[page_locate.agent] unresolved title={!r} path_titles={} stop={}", - residual.node.title, - residual.path_titles, - result.stop_reason, - ) - continue - base = result.match - agent_evidence = dict(base.evidence.get("page_locate_agent", {})) - agent_evidence["path_titles"] = list(residual.path_titles) - agent_evidence["stop_reason"] = result.stop_reason - matches[residual.path_titles] = TitleMatch( - page=base.page, - confidence=base.confidence, - source=base.source, - matched_line=base.matched_line, - score=base.score, - candidates=base.candidates, - evidence={**base.evidence, "page_locate_agent": agent_evidence}, - ) - return matches - - -def grep_title_page_candidates( - *, - title: str, - scope_pages: list[int], - page_texts: dict[int, str], - printed_page: int | None = None, - page_offset_hint: int | None = None, - limit: int | None = None, -) -> list[TitleMatch]: - return collect_title_candidate_matches( - title, - scope_pages=scope_pages, - page_texts=page_texts, - printed_page=printed_page, - page_offset_hint=page_offset_hint, - limit=limit, - ) - - def verify_section_page_choice( *, ctx: ToolContext | None, diff --git a/apps/worker/app/services/document_agent/structure/page_locate_subagent.py b/apps/worker/app/services/document_agent/structure/page_locate_subagent.py deleted file mode 100644 index c3826cb2..00000000 --- a/apps/worker/app/services/document_agent/structure/page_locate_subagent.py +++ /dev/null @@ -1,446 +0,0 @@ -"""Per-title ReAct sub-agent for residual hierarchy page location. - -This is intentionally NOT a fixed pipeline. For every title that strict anchoring -could not resolve, an LLM runs a bounded reasoning loop and decides — round by -round — which tool to call: - -* ``grep.title_pages`` — find candidate body pages for a (possibly rewritten) - query. The agent may shorten the title to a distinctive core when the full - title carries trailing document-reference codes / ``《》`` wrappers / version - notes that never appear contiguously in the body, or when the heading is split - across lines. -* ``verify.section_page`` — render candidate pages and ask a VLM which one truly - *starts* the section (vs a TOC row, header/footer, or a body citation). - -The agent then ``submit``s a confirmed start page (or gives up). Tools are -dispatched through the shared registry, so the agent reuses the exact same -grep/VLM primitives the rest of the profile agent uses. -""" - -from __future__ import annotations - -import json -from dataclasses import dataclass, field -from typing import Any, Callable, cast - -from loguru import logger - -from app.services.document_agent.manifest import ToolContext -from app.services.document_agent.registry import REGISTRY -from app.services.document_agent.structure.hierarchy_locator import TitleMatch - -DecideFn = Callable[..., dict[str, Any] | None] - -_ALLOWED_ACTIONS = {"grep", "verify", "submit", "give_up"} -VLM_MATCH_DEFAULT_CONFIDENCE = 0.7 -HEURISTIC_MATCH_DEFAULT_CONFIDENCE = 0.55 - - -@dataclass(frozen=True) -class SubAgentConfig: - max_rounds: int = 6 - candidate_cap: int = 8 - verify_page_cap: int = 4 - decide_max_tokens: int = 500 - - -@dataclass -class SubAgentResult: - match: TitleMatch | None - transcript: list[dict[str, Any]] = field(default_factory=list) - rounds: int = 0 - stop_reason: str = "exhausted" - - -class PageLocateSubAgent: - """Bounded ReAct loop that locates ONE residual title via grep + VLM tools.""" - - def __init__( - self, - *, - ctx: ToolContext | None, - scope_pages: list[int], - page_count: int, - config: SubAgentConfig | None = None, - page_offset_hint: int | None = None, - decide: DecideFn | None = None, - ) -> None: - self.ctx = ctx - self.page_count = page_count - self.scope_pages = sorted({p for p in scope_pages if 1 <= p <= page_count}) - self.config = config or SubAgentConfig() - self.page_offset_hint = page_offset_hint - self._decide = decide or decide_next_action - # Ensure the page-locate tools are registered. Importing the light - # ``structure`` module (not the heavy ``tools`` package) avoids pulling - # in S3/DB-dependent modules and avoids import cycles. - import app.services.document_agent.structure.page_locate_tools # noqa: F401 - - def locate(self, *, title: str, printed_page: int | None = None) -> SubAgentResult: - if not self.scope_pages: - return SubAgentResult(match=None, stop_reason="empty_scope") - - transcript: list[dict[str, Any]] = [] - seen_grep: dict[int, dict[str, Any]] = {} - last_verify: dict[str, Any] | None = None - - for round_index in range(self.config.max_rounds): - decision = self._decide( - ctx=self.ctx, - title=title, - scope_pages=self.scope_pages, - printed_page=printed_page, - page_offset_hint=self.page_offset_hint, - observations=transcript, - seen_pages=sorted(seen_grep.keys()), - last_verify=last_verify, - round_index=round_index, - max_rounds=self.config.max_rounds, - ) - if decision is None: - break - - action = str(decision.get("action") or "") - entry: dict[str, Any] = {"round": round_index, "decision": decision} - transcript.append(entry) - - if action == "grep": - query = str(decision.get("query_title") or title).strip() or title - pages = self._coerce_pages(decision.get("pages")) or self.scope_pages - result = self._dispatch( - "grep.title_pages", - { - "title": query, - "pages": pages, - "candidate_cap": self.config.candidate_cap, - }, - ) - candidates = [] - if result.status == "ok" and result.payload: - candidates = list(result.payload.get("candidates") or []) - for cand in candidates: - seen_grep[int(cand["page"])] = cand - entry["observation"] = { - "tool": "grep.title_pages", - "query": query, - "candidates": candidates[: self.config.candidate_cap], - } - - elif action == "verify": - pages = self._coerce_pages(decision.get("pages")) or sorted(seen_grep.keys()) - pages = pages[: max(self.config.verify_page_cap, 1)] - if not pages: - entry["observation"] = { - "tool": "verify.section_page", - "error": "no candidate pages to verify", - } - continue - result = self._dispatch("verify.section_page", {"title": title, "pages": pages}) - choice = result.payload if (result.status == "ok" and result.payload) else {} - last_verify = choice - entry["observation"] = { - "tool": "verify.section_page", - "pages": pages, - "choice": choice, - } - - elif action == "submit": - selected = decision.get("selected_page") - if selected is None: - return SubAgentResult(None, transcript, round_index + 1, "submit_null") - page = int(selected) - if page not in self.scope_pages: - entry["observation"] = {"error": f"submit page {page} outside scope"} - continue - match = self._build_match( - title=title, - page=page, - seen_grep=seen_grep, - last_verify=last_verify, - decision=decision, - transcript=transcript, - ) - return SubAgentResult(match, transcript, round_index + 1, "submit") - - elif action == "give_up": - return SubAgentResult(None, transcript, round_index + 1, "give_up") - - else: - entry["observation"] = {"error": f"unknown action {action!r}"} - - # Out of rounds / budget: accept a VLM-confirmed page if one exists. - if last_verify and last_verify.get("selected_page") is not None: - page = int(last_verify["selected_page"]) - if page in self.scope_pages: - match = self._build_match( - title=title, - page=page, - seen_grep=seen_grep, - last_verify=last_verify, - decision={"confidence": last_verify.get("confidence")}, - transcript=transcript, - ) - return SubAgentResult(match, transcript, self.config.max_rounds, "round_cap_with_verify") - return SubAgentResult(None, transcript, len(transcript), "exhausted") - - def _dispatch(self, name: str, args: dict[str, Any]): - assert self.ctx is not None, "dispatch requires a non-None ToolContext" - return REGISTRY.dispatch(name, self.ctx, args) - - def _coerce_pages(self, raw: Any) -> list[int]: - if not isinstance(raw, (list, tuple)): - return [] - pages: list[int] = [] - for value in raw: - try: - page = int(value) - except (TypeError, ValueError): - continue - if page in self.scope_pages and page not in pages: - pages.append(page) - return pages - - def _build_match( - self, - *, - title: str, - page: int, - seen_grep: dict[int, dict[str, Any]], - last_verify: dict[str, Any] | None, - decision: dict[str, Any], - transcript: list[dict[str, Any]], - ) -> TitleMatch: - grep_hit = seen_grep.get(page) or {} - vlm_confirmed = bool(last_verify and last_verify.get("selected_page") == page) - if vlm_confirmed: - assert last_verify is not None # guaranteed by vlm_confirmed check - source = cast(Any, last_verify.get("source") or "agent_vlm") - confidence = float(last_verify.get("confidence") or VLM_MATCH_DEFAULT_CONFIDENCE) - reason = last_verify.get("reason") - else: - source = cast(Any, "agent_heuristic") - confidence = float( - decision.get("confidence") - or grep_hit.get("confidence") - or HEURISTIC_MATCH_DEFAULT_CONFIDENCE - ) - reason = decision.get("reason") - return TitleMatch( - page=page, - confidence=confidence, - source=source, - matched_line=str(grep_hit.get("matched_line") or ""), - score=float(grep_hit.get("score") or confidence), - candidates=sorted(seen_grep.keys()), - evidence={ - "page_locate_agent": { - "decision": source, - "selected_page": page, - "reason": reason, - "vlm_confirmed": vlm_confirmed, - "rounds": len(transcript), - "seen_pages": sorted(seen_grep.keys()), - "transcript": transcript[-8:], - } - }, - ) - - -# ── Decision policy ──────────────────────────────────────────────────────── - - -def decide_next_action( - *, - ctx: ToolContext | None, - title: str, - scope_pages: list[int], - printed_page: int | None, - page_offset_hint: int | None, - observations: list[dict[str, Any]], - seen_pages: list[int], - last_verify: dict[str, Any] | None, - round_index: int, - max_rounds: int, -) -> dict[str, Any] | None: - """Pick the next action. - - Primary path is LLM-driven (the agent may rewrite the query, expand scope, - verify with a VLM, then submit). When no reasoning model is configured we - fall back to a minimal deterministic policy (strict grep → verify → submit) - that performs no query rewriting — only the LLM is allowed to relax the - query, so offline mode never silently mis-anchors. - """ - model = None - if ctx is not None: - model = ctx.settings.get("executor_model") or ctx.settings.get("model") - if ctx is None or not model or getattr(ctx, "budget", None) is None: - return _deterministic_decide( - title=title, - observations=observations, - seen_pages=seen_pages, - last_verify=last_verify, - ) - - from shared.utils.token_estimate import estimate_tokens - - prompt = _build_prompt( - title=title, - scope_pages=scope_pages, - printed_page=printed_page, - page_offset_hint=page_offset_hint, - observations=observations, - seen_pages=seen_pages, - last_verify=last_verify, - round_index=round_index, - max_rounds=max_rounds, - ) - est = estimate_tokens(prompt) - if not ctx.budget.try_reserve("plan", est): - logger.warning("[page_locate.subagent] planner budget exhausted for title={!r}", title) - return None - try: - from shared.services.ai.llm_overrides import get_text_client - - client, model = get_text_client(requested_model=model) - raw, usage = client.chat_completion_with_usage( - messages=[{"role": "user", "content": prompt}], - model=model, - temperature=0.0, - max_tokens=500, - response_format={"type": "json_object"}, - usage_task="page_memory.page_locate_decide", - ) - ctx.budget.commit("plan", actual=usage.get("total_tokens", est), est=est) - return _normalize_decision(json.loads(raw)) - except Exception as exc: - ctx.budget.refund("plan", est=est) - logger.warning("[page_locate.subagent] decide failed for title={!r}: {}", title, exc) - return None - - -def _deterministic_decide( - *, - title: str, - observations: list[dict[str, Any]], - seen_pages: list[int], - last_verify: dict[str, Any] | None, -) -> dict[str, Any]: - has_grep = any( - entry.get("observation", {}).get("tool") == "grep.title_pages" - for entry in observations - ) - if not has_grep: - return {"action": "grep", "query_title": title, "reason": "initial strict grep"} - if last_verify is None: - if seen_pages: - return {"action": "verify", "pages": seen_pages, "reason": "verify grep candidates"} - return {"action": "give_up", "reason": "strict grep found no candidates"} - selected = last_verify.get("selected_page") - if selected is not None: - return { - "action": "submit", - "selected_page": selected, - "confidence": last_verify.get("confidence", 0.6), - "reason": "submit VLM-confirmed page", - } - return {"action": "give_up", "reason": "verify returned no page"} - - -def _normalize_decision(data: dict[str, Any]) -> dict[str, Any]: - action = str(data.get("action") or "").strip().lower() - if action not in _ALLOWED_ACTIONS: - action = "give_up" - decision: dict[str, Any] = {"action": action, "reason": data.get("reason")} - if action == "grep": - decision["query_title"] = data.get("query_title") - decision["pages"] = data.get("pages") - elif action == "verify": - decision["pages"] = data.get("pages") - elif action == "submit": - selected = data.get("selected_page") - decision["selected_page"] = None if selected in (None, "", "null") else selected - decision["confidence"] = data.get("confidence") - return decision - - -def _build_prompt( - *, - title: str, - scope_pages: list[int], - printed_page: int | None, - page_offset_hint: int | None, - observations: list[dict[str, Any]], - seen_pages: list[int], - last_verify: dict[str, Any] | None, - round_index: int, - max_rounds: int, -) -> str: - scope_desc = ( - f"{scope_pages[0]}-{scope_pages[-1]} ({len(scope_pages)} pages)" - if scope_pages - else "none" - ) - hint = "" - if printed_page is not None: - guessed = printed_page + page_offset_hint if page_offset_hint is not None else None - hint = f"\nPrinted page number in the title's TOC entry: {printed_page}" + ( - f" (rough physical-page guess: {guessed})" if guessed is not None else "" - ) - compact_obs = json.dumps( - [_compact_obs(entry) for entry in observations[-4:]], - ensure_ascii=False, - ) - return ( - "You are a sub-agent that locates the START page of ONE section title " - "inside a PDF body. Decide the single next action and return strict JSON.\n\n" - f"Title to locate: {title!r}\n" - f"Allowed body pages: {scope_desc}{hint}\n" - f"Round {round_index + 1} of {max_rounds}. " - f"Candidate pages seen so far: {seen_pages}. " - f"Last VLM verification: {json.dumps(last_verify, ensure_ascii=False)}\n" - f"Observations (most recent last): {compact_obs}\n\n" - "Available actions (pick exactly ONE, as a JSON object):\n" - '1. {"action":"grep","query_title":"","reason":"..."}\n' - " Searches body pages for the query (exact line, whitespace-insensitive " - "compact text, and token overlap) and returns candidate pages with matched " - "lines. IMPORTANT: the title may carry a trailing document-reference code " - "that never appears contiguously in the body, and the heading may be split " - "across two lines. If a strict search of the full title returns nothing, " - "retry grep with a shortened, distinctive CORE of the title (drop the " - "trailing brackets/codes/notes).\n" - '2. {"action":"verify","pages":[,...],"reason":"..."}\n' - " Renders those pages as images and asks a vision model which one truly " - "STARTS the section (not a table-of-contents row, a running header/footer, " - "or a body mention/citation). Use this to disambiguate.\n" - '3. {"action":"submit","selected_page":,"confidence":<0..1>,"reason":"..."}\n' - " Finalize. selected_page must be a page you have already seen as a " - "candidate; use null only if the section is genuinely absent.\n" - '4. {"action":"give_up","reason":"..."}\n\n' - "Rules: when there is more than one candidate or any ambiguity, verify with " - "the vision model before submitting. Never invent page numbers. " - "Return ONLY the JSON object." - ) - - -def _compact_obs(entry: dict[str, Any]) -> dict[str, Any]: - obs = entry.get("observation") or {} - tool = obs.get("tool") - if tool == "grep.title_pages": - return { - "action": "grep", - "query": obs.get("query"), - "candidates": [ - {"page": c.get("page"), "source": c.get("source"), "line": (c.get("matched_line") or "")[:60]} - for c in (obs.get("candidates") or []) - ], - } - if tool == "verify.section_page": - choice = obs.get("choice") or {} - return { - "action": "verify", - "pages": obs.get("pages"), - "selected_page": choice.get("selected_page"), - "source": choice.get("source"), - "reason": (choice.get("reason") or "")[:80], - } - return {"action": entry.get("decision", {}).get("action"), "note": obs.get("error")} diff --git a/apps/worker/app/services/document_agent/structure/page_locate_tools.py b/apps/worker/app/services/document_agent/structure/page_locate_tools.py deleted file mode 100644 index 78f319b2..00000000 --- a/apps/worker/app/services/document_agent/structure/page_locate_tools.py +++ /dev/null @@ -1,158 +0,0 @@ -"""Registry tools for page-memory residual title location. - -These live under ``structure/`` (not ``tools/``) so the page-memory ReAct -sub-agent can register and dispatch them without importing the full profiling -``tools`` package (which pulls in S3/DB-heavy modules). ``tools/page_locate.py`` -re-exports these so the profile executor and ``tools/__init__`` keep working. -""" - -from __future__ import annotations - -import time -from typing import Any - -from app.services.document_agent.manifest import ToolContext, ToolResult -from app.services.document_agent.registry import register_tool -from app.services.document_agent.structure import page_locate_agent as _pla -from app.services.document_agent.structure.hierarchy_locator import TitleMatch - -SYNTHETIC_CANDIDATE_CONFIDENCE = 0.4 - - -@register_tool( - name="grep.title_pages", - description=( - "Find candidate body pages for a section title using strict heading, " - "normalized, compact, and token grep variants." - ), - parameters={ - "type": "object", - "properties": { - "title": {"type": "string"}, - "pages": {"type": "array", "items": {"type": "integer"}}, - "candidate_cap": {"type": "integer"}, - }, - "required": ["title", "pages"], - }, -) -def grep_title_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: - start = time.monotonic() - title = str(args.get("title") or "").strip() - pages = _valid_pages(ctx, args.get("pages") or []) - if not title or not pages: - return ToolResult( - status="error", - error="grep.title_pages requires title and pages", - latency_ms=int((time.monotonic() - start) * 1000), - ) - page_texts = _get_page_texts(ctx, pages) - matches = _pla.grep_title_page_candidates( - title=title, - scope_pages=pages, - page_texts=page_texts, - limit=int(args.get("candidate_cap") or 8), - ) - return ToolResult( - status="ok", - payload={ - "title": title, - "candidates": [_match_to_payload(match) for match in matches], - }, - latency_ms=int((time.monotonic() - start) * 1000), - ) - - -@register_tool( - name="verify.section_page", - description=( - "Use VLM page screenshots to choose which candidate page starts a section." - ), - parameters={ - "type": "object", - "properties": { - "title": {"type": "string"}, - "pages": {"type": "array", "items": {"type": "integer"}}, - }, - "required": ["title", "pages"], - }, -) -def verify_section_page(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: - start = time.monotonic() - title = str(args.get("title") or "").strip() - pages = _valid_pages(ctx, args.get("pages") or []) - if not title or not pages: - return ToolResult( - status="error", - error="verify.section_page requires title and pages", - latency_ms=int((time.monotonic() - start) * 1000), - ) - page_texts = _get_page_texts(ctx, pages) - grep_by_page = { - match.page: match - for match in _pla.grep_title_page_candidates( - title=title, - scope_pages=pages, - page_texts=page_texts, - limit=len(pages), - ) - } - # Always verify the *requested* pages, even when grep found nothing on them - # (the heading may be split across lines / wrapped, so render + VLM decides). - candidates = [ - grep_by_page.get(page) - or TitleMatch( - page=page, - confidence=SYNTHETIC_CANDIDATE_CONFIDENCE, - source="agent_heuristic", - matched_line="", - score=SYNTHETIC_CANDIDATE_CONFIDENCE, - candidates=[page], - evidence={"synthesized": True}, - ) - for page in pages - ] - choice = _pla.verify_section_page_choice( - ctx=ctx, - title=title, - candidate_matches=candidates, - candidate_page_cap=len(pages), - ) - return ToolResult( - status="ok", - payload=choice, - latency_ms=int((time.monotonic() - start) * 1000), - tokens_used=int(choice.get("tokens_used") or 0), - ) - - -def _valid_pages(ctx: ToolContext, raw_pages: Any) -> list[int]: - page_count = max(int(ctx.blackboard.page_count or 0), 0) - return sorted( - { - int(page) - for page in raw_pages - if 1 <= int(page) <= page_count - } - ) - - -def _get_page_texts(ctx: ToolContext, pages: list[int]) -> dict[int, str]: - cache = getattr(ctx.blackboard, "page_full_text_cache", {}) - missing = [page for page in pages if page not in cache] - if missing: - from app.services.document_agent.pdf_text import read_page_texts - - cache.update(read_page_texts(ctx.pdf_path, missing, timeout=300)) - return {page: cache.get(page, "") for page in pages} - - -def _match_to_payload(match: Any) -> dict[str, Any]: - return { - "page": match.page, - "confidence": match.confidence, - "source": match.source, - "matched_line": match.matched_line, - "score": match.score, - "candidates": match.candidates, - "evidence": match.evidence, - } diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index ac5c2990..177ff879 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -6,7 +6,6 @@ from . import find_toc_anchor_pages as find_toc_anchor_pages # noqa: F401 from . import grep_text as grep_text # noqa: F401 from . import inspect_pages as inspect_pages # noqa: F401 -from . import page_locate as page_locate # noqa: F401 from . import propose_shard_plan as propose_shard_plan # noqa: F401 from . import validate_anatomy_map as validate_anatomy_map # noqa: F401 from . import verdict as verdict # noqa: F401 diff --git a/apps/worker/app/services/document_agent/tools/page_locate.py b/apps/worker/app/services/document_agent/tools/page_locate.py deleted file mode 100644 index 0e67f143..00000000 --- a/apps/worker/app/services/document_agent/tools/page_locate.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Page-memory residual title-location tools. - -The implementations live in ``structure/page_locate_tools.py`` so the -page-memory sub-agent can register/dispatch them without importing this heavy -``tools`` package. Importing this module (e.g. via ``tools/__init__``) ensures -the tools are registered for the profile executor as well. -""" - -from __future__ import annotations - -from app.services.document_agent.structure.page_locate_tools import ( # noqa: F401 - grep_title_pages, - verify_section_page, -) - -__all__ = ["grep_title_pages", "verify_section_page"] diff --git a/apps/worker/app/services/page_memory/_serialization.py b/apps/worker/app/services/page_memory/_serialization.py index c32fe72e..a963076c 100644 --- a/apps/worker/app/services/page_memory/_serialization.py +++ b/apps/worker/app/services/page_memory/_serialization.py @@ -9,7 +9,6 @@ from app.services.page_memory._utils import ( collapse_page_ranges, page_scope_info, - sort_skeletons, ) from shared.services.chunks.path_segments import split_escaped_document_path @@ -93,8 +92,14 @@ def serialize_skeletons(skeletons: list[Any]) -> list[dict[str, Any]]: def build_hierarchy_tree(skeletons: list[Any]) -> dict[str, Any]: + """Build a nested title tree preserving ``skeletons`` list order. + + Sibling key order follows first-seen order in ``skeletons`` (dict + insertion order). Callers that need cross-page ordering should + ``sort_skeletons`` first; same-page order must remain VLM/TOC order. + """ hierarchy: dict[str, Any] = {} - for skel in sort_skeletons(skeletons): + for skel in skeletons: parts = split_escaped_document_path( getattr(skel, "section_path", "") or "" ) @@ -112,6 +117,7 @@ def serialize_hierarchy_artifact( *, scope_manifest_data: dict[str, Any] | None = None, ) -> dict[str, Any]: + # Keep nodes and HIERARCHY on the same list order — do not re-sort here. nodes = serialize_skeletons(skeletons) artifact: dict[str, Any] = { "HIERARCHY": build_hierarchy_tree(skeletons), diff --git a/apps/worker/app/services/page_memory/_utils.py b/apps/worker/app/services/page_memory/_utils.py index 219889bf..23d89f7f 100644 --- a/apps/worker/app/services/page_memory/_utils.py +++ b/apps/worker/app/services/page_memory/_utils.py @@ -2,26 +2,22 @@ from __future__ import annotations -import re from dataclasses import dataclass, replace from typing import Any -_NATURAL_RE = re.compile(r"(\d+)") - - -def _natural_key(path: str) -> list[int | str]: - """Split a string into a list of int/str segments for natural ordering.""" - return [int(c) if c.isdigit() else c.lower() for c in _NATURAL_RE.split(path)] - def sort_skeletons(skeletons: list[Any]) -> list[Any]: + """Order skeletons by start page only. + + Same-page relative order is preserved (Python ``sorted`` is stable). That + relative order is authoritative: coarse TOC emit order, or VLM observed + title order from fine hierarchy. Do **not** tie-break by ``section_path`` / + title — alphabetical path order reverses real reading order (e.g. Open + access before Separation on the same page). + """ return sorted( skeletons, - key=lambda item: ( - int(getattr(item, "start_page", 0) or 0), - int(getattr(item, "level", 0) or 0), - _natural_key(str(getattr(item, "section_path", "") or "")), - ), + key=lambda item: int(getattr(item, "start_page", 0) or 0), ) diff --git a/apps/worker/app/services/page_memory/fine_hierarchy.py b/apps/worker/app/services/page_memory/fine_hierarchy.py index 92ea603e..96fd18de 100644 --- a/apps/worker/app/services/page_memory/fine_hierarchy.py +++ b/apps/worker/app/services/page_memory/fine_hierarchy.py @@ -7,8 +7,9 @@ under each coarse TOC leaf. Boundary trimming is code-side: VLM extracts all outline titles in reading -order; ``_collect_candidates`` then drops titles at/before the coarse start -anchor and at/after the next-leaf end anchor. +order; ``_collect_candidates`` drops titles at/before the coarse start anchor +and at/after the next-skeleton end anchor. If the end anchor is missing, +candidates on the scan boundary page are dropped (page-level fallback). """ from __future__ import annotations @@ -22,7 +23,7 @@ from app.services.page_memory.page_tagger import PageTagResult from app.services.page_memory.skeleton_extractor import SectionSkeleton -from app.services.page_memory._utils import page_scope_info +from app.services.page_memory._utils import page_scope_info, sort_skeletons from shared.services.ai.llm_overrides import get_text_client from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response @@ -44,7 +45,7 @@ def refine_fat_leaf_skeletons( For each coarse skeleton that overlaps fat-leaf pages: 1. Collect ``observed_titles`` in reading order and trim by start/end - coarse anchors. + coarse anchors (tail miss → drop the boundary page). 2. Ask the page-memory hierarchy prompt to assign relative levels. 3. Rebuild the nested tree and graft it under the coarse leaf. @@ -127,24 +128,21 @@ def compute_fat_leaf_pages( def build_next_title_by_path( skeletons: list[SectionSkeleton], ) -> dict[str, str | None]: - """Map each leaf ``section_path`` to the next leaf's title (by start_page).""" - ordered = sorted( - skeletons, - key=lambda item: ( - int(getattr(item, "start_page", 0) or 0), - str(getattr(item, "section_path", "") or ""), - ), - ) + """Map each skeleton ``section_path`` to the next title in reading order. + + Order matches ``sort_skeletons``: by ``start_page`` only, stable so + same-page relative order (TOC emit / parent-before-child) is preserved. + The next title is the immediately following skeleton — including a + same-page sibling or a ``parent_self_only`` node's first child — used as + the tail trim anchor in ``_trim_by_coarse_anchors``. + """ + ordered = sort_skeletons(skeletons) next_map: dict[str, str | None] = {} for index, skeleton in enumerate(ordered): next_title: str | None = None - start_page = int(getattr(skeleton, "start_page", 0) or 0) - for later in ordered[index + 1 :]: - later_start = int(getattr(later, "start_page", 0) or 0) - if later_start > start_page: - title = str(getattr(later, "title", "") or "").strip() - next_title = title or None - break + if index + 1 < len(ordered): + title = str(getattr(ordered[index + 1], "title", "") or "").strip() + next_title = title or None next_map[str(skeleton.section_path)] = next_title return next_map @@ -163,11 +161,14 @@ def _collect_candidates( ) -> list[dict[str, Any]]: """Gather VLM-observed titles within a skeleton's page range. - Preserves page ascending order and within-page VLM reading order, then - trims by coarse start/end anchors: + Preserves page ascending order and within-page VLM reading order + (``observed_titles`` as returned by the VLM; prominence is metadata only). + Then trims by coarse start/end anchors: - Drop the start anchor and everything before it. - - Drop the end anchor (next leaf title) and everything after it. + - Drop the end anchor (next skeleton title) and everything after it. + - If the end anchor is missing, drop all candidates on ``exclusive_end`` + (page-level fallback so the next section cannot bleed in). Returns list of ``{id, heading, page, prominence}``. """ @@ -180,7 +181,6 @@ def _collect_candidates( tag = tag_by_page.get(page) if tag is None or not tag.observed_titles: continue - # Keep VLM reading order; do not re-sort by prominence. for title_entry in tag.observed_titles: text = str(title_entry.get("text", "") or "").strip() if not text or len(text) < 2: @@ -200,6 +200,7 @@ def _collect_candidates( start_title=start_title, end_title=end_title, section_path=skeleton.section_path, + boundary_page=exclusive_end, ) candidates: list[dict[str, Any]] = [] @@ -226,8 +227,14 @@ def _trim_by_coarse_anchors( start_title: str, end_title: str | None, section_path: str, + boundary_page: int, ) -> list[dict[str, Any]]: - """Hard-trim reading-order titles by start/end coarse anchors.""" + """Hard-trim reading-order titles by start/end coarse anchors. + + Tail miss policy: if ``end_title`` is set but not found among candidates, + drop every entry on ``boundary_page`` (the scope's last scanned page). + Head miss stays lenient — only drop exact parent-title duplicates. + """ start_key = _title_key(start_title) end_key = _title_key(end_title) if end_title else "" @@ -243,7 +250,6 @@ def _trim_by_coarse_anchors( "keeping all candidates before end trim", section_path, ) - # Fallback: drop exact parent-title duplicates if present. trimmed = [item for item in trimmed if item.get("key") != start_key] else: trimmed = trimmed[start_idx + 1 :] @@ -256,10 +262,14 @@ def _trim_by_coarse_anchors( if end_idx is None: logger.warning( "[page_memory.fine_hierarchy] end anchor {!r} not found for {}; " - "skipping back trim", + "dropping all candidates on boundary page {}", end_title, section_path, + boundary_page, ) + trimmed = [ + item for item in trimmed if int(item.get("page") or 0) != boundary_page + ] else: trimmed = trimmed[:end_idx] @@ -427,6 +437,18 @@ def _run_hierarchy_on_candidates( def _exclusive_end(skeletons: list[SectionSkeleton], index: int) -> int: + """Last page to *scan* for title candidates under ``skeletons[index]``. + + When later starts are visible in ``skeletons``, the span stops before the + next start (``next.start - 1``). For a single-leaf scope list the closed + ``end_page`` is returned, which may be the shared boundary page with the + next coarse leaf — that page is scanned here; candidate ownership on it is + decided later by ``_trim_by_coarse_anchors`` via the global next-title + anchor (not by this helper). + + Distinct from ``node_assembler._exclusive_end``, which decides body-text + ownership / summary slices at assembly time. + """ skeleton = skeletons[index] later_starts = [ skel.start_page diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index bf6be444..43f66402 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -22,7 +22,10 @@ TitleMatch, TitleNode, extract_toc_nodes, + first_leaf_start_under, iter_leaf_title_nodes, + last_leaf_start_under, + locate_title_compact_strict, resolve_hierarchy_page_ranges, ) from app.services.document_parser.structure.body_boundary import ( @@ -146,15 +149,6 @@ def extract_section_skeletons( ] nodes = toc_nodes - # TODO(page-memory): locate null-page coarse TOC parents (self-only spans). - # Leaf-only closed-closed ranges currently end at the next leaf start, so - # content under a null-page parent (e.g. "Section A Governing requirements") - # before its first child is attributed to the previous leaf. Candidate fix: - # after calibration, place each null-page coarse node A in - # (last leaf under previous same-level coarse, first leaf under A], then - # bound fine-hierarchy scopes with current_leaf → next_coarse. - # See skill pdf-page-based-track Open TODOs. - # Collapse degenerate single-child intermediate chains before locate. # Rule: only merge a parent with its only child when that child is NOT a # leaf (i.e. the child still has children of its own). This preserves the @@ -201,7 +195,7 @@ def extract_section_skeletons( ) ] - # Phase A3: try offset-guided bulk anchoring before expensive residual agent. + # Phase A3: offset-guided bulk anchoring for printed-page leaves. offset_matches: dict[tuple[str, ...], TitleMatch] | None = None if offset_hint is not None and ctx is not None: offset_matches = _offset_guided_anchoring( @@ -230,13 +224,30 @@ def extract_section_skeletons( } resolve_nodes = nodes + match_overrides, null_page_report = locate_null_page_parent_overrides( + nodes=resolve_nodes, + match_overrides=match_overrides, + page_texts=page_texts, + body_pages=primary_body_pages, + ctx=ctx, + ) + locate_summary["null_page_parent_locate"] = { + "attempted": len(null_page_report), + "located": sum(1 for row in null_page_report if row.get("page") is not None), + "unresolved": sum( + 1 for row in null_page_report if row.get("result") == "unresolved" + ), + "visual_verify_calls": sum( + int(row.get("visual_verify_calls") or 0) for row in null_page_report + ), + "entries": null_page_report, + } ranges = resolve_hierarchy_page_ranges( resolve_nodes, page_count=primary_page_count, page_texts=page_texts, body_pages=primary_body_pages, - page_offset_hint=offset_hint, match_overrides=match_overrides, ) if not ranges: @@ -490,8 +501,207 @@ def _body_pages(*, anatomy: Any | None, page_count: int) -> list[int]: return [page for page in range(1, page_count + 1) if page not in excluded] -# ── VLM offset calibration (Phase A1) ─────────────────────────────────────── +# ── Null-page parent locate (compact-strict + RTL visual) ─────────────────── + +_NULL_PARENT_VISUAL_CONFIDENCE = 0.6 + + +def locate_null_page_parent_overrides( + *, + nodes: list[TitleNode], + match_overrides: dict[tuple[str, ...], TitleMatch], + page_texts: dict[int, str], + body_pages: list[int], + ctx: ToolContext | None, +) -> tuple[dict[tuple[str, ...], TitleMatch], list[dict[str, Any]]]: + """Locate TOC parents with ``printed_page=None`` into ``match_overrides``. + + Window for parent P: ``[last leaf start under previous same-level sibling, + first leaf start under P]``. Text path is compact→strict unique page; on + miss/ambiguity, scan right→left with ``verify_section_page_choice``. + + Returns ``(overrides, report)`` where *report* lists every null-page parent + attempt (for debug / LLM-call accounting). + """ + if not nodes or not body_pages: + return dict(match_overrides), [] + + out = dict(match_overrides) + body_set = set(body_pages) + parent_scope_start = body_pages[0] + report: list[dict[str, Any]] = [] + + def walk( + sibling_nodes: list[TitleNode], + parent_titles: tuple[str, ...], + scope_start: int, + ) -> None: + for index, node in enumerate(sibling_nodes): + path_titles = (*parent_titles, node.title) + if ( + node.children + and node.printed_page is None + and path_titles not in out + ): + if index > 0: + left = last_leaf_start_under( + sibling_nodes[index - 1], parent_titles, out + ) + if left is None: + left = scope_start + else: + left = scope_start + right = first_leaf_start_under(node, parent_titles, out) + entry: dict[str, Any] = { + "path_titles": list(path_titles), + "title": node.title, + "printed_page": None, + "window": None, + "result": "skipped_no_right", + "page": None, + "accept": None, + "visual_verify_calls": 0, + } + if right is None or right < left: + report.append(entry) + logger.info( + "[page_memory.skeleton] null-page parent skipped: " + "title={!r} reason=no_located_first_child left={}", + node.title, + left, + ) + else: + entry["window"] = [left, right] + scope_pages = [ + page for page in body_pages if left <= page <= right + ] + match = locate_title_compact_strict( + node.title, + scope_pages=scope_pages, + page_texts=page_texts, + ) + visual_calls = 0 + if match is None and ctx is not None: + match, visual_calls = _visual_rtl_locate_parent( + title=node.title, + left=left, + right=right, + body_set=body_set, + ctx=ctx, + ) + entry["visual_verify_calls"] = visual_calls + if match is not None and match.page in body_set: + out[path_titles] = match + entry["result"] = str(match.evidence.get("accept") or match.source) + entry["page"] = match.page + entry["accept"] = match.evidence.get("accept") + logger.info( + "[page_memory.skeleton] null-page parent located: " + "title={!r} page={} window={} accept={} visual_calls={}", + node.title, + match.page, + [left, right], + match.evidence.get("accept"), + visual_calls, + ) + else: + entry["result"] = "unresolved" + logger.info( + "[page_memory.skeleton] null-page parent unresolved: " + "title={!r} window={} visual_calls={}", + node.title, + [left, right], + visual_calls, + ) + report.append(entry) + if node.children: + child_scope_start = ( + out[path_titles].page if path_titles in out else scope_start + ) + walk(node.children, path_titles, child_scope_start) + + walk(nodes, (), parent_scope_start) + logger.info( + "[page_memory.skeleton] null-page parent locate summary: " + "attempted={} located={} unresolved={} visual_verify_calls={}", + len(report), + sum(1 for row in report if row.get("page") is not None), + sum(1 for row in report if row.get("result") == "unresolved"), + sum(int(row.get("visual_verify_calls") or 0) for row in report), + ) + return out, report + + +def _visual_rtl_locate_parent( + *, + title: str, + left: int, + right: int, + body_set: set[int], + ctx: ToolContext, +) -> tuple[TitleMatch | None, int]: + """Confirm parent title from right boundary toward left via VLM verify.""" + visual_calls = 0 + for page in range(right, left - 1, -1): + if page not in body_set: + continue + candidate = TitleMatch( + page=page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[page], + evidence={"null_page_parent_probe": True}, + ) + visual_calls += 1 + result = verify_section_page_choice( + ctx=ctx, + title=title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + selected = result.get("selected_page") + confidence = float(result.get("confidence") or 0.0) + if selected != page or confidence < _NULL_PARENT_VISUAL_CONFIDENCE: + continue + if result.get("source") == "agent_vlm": + return ( + TitleMatch( + page=page, + confidence=confidence, + source="agent_vlm", + matched_line="", + score=confidence, + candidates=[page], + evidence={ + "accept": "visual_rtl", + "reason": result.get("reason", ""), + "visual_verify_calls": visual_calls, + }, + ), + visual_calls, + ) + return ( + TitleMatch( + page=page, + confidence=confidence, + source="agent_heuristic", + matched_line="", + score=confidence, + candidates=[page], + evidence={ + "accept": "visual_rtl", + "reason": result.get("reason", ""), + "visual_verify_calls": visual_calls, + }, + ), + visual_calls, + ) + return None, visual_calls + +# ── VLM offset calibration (Phase A1) ─────────────────────────────────────── _CALIBRATION_WINDOW_PAGES = 10 _CALIBRATION_LEAF_PROBE_COUNT = 3 @@ -991,12 +1201,30 @@ def _resolve_pending_tocs( "reason": "offset_guided_anchoring_skipped_or_empty", } + match_overrides, null_page_report = locate_null_page_parent_overrides( + nodes=nodes, + match_overrides=match_overrides, + page_texts=page_texts, + body_pages=toc_body_pages, + ctx=ctx, + ) + locate_summary["null_page_parent_locate"] = { + "attempted": len(null_page_report), + "located": sum(1 for row in null_page_report if row.get("page") is not None), + "unresolved": sum( + 1 for row in null_page_report if row.get("result") == "unresolved" + ), + "visual_verify_calls": sum( + int(row.get("visual_verify_calls") or 0) for row in null_page_report + ), + "entries": null_page_report, + } + ranges = resolve_hierarchy_page_ranges( nodes, page_count=toc_scope_end, page_texts=page_texts, body_pages=toc_body_pages, - page_offset_hint=offset, match_overrides=match_overrides, ) diff --git a/apps/worker/tests/contract/test_page_memory_collapse_single_child_contract.py b/apps/worker/tests/contract/test_page_memory_collapse_single_child_contract.py index 0401b2f0..dec0f13c 100644 --- a/apps/worker/tests/contract/test_page_memory_collapse_single_child_contract.py +++ b/apps/worker/tests/contract/test_page_memory_collapse_single_child_contract.py @@ -213,8 +213,8 @@ def test_generic_label_merges_regardless_of_content() -> None: # ── Test 5: result is sorted consistently ──────────────────────────── -def test_result_sorted_by_start_page_level_path() -> None: - """Output respects ``_sort_skeletons`` ordering so downstream is stable.""" +def test_result_sorted_by_start_page_preserving_same_page_order() -> None: + """Output is ordered by start_page; same-page relative order is preserved.""" skeletons = [ _skel( "doc.pdf/Z", diff --git a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py index 75eb17b9..c78ba9dc 100644 --- a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py +++ b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py @@ -239,3 +239,110 @@ def test_refine_fat_leaf_keeps_slash_in_title_as_single_path_segment( assert refined[0].section_path == "manual.pdf/Index/Symbols\u2215Numbers" assert refined[1].parent_path == "manual.pdf/Index" assert refined[0].section_path.count("/") == 2 + + +def test_build_next_title_by_path_preserves_parent_before_same_page_child() -> None: + parent = SectionSkeleton( + section_path="demo.pdf/Section Z Parent", + level=1, + start_page=10, + end_page=12, + title="Section Z Parent", + parent_path="demo.pdf", + evidence={"skeleton_kind": "parent_self_only"}, + ) + # Alphabetically sorts before the parent path; stable start_page order + # must still keep emit order (parent first). + child = SectionSkeleton( + section_path="demo.pdf/Section Z Parent/A First Child", + level=2, + start_page=10, + end_page=12, + title="A First Child", + parent_path="demo.pdf/Section Z Parent", + ) + sibling = SectionSkeleton( + section_path="demo.pdf/Later", + level=1, + start_page=13, + end_page=15, + title="Later", + parent_path="demo.pdf", + ) + + next_map = fine_hierarchy.build_next_title_by_path([parent, child, sibling]) + + assert next_map[parent.section_path] == "A First Child" + assert next_map[child.section_path] == "Later" + assert next_map[sibling.section_path] is None + + +def test_trim_drops_boundary_page_when_end_anchor_missing() -> None: + raw = [ + {"heading": "Keep", "page": 10, "key": "keep"}, + {"heading": "Boundary Leftover", "page": 12, "key": "boundaryleftover"}, + {"heading": "Also Boundary", "page": 12, "key": "alsoboundary"}, + ] + + trimmed = fine_hierarchy._trim_by_coarse_anchors( + raw, + start_title="", + end_title="Next Section Title", + section_path="demo.pdf/Current", + boundary_page=12, + ) + + assert [item["heading"] for item in trimmed] == ["Keep"] + + +def test_refine_fat_leaf_drops_boundary_page_on_tail_miss(monkeypatch) -> None: + previous = SectionSkeleton( + section_path="demo.pdf/History", + level=2, + start_page=14, + end_page=23, + title="History", + parent_path="demo.pdf", + ) + next_section = SectionSkeleton( + section_path="demo.pdf/List of amendments", + level=2, + start_page=23, + end_page=30, + title="List of amendments", + parent_path="demo.pdf", + ) + tags = [ + PageTagResult( + page_index=22, + observed_titles=[{"text": "NCC 2022", "prominence": 1.0}], + ), + PageTagResult( + page_index=23, + observed_titles=[ + {"text": "Not The Next Title", "prominence": 1.0}, + {"text": "Another Boundary Title", "prominence": 0.9}, + ], + ), + ] + + monkeypatch.setattr( + fine_hierarchy, + "get_text_client", + lambda requested_model=None: ( + _FakeClient([{"id": 1, "level": 1}]), + requested_model, + ), + ) + + refined = fine_hierarchy.refine_fat_leaf_skeletons( + coarse_skeletons=[previous], + tag_results=tags, + fat_leaf_pages={22, 23}, + next_title_by_path={previous.section_path: next_section.title}, + model_name="test-model", + ) + + assert [item.title for item in refined] == ["NCC 2022"] + assert all("amendment" not in item.title.casefold() for item in refined) + assert all(item.start_page != 23 for item in refined) diff --git a/apps/worker/tests/unit/test_null_page_parent_locate.py b/apps/worker/tests/unit/test_null_page_parent_locate.py new file mode 100644 index 00000000..edbc2915 --- /dev/null +++ b/apps/worker/tests/unit/test_null_page_parent_locate.py @@ -0,0 +1,211 @@ +"""Unit tests for null-page parent locate (compact-strict + self-only emit).""" + +from __future__ import annotations + +import os +from typing import Any +from unittest.mock import MagicMock + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from app.services.document_agent.structure.hierarchy_locator import ( + TitleMatch, + TitleNode, + locate_title_compact_strict, + resolve_hierarchy_page_ranges, +) +from app.services.page_memory.skeleton_extractor import ( + locate_null_page_parent_overrides, +) + + +def _leaf_match(page: int) -> TitleMatch: + return TitleMatch( + page=page, + confidence=0.9, + source="agent_vlm", + matched_line="", + score=0.9, + candidates=[page], + evidence={}, + ) + + +def _example_tree() -> list[TitleNode]: + return [ + TitleNode( + title="Chapter 1", + level=1, + printed_page=10, + children=[ + TitleNode(title="1.1 Intro", level=2, printed_page=10), + TitleNode(title="1.2 Details", level=2, printed_page=12), + ], + ), + TitleNode( + title="Chapter 2 Overview", + level=1, + printed_page=None, + children=[ + TitleNode(title="2.1 Setup", level=2, printed_page=15), + TitleNode(title="2.2 Results", level=2, printed_page=18), + ], + ), + ] + + +def _example_leaf_overrides() -> dict[tuple[str, ...], TitleMatch]: + return { + ("Chapter 1", "1.1 Intro"): _leaf_match(12), + ("Chapter 1", "1.2 Details"): _leaf_match(14), + ("Chapter 2 Overview", "2.1 Setup"): _leaf_match(17), + ("Chapter 2 Overview", "2.2 Results"): _leaf_match(20), + } + + +def test_compact_strict_matches_cross_line_title() -> None: + match = locate_title_compact_strict( + "Chapter 2 Overview", + scope_pages=[14, 15, 16, 17], + page_texts={ + 14: "1.2 Details body", + 15: "Chapter 2\nOverview\nIntro paragraph", + 16: "more intro", + 17: "2.1 Setup", + }, + ) + assert match is not None + assert match.page == 15 + assert match.evidence.get("accept") == "compact_strict_unique" + + +def test_compact_strict_rejects_ambiguous_pages() -> None: + match = locate_title_compact_strict( + "Chapter 2 Overview", + scope_pages=[14, 15, 16, 17], + page_texts={ + 14: "x", + 15: "Chapter 2\nOverview", + 16: "Chapter 2 Overview again", + 17: "2.1 Setup", + }, + ) + assert match is None + + +def test_resolve_emits_parent_self_only_and_truncates_prev_leaf() -> None: + nodes = _example_tree() + overrides = _example_leaf_overrides() + overrides[("Chapter 2 Overview",)] = TitleMatch( + page=15, + confidence=0.92, + source="anchored", + matched_line="chapter2overview", + score=0.96, + candidates=[15], + evidence={"accept": "compact_strict_unique"}, + ) + page_texts = {page: "" for page in range(12, 21)} + ranges = resolve_hierarchy_page_ranges( + nodes, + page_count=20, + page_texts=page_texts, + body_pages=list(range(12, 21)), + match_overrides=overrides, + ) + by_path = {item.path_titles: item for item in ranges} + + details = by_path[("Chapter 1", "1.2 Details")] + assert details.end_page == 15 + + parent = by_path[("Chapter 2 Overview",)] + assert parent.start_page == 15 + assert parent.end_page == 17 + assert parent.evidence.get("skeleton_kind") == "parent_self_only" + + setup = by_path[("Chapter 2 Overview", "2.1 Setup")] + assert setup.start_page == 17 + + +def test_locate_null_page_parent_overrides_uses_compact_window() -> None: + nodes = _example_tree() + overrides, report = locate_null_page_parent_overrides( + nodes=nodes, + match_overrides=_example_leaf_overrides(), + page_texts={ + 14: "1.2 body", + 15: "Chapter 2\nOverview", + 16: "intro", + 17: "2.1 Setup", + 20: "2.2 Results", + }, + body_pages=list(range(12, 21)), + ctx=None, + ) + parent = overrides[("Chapter 2 Overview",)] + assert parent.page == 15 + assert parent.evidence.get("accept") == "compact_strict_unique" + assert any(row["title"] == "Chapter 2 Overview" and row["page"] == 15 for row in report) + + +def test_locate_null_page_parent_visual_rtl_when_text_ambiguous( + monkeypatch: Any, +) -> None: + nodes = _example_tree() + calls: list[int] = [] + + def _fake_verify( + *, + ctx: Any, + title: str, + candidate_matches: list[TitleMatch], + candidate_page_cap: int, + ) -> dict[str, Any]: + page = candidate_matches[0].page + if title == "Chapter 2 Overview": + calls.append(page) + if title == "Chapter 2 Overview" and page == 15: + return { + "selected_page": 15, + "confidence": 0.8, + "source": "agent_vlm", + "reason": "title heading", + } + return { + "selected_page": None, + "confidence": 0.2, + "source": "agent_vlm", + "reason": "not here", + } + + monkeypatch.setattr( + "app.services.page_memory.skeleton_extractor.verify_section_page_choice", + _fake_verify, + ) + overrides, report = locate_null_page_parent_overrides( + nodes=nodes, + match_overrides=_example_leaf_overrides(), + page_texts={ + 14: "x", + 15: "Chapter 2\nOverview", + 16: "Chapter 2 Overview", + 17: "2.1 Setup", + 20: "2.2", + }, + body_pages=list(range(12, 21)), + ctx=MagicMock(), + ) + parent = overrides[("Chapter 2 Overview",)] + assert parent.page == 15 + assert parent.evidence.get("accept") == "visual_rtl" + assert calls[0] == 17 + assert 15 in calls + assert calls.index(15) > calls.index(17) + entry = next(row for row in report if row["title"] == "Chapter 2 Overview") + assert entry["visual_verify_calls"] >= 1 + assert entry["result"] == "visual_rtl" diff --git a/apps/worker/tests/unit/test_sort_skeletons_preserves_vlm_order.py b/apps/worker/tests/unit/test_sort_skeletons_preserves_vlm_order.py new file mode 100644 index 00000000..8fea71b4 --- /dev/null +++ b/apps/worker/tests/unit/test_sort_skeletons_preserves_vlm_order.py @@ -0,0 +1,94 @@ +"""sort_skeletons must not alphabetize same-page VLM/TOC order.""" + +from __future__ import annotations + +from app.services.page_memory._serialization import ( + build_hierarchy_tree, + serialize_hierarchy_artifact, +) +from app.services.page_memory._utils import sort_skeletons +from app.services.page_memory.skeleton_extractor import SectionSkeleton + + +def _skel(path: str, title: str, *, start_page: int, level: int = 3) -> SectionSkeleton: + parent = "/".join(path.split("/")[:-1]) or None + return SectionSkeleton( + section_path=path, + title=title, + level=level, + start_page=start_page, + end_page=start_page, + parent_path=parent, + evidence={}, + ) + + +def test_same_page_siblings_keep_input_order_not_alphabetical() -> None: + base = "doc.pdf/Section D/Part D3 Construction of exits" + # VLM / reading order: Separation (D3D5) before Open access (D3D6) + page_197 = [ + _skel( + f"{base}/Non-fire-isolated stairways and ramps", + "Non-fire-isolated stairways and ramps", + start_page=197, + ), + _skel( + f"{base}/Separation of rising and descending stair flights", + "Separation of rising and descending stair flights", + start_page=197, + ), + _skel( + f"{base}/Open access ramps and balconies", + "Open access ramps and balconies", + start_page=197, + ), + _skel(f"{base}/Smoke lobbies", "Smoke lobbies", start_page=197), + ] + # Later page first in input — page sort must move it after, without + # alphabetizing the same-page block. + mixed = [ + _skel(f"{base}/Landings", "Landings", start_page=201), + *page_197, + ] + titles = [s.title for s in sort_skeletons(mixed)] + assert titles == [ + "Non-fire-isolated stairways and ramps", + "Separation of rising and descending stair flights", + "Open access ramps and balconies", + "Smoke lobbies", + "Landings", + ] + # Alphabetical would put Open before Separation — must not happen + assert titles.index("Separation of rising and descending stair flights") < titles.index( + "Open access ramps and balconies" + ) + + +def test_hierarchy_tree_matches_node_list_order() -> None: + base = "doc.pdf/Part D3" + skeletons = [ + _skel( + f"{base}/Separation of rising and descending stair flights", + "Separation of rising and descending stair flights", + start_page=197, + ), + _skel( + f"{base}/Open access ramps and balconies", + "Open access ramps and balconies", + start_page=197, + ), + ] + tree = build_hierarchy_tree(skeletons) + assert list(tree["Part D3"].keys()) == [ + "Separation of rising and descending stair flights", + "Open access ramps and balconies", + ] + + artifact = serialize_hierarchy_artifact(skeletons) + assert list(artifact["HIERARCHY"]["Part D3"].keys()) == [ + n["title"] for n in artifact["nodes"] + ] + assert [n["title"] for n in artifact["nodes"]] == [ + "Separation of rising and descending stair flights", + "Open access ramps and balconies", + ] diff --git a/packages/shared-python/shared/models/schemas/page_memory_config.py b/packages/shared-python/shared/models/schemas/page_memory_config.py index 8e01fce3..a1bc6090 100644 --- a/packages/shared-python/shared/models/schemas/page_memory_config.py +++ b/packages/shared-python/shared/models/schemas/page_memory_config.py @@ -29,11 +29,6 @@ class PageMemoryConfig: table_engine: str = "tabula" table_merge_enabled: bool = True node_summary_max_pages: int = 5 - page_locate_residual_agent_limit: int = 50 - page_locate_max_emit_depth: int = 5 - page_locate_min_emit_depth: int = 2 - page_locate_vlm_candidate_page_cap: int = 4 - page_locate_full_leaf_sections: bool = False scan_direction: str = "top_to_bottom_left_to_right" @classmethod @@ -128,26 +123,6 @@ def from_mapping(cls, value: object) -> Self: value.get("node_summary_max_pages"), default.node_summary_max_pages, ), - page_locate_residual_agent_limit=_as_int( - value.get("page_locate_residual_agent_limit"), - default.page_locate_residual_agent_limit, - ), - page_locate_max_emit_depth=_as_int( - value.get("page_locate_max_emit_depth"), - default.page_locate_max_emit_depth, - ), - page_locate_min_emit_depth=_as_int( - value.get("page_locate_min_emit_depth"), - default.page_locate_min_emit_depth, - ), - page_locate_vlm_candidate_page_cap=_as_int( - value.get("page_locate_vlm_candidate_page_cap"), - default.page_locate_vlm_candidate_page_cap, - ), - page_locate_full_leaf_sections=_as_bool( - value.get("page_locate_full_leaf_sections"), - default.page_locate_full_leaf_sections, - ), scan_direction=_as_str( value.get("scan_direction"), default.scan_direction ),