From 834f8d2ce4e2e8ac11e9f0908284ad5f62ab02c2 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 20 Jul 2026 15:04:41 +0800 Subject: [PATCH 1/4] refactor: remove residual page location agent and related tools This commit eliminates the page locate sub-agent and its associated tools, streamlining the hierarchy locator module. The VLM page verification process is now directly integrated into the hierarchy locator, enhancing efficiency and reducing complexity. Unused configurations and functions related to the residual agent have also been removed from the page memory configuration. --- .../structure/hierarchy_locator.py | 364 +------------- .../structure/page_locate_agent.py | 242 +--------- .../structure/page_locate_subagent.py | 446 ------------------ .../structure/page_locate_tools.py | 158 ------- .../services/document_agent/tools/__init__.py | 1 - .../document_agent/tools/page_locate.py | 16 - .../models/schemas/page_memory_config.py | 25 - 7 files changed, 14 insertions(+), 1238 deletions(-) delete mode 100644 apps/worker/app/services/document_agent/structure/page_locate_subagent.py delete mode 100644 apps/worker/app/services/document_agent/structure/page_locate_tools.py delete mode 100644 apps/worker/app/services/document_agent/tools/page_locate.py 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..bf55008b 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -1,14 +1,13 @@ """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. +This module is intentionally deterministic: it performs strict title anchoring +and range assembly. Titles are accepted only via strict exact anchoring or +offset-guided ``match_overrides``; weaker candidate collection has been removed +in favor of VLM-verified calibration upstream. """ from __future__ import annotations -import re from dataclasses import dataclass, field from typing import Any, Literal @@ -19,10 +18,6 @@ TitleMatchSource = Literal[ "anchored", - "page_compact", - "normalized", - "token", - "printed_prior", "h1_result", "agent_vlm", "agent_heuristic", @@ -80,47 +75,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, *, @@ -141,80 +95,6 @@ def locate_title_strict_exact( ) -def collect_title_candidate_matches( - 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 [] - - 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() - ] - - 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, - }, - ) - ) - - 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, - ) - if limit is not None: - return matches[: max(int(limit), 0)] - return matches - - def resolve_hierarchy_page_ranges( nodes: list[TitleNode], *, @@ -223,7 +103,6 @@ def resolve_hierarchy_page_ranges( 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. @@ -250,7 +129,6 @@ def resolve_hierarchy_page_ranges( 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 @@ -276,7 +154,6 @@ def _resolve_siblings( 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]] = [] @@ -294,14 +171,6 @@ def _resolve_siblings( 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, @@ -323,7 +192,6 @@ def _resolve_siblings( 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: @@ -357,7 +225,6 @@ def _resolve_siblings( page_texts=page_texts, page_offset_hint=page_offset_hint, match_overrides=match_overrides, - use_weak_fallback=use_weak_fallback, resolved=resolved, ) continue @@ -385,7 +252,6 @@ def _find_next_located_sibling( 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) @@ -401,14 +267,6 @@ def _find_next_located_sibling( scope_pages=pages, page_texts=page_texts, ) - 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, @@ -487,50 +345,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, @@ -628,93 +442,6 @@ 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], *, @@ -740,11 +467,9 @@ def sort_key(hit: _LineHit) -> tuple[float, int, int, int]: 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, @@ -763,34 +488,6 @@ def sort_key(hit: _LineHit) -> tuple[float, int, int, int]: ) -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 +507,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/packages/shared-python/shared/models/schemas/page_memory_config.py b/packages/shared-python/shared/models/schemas/page_memory_config.py index bbac71ed..fd5d2ad5 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 @classmethod def default(cls) -> Self: @@ -127,26 +122,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, - ), ) def to_dict(self) -> dict[str, object]: From 1b05f798116aeca5eddce4d7072e62fc9422d665 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 20 Jul 2026 21:44:32 +0800 Subject: [PATCH 2/4] feat: null-page TOC parent locate and VLM-order skeleton sort Add window-bounded locate for coarse TOC parents whose printed_page is null (compact-strict text match with visual RTL fallback), emitting parent self-only skeletons so interstitial content is not swallowed. Sort skeletons by start_page only and build the hierarchy tree from list order, so same-page siblings keep VLM/TOC reading order instead of being reordered alphabetically by section_path. Co-authored-by: Cursor --- .../structure/hierarchy_locator.py | 213 ++++++++++----- .../services/page_memory/_serialization.py | 10 +- .../worker/app/services/page_memory/_utils.py | 22 +- .../page_memory/skeleton_extractor.py | 245 +++++++++++++++++- ...e_memory_collapse_single_child_contract.py | 4 +- .../unit/test_null_page_parent_locate.py | 211 +++++++++++++++ ...test_sort_skeletons_preserves_vlm_order.py | 94 +++++++ 7 files changed, 715 insertions(+), 84 deletions(-) create mode 100644 apps/worker/tests/unit/test_null_page_parent_locate.py create mode 100644 apps/worker/tests/unit/test_sort_skeletons_preserves_vlm_order.py 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 bf55008b..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,13 +1,14 @@ """Locate hierarchy titles on PDF pages and resolve page ranges. -This module is intentionally deterministic: it performs strict title anchoring -and range assembly. Titles are accepted only via strict exact anchoring or -offset-guided ``match_overrides``; weaker candidate collection has been removed -in favor of VLM-verified calibration upstream. +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 +import re from dataclasses import dataclass, field from typing import Any, Literal @@ -89,26 +90,96 @@ 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 locate_title_compact_strict( + title: str, + *, + scope_pages: list[int], + page_texts: dict[int, str], +) -> TitleMatch | None: + """Locate *title* after cross-line compact cleanup; accept only a unique page. + + 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 + + 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] + + 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"}, + ) + + +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( nodes: list[TitleNode], *, 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, ) -> 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 [] @@ -127,7 +198,6 @@ 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 {}, resolved=resolved, ) @@ -152,7 +222,6 @@ 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], resolved: list[ResolvedHierarchyRange], ) -> None: @@ -162,19 +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 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: @@ -190,7 +253,6 @@ 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, parent_titles=parent_titles, ) @@ -217,13 +279,31 @@ 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, resolved=resolved, ) @@ -242,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], @@ -250,27 +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], 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 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 @@ -282,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) @@ -315,6 +409,7 @@ def _infer_start_from_descendant_overrides( evidence={ "inferred_from": "descendant_leaf_override", "original_confidence": min_match.confidence, + "status": "degraded", }, ) @@ -418,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], @@ -446,23 +545,13 @@ 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 = { @@ -481,8 +570,6 @@ 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 {}), }, ) diff --git a/apps/worker/app/services/page_memory/_serialization.py b/apps/worker/app/services/page_memory/_serialization.py index b09e3420..76a4f950 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, ) @@ -92,8 +91,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 = str(getattr(skel, "section_path", "") or "").split("/") section_parts = parts[1:] if len(parts) > 1 else parts current = hierarchy @@ -109,6 +114,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/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 4b917cc2..4d13af19 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 ( @@ -188,7 +191,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( @@ -217,13 +220,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: @@ -473,8 +493,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 @@ -974,12 +1193,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/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", + ] From 242b6deafae179d2f42d10b5d80c4f3f4cb71db1 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 20 Jul 2026 21:51:20 +0800 Subject: [PATCH 3/4] chore: drop stale null-page TODO and duplicate debug import after merge The null-page parent locate is already implemented on this branch, so the placeholder TODO reintroduced by the merge is misleading. Also remove the duplicated page-memory serialization import left by the merge. Co-authored-by: Cursor --- .../app/services/page_memory/skeleton_extractor.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index aa7b050b..43f66402 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -149,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 From f409c4d95e7a19d4c2cc2c6131e7be1c0a8633e3 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 20 Jul 2026 22:35:41 +0800 Subject: [PATCH 4/4] fix: drop fine-hierarchy boundary page when next-title anchor misses Prevent next-section titles from bleeding into the current fat leaf when the VLM fails to observe the coarse end anchor. Align next-title ordering with stable reading order so same-page parent/child anchors stay correct. Co-authored-by: Cursor --- .../services/page_memory/fine_hierarchy.py | 74 +++++++----- ...est_page_memory_fine_hierarchy_contract.py | 107 ++++++++++++++++++ 2 files changed, 155 insertions(+), 26 deletions(-) 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/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)