diff --git a/apps/api/tests/contract/test_agentic_discovery_selection_contract.py b/apps/api/tests/contract/test_agentic_discovery_selection_contract.py index a27dc654..09858a1d 100644 --- a/apps/api/tests/contract/test_agentic_discovery_selection_contract.py +++ b/apps/api/tests/contract/test_agentic_discovery_selection_contract.py @@ -1,4 +1,13 @@ from shared.services.retrieval.agentic.navigation.actions import build_legal_actions +from shared.services.retrieval.agentic.navigation.state import RejectionRecord + + +def _rejected(paths: dict[str, str]) -> dict[str, RejectionRecord]: + """Build a rejection ledger from {path: reason}.""" + return { + path: RejectionRecord(path=path, reason=reason, step=1, detail="") + for path, reason in paths.items() + } def test_discovery_hint_is_projected_as_collect_action() -> None: @@ -14,8 +23,7 @@ def test_discovery_hint_is_projected_as_collect_action() -> None: "chunk_type": "text", } ], - rejected_paths=set(), - rejected_collect_paths=set(), + rejected={}, total_images=0, total_tables=0, budget_snapshot=None, @@ -47,8 +55,7 @@ def test_discovery_hint_under_collected_path_is_not_repeated() -> None: "discovery_score": 0.82, } ], - rejected_paths=set(), - rejected_collect_paths=set(), + rejected={}, total_images=0, total_tables=0, budget_snapshot=None, @@ -57,7 +64,8 @@ def test_discovery_hint_under_collected_path_is_not_repeated() -> None: assert action_set.collect == [] -def test_discovery_hint_under_rejected_collect_path_is_not_repeated() -> None: +def test_discovery_hint_under_tool_adjudicated_path_is_not_repeated() -> None: + """tool_adjudicated rejections are not revived by discovery this round.""" action_set = build_legal_actions( items=[], current_scope=None, @@ -69,11 +77,239 @@ def test_discovery_hint_under_rejected_collect_path_is_not_repeated() -> None: "discovery_score": 0.7, } ], - rejected_paths=set(), - rejected_collect_paths={"1、2016:机构行为助推行情演绎"}, + rejected=_rejected({ + "1、2016:机构行为助推行情演绎": "tool_adjudicated", + }), total_images=0, total_tables=0, budget_snapshot=None, ) assert action_set.collect == [] + + +# ─── NavigationState ledger (Phase 0) ─────────────────────────────────────── + + +def test_mark_rejected_collect_records_tool_adjudicated_reason() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.mark_rejected_collect("Chapter 1", step=3, detail="no matching asset") + + assert "Chapter 1" in state.rejected + record = state.rejected["Chapter 1"] + assert record.reason == "tool_adjudicated" + assert record.step == 3 + assert record.detail == "no matching asset" + + +def test_mark_rejected_if_unproductive_records_navigational_abandon() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.mark_rejected_if_unproductive("Chapter 2", step=5, detail="back_from_unproductive") + + assert state.rejected["Chapter 2"].reason == "navigational_abandon" + + +def test_tool_adjudicated_overrides_weak_abandon_record() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.mark_rejected_if_unproductive("Chapter 3", step=2) + state.mark_rejected_collect("Chapter 3", step=4, detail="asset mismatch") + + # Stronger reason wins. + assert state.rejected["Chapter 3"].reason == "tool_adjudicated" + assert state.rejected["Chapter 3"].step == 4 + + +def test_weak_abandon_does_not_overwrite_strong_record() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.mark_rejected_collect("Chapter 4", step=1) + state.mark_rejected_if_unproductive("Chapter 4", step=5) + + assert state.rejected["Chapter 4"].reason == "tool_adjudicated" + assert state.rejected["Chapter 4"].step == 1 + + +def test_coverage_helpers_derive_from_collected_paths() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.add_collected( + {"path": "A", "hydrate_mode": "chunks", "confidence": 0.9}, + step=1, + scope_context=None, + ) + state.add_collected( + {"path": "B", "hydrate_mode": "outline", "confidence": 0.6}, + step=2, + scope_context=None, + ) + # A path upgraded from outline to full counts as covered, not outline. + state.add_collected( + {"path": "C", "hydrate_mode": "outline", "confidence": 0.5}, + step=3, + scope_context=None, + ) + state.add_collected( + {"path": "C", "hydrate_mode": "chunks", "confidence": 0.8}, + step=4, + scope_context=None, + ) + + assert state.covered_paths() == {"A", "C"} + assert state.outline_paths() == {"B"} + + +def test_snapshot_delta_records_rejection_reasons() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.step_count = 2 + state.mark_rejected_if_unproductive("X", step=2) + + state.step_count = 3 + state.mark_rejected_collect("Y", step=3, detail="asset mismatch") + + delta = state.snapshot_delta( + before_scope=None, + expanded_before=set(), + rejected_before={}, + collected_before_count=0, + ) + rejected_added = {item["path"]: item["reason"] for item in delta["rejected_added"]} + assert rejected_added == {"X": "navigational_abandon", "Y": "tool_adjudicated"} + + +def test_rejected_paths_with_reason_partitions_by_label() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.mark_rejected_collect("A", step=1) + state.mark_rejected_if_unproductive("B", step=1) + state.mark_rejected_collect("C", step=1) + + assert state.rejected_paths_with_reason("tool_adjudicated") == {"A", "C"} + assert state.rejected_paths_with_reason("navigational_abandon") == {"B"} + + +# ─── Reason-aware action filtering (T7-style regression) ──────────────────── + + +def test_tool_adjudicated_rejection_blocks_collect_even_with_discovery() -> None: + """A tool-adjudicated path stays out of COLLECT even when discovery hints it.""" + action_set = build_legal_actions( + items=[], + current_scope=None, + collected_paths=[], + expanded_scopes=set(), + discovery_hints=[{"section_path": "X", "discovery_score": 0.95}], + rejected=_rejected({"X": "tool_adjudicated"}), + total_images=0, + total_tables=0, + budget_snapshot=None, + ) + assert action_set.collect == [] + + +def test_navigational_abandon_is_revived_by_discovery_for_collect() -> None: + """A soft-abandoned path CAN still be COLLECTed when discovery signals it.""" + action_set = build_legal_actions( + items=[], + current_scope=None, + collected_paths=[], + expanded_scopes=set(), + discovery_hints=[{"section_path": "Y", "discovery_score": 0.9}], + rejected=_rejected({"Y": "navigational_abandon"}), + total_images=0, + total_tables=0, + budget_snapshot=None, + ) + assert any(action.path == "Y" for action in action_set.collect) + + +def test_navigational_abandon_suppresses_expand_without_discovery_signal() -> None: + """EXPAND is suppressed for soft-abandoned scopes lacking any discovery signal.""" + items = [{"path": "Z", "level": 1, "is_leaf": False, "chunk_count": 5}] + action_set = build_legal_actions( + items=items, + current_scope=None, + collected_paths=[], + expanded_scopes=set(), + discovery_hints=[], + rejected=_rejected({"Z": "navigational_abandon"}), + total_images=0, + total_tables=0, + budget_snapshot=None, + ) + assert any(action.path == "Z" for action in action_set.collect) + assert not any(action.path == "Z" for action in action_set.expand) + + +def test_navigational_abandon_revives_expand_with_discovery_signal() -> None: + """EXPAND is offered for soft-abandoned scopes when a discovery signal exists.""" + items = [{"path": "Z", "level": 1, "is_leaf": False, "chunk_count": 5}] + action_set = build_legal_actions( + items=items, + current_scope=None, + collected_paths=[], + expanded_scopes=set(), + discovery_hints=[{"section_path": "Z / child", "discovery_score": 0.7}], + rejected=_rejected({"Z": "navigational_abandon"}), + total_images=0, + total_tables=0, + budget_snapshot=None, + ) + assert any(action.path == "Z" for action in action_set.expand) + + +def test_covered_path_excluded_from_actions() -> None: + """Regression: a path already collected as full evidence is not re-offered.""" + items = [{"path": "A", "level": 1, "is_leaf": False, "chunk_count": 3}] + action_set = build_legal_actions( + items=items, + current_scope=None, + collected_paths=[{"path": "A", "hydrate_mode": "chunks"}], + expanded_scopes=set(), + discovery_hints=[{"section_path": "A", "discovery_score": 0.9}], + rejected={}, + total_images=0, + total_tables=0, + budget_snapshot=None, + ) + assert not any(action.path == "A" for action in action_set.collect) + assert not any(action.path == "A" for action in action_set.expand) + 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 3dabcdf0..28697c45 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -302,6 +302,10 @@ def _resolve_siblings( 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, + ) if match is None: start_page = lower_bound else: @@ -405,11 +409,58 @@ def _find_next_located_sibling( 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 +def _infer_start_from_descendant_overrides( + node: TitleNode, + parent_titles: tuple[str, ...], + 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. + """ + if not node.children or not match_overrides: + return None + leaves = iter_leaf_title_nodes([node], parent_titles=parent_titles) + min_page: int | None = None + min_match: TitleMatch | None = None + for leaf_path, _leaf_node in leaves: + m = match_overrides.get(leaf_path) + if m is None: + continue + if m.page not in scope_pages: + continue + if min_page is None or m.page < min_page: + min_page = m.page + min_match = m + if min_match is None: + return None + return TitleMatch( + page=min_match.page, + confidence=min(min_match.confidence, 0.80), + source=min_match.source, + matched_line="", + score=min(min_match.score, 0.80), + candidates=[min_match.page], + evidence={ + "inferred_from": "descendant_leaf_override", + "original_confidence": min_match.confidence, + }, + ) + + def _match_override( path_titles: tuple[str, ...], match_overrides: dict[tuple[str, ...], TitleMatch], 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 d31014f7..ad033f54 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 @@ -132,6 +132,22 @@ def prepare(self, nodes: list[TitleNode]) -> PageLocatePrepareResult: 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], @@ -139,9 +155,10 @@ def _classify_residuals( 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=self.body_pages, + scope_pages=scope, page_texts=self.page_texts, ) if match is not None: @@ -180,9 +197,10 @@ def _resolve_residuals( 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=self.body_pages, + scope_pages=scope, page_count=self.page_count, config=sub_config, page_offset_hint=self.page_offset_hint, diff --git a/apps/worker/app/services/page_memory/_utils.py b/apps/worker/app/services/page_memory/_utils.py index 816dee91..27d9d6c5 100644 --- a/apps/worker/app/services/page_memory/_utils.py +++ b/apps/worker/app/services/page_memory/_utils.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from dataclasses import dataclass, replace from typing import Any _NATURAL_RE = re.compile(r"(\d+)") @@ -57,3 +58,162 @@ def page_scope_info(pages: Any) -> dict[str, Any]: "page_count": len(normalized), "page_ranges": collapse_page_ranges(normalized), } + + +@dataclass(frozen=True) +class CoarseScope: + scope_id: str + skeletons: list[Any] + strategy: str + start_page: int + end_page: int + + +def build_hierarchy_scopes( + *, + skeletons: list[Any], + filename: str, + page_count: int, +) -> list[CoarseScope]: + """Split skeleton list into coarse scopes for independent processing. + + Groups skeletons by top-level section_path prefix. Skeletons not claimed + by any top-node are grouped by their root ancestor (first path segment + after filename) as orphan scopes. + """ + if not skeletons: + return [] + + root_fallback = all( + getattr(item, "title", "") == "Root" + or (getattr(item, "evidence", {}) or {}).get("source") == "fallback_root" + for item in skeletons + ) + if root_fallback: + ordered = sort_skeletons(skeletons) + return [ + CoarseScope( + scope_id=scope_id_for_pages(1, page_count), + skeletons=ordered, + strategy="fallback_root", + start_page=1, + end_page=page_count, + ) + ] + + min_level = min(int(getattr(item, "level", 0) or 0) for item in skeletons) + top_nodes = [ + item + for item in skeletons + if int(getattr(item, "level", 0) or 0) == min_level + or not str(getattr(item, "parent_path", "") or "").startswith(f"{filename}/") + ] + seen_top_paths: set[str] = set() + unique_top_nodes: list[Any] = [] + for item in sort_skeletons(top_nodes): + if item.section_path in seen_top_paths: + continue + seen_top_paths.add(item.section_path) + unique_top_nodes.append(item) + + scopes: list[CoarseScope] = [] + for index, top in enumerate(unique_top_nodes): + prefix = f"{top.section_path}/" + members = [ + item + for item in skeletons + if item.section_path == top.section_path + or str(item.section_path).startswith(prefix) + ] + if not members: + continue + start_page = max(1, int(getattr(top, "start_page", 1) or 1)) + next_top_start = ( + int(getattr(unique_top_nodes[index + 1], "start_page", page_count + 1) or page_count + 1) + if index + 1 < len(unique_top_nodes) + else page_count + 1 + ) + end_page = min( + page_count, + max( + start_page, + min( + int(getattr(top, "end_page", page_count) or page_count), + next_top_start - 1, + ), + ), + ) + bounded_members = [ + replace( + item, + start_page=max( + start_page, + int(getattr(item, "start_page", start_page) or start_page), + ), + end_page=min( + end_page, + int(getattr(item, "end_page", end_page) or end_page), + ), + ) + for item in members + if int(getattr(item, "start_page", 1) or 1) <= end_page + and int(getattr(item, "end_page", end_page) or end_page) >= start_page + ] + bounded_members = sort_skeletons(bounded_members) + scopes.append( + CoarseScope( + scope_id=scope_id_for_pages(start_page, end_page), + skeletons=bounded_members, + strategy=f"coarse_scope_{index + 1}", + start_page=start_page, + end_page=end_page, + ) + ) + + if scopes: + claimed_paths: set[str] = set() + for scope in scopes: + for item in scope.skeletons: + claimed_paths.add(item.section_path) + orphans = [item for item in skeletons if item.section_path not in claimed_paths] + if orphans: + root_groups: dict[str, list[Any]] = {} + for item in orphans: + parts = str(getattr(item, "section_path", "") or "").split("/") + root_key = "/".join(parts[:2]) if len(parts) >= 2 else item.section_path + root_groups.setdefault(root_key, []).append(item) + for root_key, group in sorted(root_groups.items()): + group = sort_skeletons(group) + sp = max(1, min(int(getattr(g, "start_page", 1) or 1) for g in group)) + ep = min(page_count, max(int(getattr(g, "end_page", page_count) or page_count) for g in group)) + scopes.append( + CoarseScope( + scope_id=scope_id_for_pages(sp, ep), + skeletons=group, + strategy="orphan_scope", + start_page=sp, + end_page=ep, + ) + ) + return scopes + + ordered = sort_skeletons(skeletons) + all_pages: set[int] = set() + for item in ordered: + sp = max(1, int(getattr(item, "start_page", 1) or 1)) + ep = min(page_count, int(getattr(item, "end_page", sp) or sp)) + if ep >= sp: + all_pages.update(range(sp, ep + 1)) + pages = sorted(all_pages) or list(range(1, page_count + 1)) + return [ + CoarseScope( + scope_id=scope_id_for_pages( + pages[0] if pages else 1, + pages[-1] if pages else page_count, + ), + skeletons=ordered, + strategy="full_coarse_hierarchy", + start_page=pages[0] if pages else 1, + end_page=pages[-1] if pages else page_count, + ) + ] diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 9e697aed..d1397bc3 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -2,7 +2,7 @@ import os import shutil -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -21,9 +21,10 @@ from app.services.document_parser.support.stage_profiler import stage_timer from app.services.page_memory.normalizer import normalize_to_pdf from app.services.page_memory._utils import ( + CoarseScope, + build_hierarchy_scopes, collapse_page_ranges, page_scope_info, - scope_id_for_pages, sort_skeletons, ) from app.services.page_memory._serialization import ( @@ -52,13 +53,7 @@ class PageMemoryInput: ) -@dataclass(frozen=True) -class _HierarchyScope: - scope_id: str - skeletons: list[Any] - strategy: str - start_page: int - end_page: int +_HierarchyScope = CoarseScope @dataclass(frozen=True) @@ -263,7 +258,6 @@ def _build_page_dataframe( filename=filename, page_texts=page_texts, ctx=ctx, - page_memory_config=page_memory_config, ) else: skeletons = [] @@ -505,111 +499,11 @@ def _build_hierarchy_scopes( filename: str, page_count: int, ) -> list[_HierarchyScope]: - if not skeletons: - return [] - - root_fallback = all( - getattr(item, "title", "") == "Root" - or (getattr(item, "evidence", {}) or {}).get("source") == "fallback_root" - for item in skeletons + return build_hierarchy_scopes( + skeletons=skeletons, + filename=filename, + page_count=page_count, ) - if root_fallback: - ordered = sort_skeletons(skeletons) - return [ - _HierarchyScope( - scope_id=scope_id_for_pages(1, page_count), - skeletons=ordered, - strategy="fallback_root", - start_page=1, - end_page=page_count, - ) - ] - - min_level = min(int(getattr(item, "level", 0) or 0) for item in skeletons) - top_nodes = [ - item - for item in skeletons - if int(getattr(item, "level", 0) or 0) == min_level - or not str(getattr(item, "parent_path", "") or "").startswith(f"{filename}/") - ] - seen_top_paths: set[str] = set() - unique_top_nodes: list[Any] = [] - for item in sort_skeletons(top_nodes): - if item.section_path in seen_top_paths: - continue - seen_top_paths.add(item.section_path) - unique_top_nodes.append(item) - - scopes: list[_HierarchyScope] = [] - for index, top in enumerate(unique_top_nodes): - prefix = f"{top.section_path}/" - members = [ - item - for item in skeletons - if item.section_path == top.section_path - or str(item.section_path).startswith(prefix) - ] - if not members: - continue - start_page = max(1, int(getattr(top, "start_page", 1) or 1)) - next_top_start = ( - int(getattr(unique_top_nodes[index + 1], "start_page", page_count + 1) or page_count + 1) - if index + 1 < len(unique_top_nodes) - else page_count + 1 - ) - end_page = min( - page_count, - max( - start_page, - min( - int(getattr(top, "end_page", page_count) or page_count), - next_top_start - 1, - ), - ), - ) - bounded_members = [ - replace( - item, - start_page=max( - start_page, - int(getattr(item, "start_page", start_page) or start_page), - ), - end_page=min( - end_page, - int(getattr(item, "end_page", end_page) or end_page), - ), - ) - for item in members - if int(getattr(item, "start_page", 1) or 1) <= end_page - and int(getattr(item, "end_page", end_page) or end_page) >= start_page - ] - bounded_members = sort_skeletons(bounded_members) - scopes.append( - _HierarchyScope( - scope_id=scope_id_for_pages(start_page, end_page), - skeletons=bounded_members, - strategy=f"coarse_scope_{index + 1}", - start_page=start_page, - end_page=end_page, - ) - ) - if scopes: - return scopes - - ordered = sort_skeletons(skeletons) - pages = _derive_hierarchy_page_scope(skeletons=ordered, page_count=page_count) - return [ - _HierarchyScope( - scope_id=scope_id_for_pages( - pages[0] if pages else 1, - pages[-1] if pages else page_count, - ), - skeletons=ordered, - strategy="full_coarse_hierarchy", - start_page=pages[0] if pages else 1, - end_page=pages[-1] if pages else page_count, - ) - ] def _allocate_asset_pages( diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 7a7d4ad6..d0778137 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -16,13 +16,14 @@ ToolContext, ) from app.services.document_agent.structure.page_locate_agent import ( - PageLocateConfig, - PageLocateResidualAgent, + verify_section_page_choice, ) from app.services.document_agent.structure.hierarchy_locator import ( ResolvedHierarchyRange, + TitleMatch, TitleNode, extract_toc_nodes, + iter_leaf_title_nodes, resolve_hierarchy_page_ranges, ) from app.services.document_parser.structure.body_boundary import ( @@ -30,11 +31,62 @@ normalize_heading_text, ) from loguru import logger -from shared.models.schemas.page_memory_config import PageMemoryConfig _FRONT_TOC_REGION_GAP_PAGES = 5 +def _prune_out_of_scope_nodes( + nodes: list[TitleNode], + *, + offset: int, + page_count: int, +) -> tuple[list[TitleNode], int]: + """Remove leaf nodes whose printed_page + offset exceeds page_count. + + Bottom-up: prune out-of-scope leaves, then remove intermediate nodes + that become childless after pruning. Returns (pruned_tree, removed_count). + """ + from dataclasses import replace as _replace + + removed = 0 + + def _prune(node: TitleNode) -> TitleNode | None: + nonlocal removed + if not node.children: + if node.printed_page is not None: + expected = node.printed_page + offset + if expected > page_count or expected < 1: + removed += 1 + return None + return node + pruned_children = [] + for child in node.children: + result = _prune(child) + if result is not None: + pruned_children.append(result) + if not pruned_children: + removed += 1 + return None + return _replace(node, children=pruned_children) + + pruned = [] + for node in nodes: + result = _prune(node) + if result is not None: + pruned.append(result) + + if removed: + logger.info( + "[page_memory.skeleton] pruned {} out-of-scope TOC nodes " + "(printed_page + offset={} exceeds page_count={})", + removed, + offset, + page_count, + ) + + return pruned, removed + + @dataclass(frozen=True) class SectionSkeleton: section_path: str @@ -56,7 +108,6 @@ def extract_section_skeletons( page_texts: dict[int, str], ctx: ToolContext | None = None, hierarchy_nodes: list[TitleNode] | None = None, - page_memory_config: PageMemoryConfig | None = None, ) -> list[SectionSkeleton]: """Convert PageAnatomyMap hierarchy evidence into section skeletons. @@ -69,10 +120,12 @@ def extract_section_skeletons( return [_root_skeleton(root_path=root_path, filename=filename, page_count=0)] toc_selection: dict[str, Any] = {} + pending_tocs: list[dict[str, Any]] = [] + toc_hierarchies: list[dict[str, Any]] | None = None if hierarchy_nodes: nodes = hierarchy_nodes else: - toc_hierarchies, toc_selection = _select_global_toc_hierarchies( + toc_hierarchies, pending_tocs, toc_selection = _select_global_toc_hierarchies( anatomy=anatomy, filename=filename, ) @@ -91,30 +144,86 @@ def extract_section_skeletons( # 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 - # original leaf title so PageLocateResidualAgent can find it in the PDF. + # original leaf title so offset-guided anchoring can find it in the PDF. nodes = _collapse_intermediate_single_child_chains(nodes) body_pages = _body_pages(anatomy=anatomy, page_count=page_count) - offset_hint = _estimate_page_offset(nodes=nodes, anatomy=anatomy) - locate_result = PageLocateResidualAgent( + + # When pending TOCs exist, limit primary scope so the last sibling's + # end_page doesn't extend into the pending TOC region. + primary_page_count = page_count + primary_body_pages = body_pages + if pending_tocs: + pending_starts: list[int] = [] + for t in pending_tocs: + start = _toc_range_start(t) + if start is not None: + pending_starts.append(start) + if pending_starts: + primary_page_count = min(pending_starts) - 1 + primary_body_pages = [p for p in body_pages if p <= primary_page_count] + + offset_hint, calibration_overrides = _calibrate_offset_via_vlm( + nodes=nodes, + toc_hierarchies=toc_hierarchies if not hierarchy_nodes else None, ctx=ctx, page_texts=page_texts, - body_pages=body_pages, page_count=page_count, - page_offset_hint=offset_hint, - config=( - PageLocateConfig.from_page_memory_config(page_memory_config) - if page_memory_config is not None - else None - ), - ).prepare(nodes) + ) + + # Prune TOC nodes whose printed_page + offset exceeds the physical PDF. + pruned_count = 0 + if offset_hint is not None: + nodes, pruned_count = _prune_out_of_scope_nodes( + nodes, offset=offset_hint, page_count=page_count, + ) + if not nodes: + return [ + _root_skeleton( + root_path=root_path, + filename=filename, + page_count=page_count, + reason="all_toc_nodes_out_of_scope", + ) + ] + + # Phase A3: try offset-guided bulk anchoring before expensive residual agent. + offset_matches: dict[tuple[str, ...], TitleMatch] | None = None + if offset_hint is not None and ctx is not None: + offset_matches = _offset_guided_anchoring( + nodes=nodes, + offset=offset_hint, + ctx=ctx, + page_count=page_count, + calibration_overrides=calibration_overrides, + ) + + if offset_matches is not None: + match_overrides = offset_matches + locate_summary: dict[str, Any] = { + "agent": "offset_guided_bulk", + "offset": offset_hint, + "bulk_count": len(offset_matches), + "pruned_out_of_scope": pruned_count, + } + else: + match_overrides = calibration_overrides + locate_summary = { + "agent": "offset_only", + "offset": offset_hint, + "reason": "offset_guided_anchoring_skipped_or_empty", + "pruned_out_of_scope": pruned_count, + } + resolve_nodes = nodes + + ranges = resolve_hierarchy_page_ranges( - locate_result.nodes, - page_count=page_count, + resolve_nodes, + page_count=primary_page_count, page_texts=page_texts, - body_pages=body_pages, + body_pages=primary_body_pages, page_offset_hint=offset_hint, - match_overrides=locate_result.match_overrides, + match_overrides=match_overrides, ) if not ranges: return [ @@ -131,11 +240,25 @@ def extract_section_skeletons( item, filename=filename, page_count=page_count, - locate_summary=locate_result.summary, + locate_summary=locate_summary, toc_selection=toc_selection, ) for item in ranges ] + + # Phase B: graft pending TOCs (appendix / parallel sections) + if pending_tocs: + secondary_skeletons = _resolve_pending_tocs( + pending_tocs=pending_tocs, + primary_ranges=ranges, + ctx=ctx, + page_texts=page_texts, + page_count=page_count, + filename=filename, + body_pages=body_pages, + ) + skeletons.extend(secondary_skeletons) + _log_unlocated_title_warnings(filename=filename, skeletons=skeletons) return skeletons @@ -177,11 +300,11 @@ def _range_to_skeleton( # Motivation: TOC hierarchies often contain "structural" intermediate nodes # (category codes, volume identifiers) that add depth but carry no locatable # text. Compressing them before locate keeps emit_depth small and lets the -# VLM/grep focus on meaningful leaf titles. +# offset-guided anchoring focus on meaningful leaf titles. # # Critical invariant: a node whose only child is a LEAF (no grandchildren) is # NOT merged, so the leaf's original title survives unchanged into -# PageLocateResidualAgent. Only pure-intermediate chains are compressed. +# offset-guided anchoring. Only pure-intermediate chains are compressed. def _collapse_intermediate_single_child_chains( @@ -256,17 +379,19 @@ def _select_global_toc_hierarchies( *, anatomy: Any | None, filename: str, -) -> tuple[list[dict[str, Any]] | None, dict[str, Any]]: - """Keep the front/global TOC cluster and skip later embedded TOCs. +) -> tuple[list[dict[str, Any]] | None, list[dict[str, Any]], dict[str, Any]]: + """Split TOC hierarchies into primary (front cluster) and pending (for probe). + + Profile-time TOC extraction can find multiple TOCs in a long document. + The front cluster is selected by physical page proximity. Remaining TOCs + are returned as *pending* for downstream independent calibration rather + than being unconditionally discarded. - Profile-time TOC extraction can find local TOCs inside a long document - (for example, an embedded standard with its own English outline). Page - memory C4 currently emits a document-level skeleton, so later page-based - TOC regions must not be concatenated as root siblings. + Returns (primary_hierarchies, pending_hierarchies, summary). """ hierarchies = list(_toc_hierarchies(anatomy) or []) if len(hierarchies) <= 1: - return (hierarchies or None), {} + return (hierarchies or None), [], {} page_based = [ hierarchy @@ -274,11 +399,11 @@ def _select_global_toc_hierarchies( if hierarchy.get("toc_range_unit") == "page" and _toc_range_start(hierarchy) is not None ] if not page_based or len(page_based) != len(hierarchies): - return hierarchies, {} + return hierarchies, [], {} sorted_items = sorted(enumerate(hierarchies), key=lambda item: _toc_range_start(item[1]) or 0) selected_indices: set[int] = set() - skipped: list[dict[str, Any]] = [] + pending_indices: list[int] = [] cluster_end: int | None = None for original_index, hierarchy in sorted_items: @@ -295,34 +420,29 @@ def _select_global_toc_hierarchies( selected_indices.add(original_index) cluster_end = max(cluster_end, end) continue - skipped.append( - { - "index": original_index, - "toc_range": [start, end], - "scan_range": hierarchy.get("scan_range"), - "reason": "embedded_toc_region_outside_front_cluster", - } - ) + pending_indices.append(original_index) selected = [ hierarchy for index, hierarchy in enumerate(hierarchies) if index in selected_indices ] - if skipped: - logger.warning( - "[page_memory.skeleton] skipped embedded toc regions filename={} skipped={}", + pending = [hierarchies[i] for i in pending_indices] + + if pending: + logger.info( + "[page_memory.skeleton] toc split: primary={} pending={} filename={}", + len(selected), + len(pending), filename, - skipped, ) summary = { - "strategy": "front_page_toc_cluster", + "strategy": "front_cluster_with_pending", "input_count": len(hierarchies), - "selected_count": len(selected), - "skipped_count": len(skipped), - "skipped": skipped, + "primary_count": len(selected), + "pending_count": len(pending), } - return (selected or None), summary + return (selected or None), pending, summary def _toc_range_start(hierarchy: dict[str, Any]) -> int | None: @@ -364,35 +484,594 @@ 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] -def _estimate_page_offset(*, nodes: list[TitleNode], anatomy: Any | None) -> int | None: - printed_by_title: dict[str, int] = {} - for node in _walk_nodes(nodes): +# ── VLM offset calibration (Phase A1) ─────────────────────────────────────── + +_CALIBRATION_WINDOW_PAGES = 10 +_CALIBRATION_LEAF_PROBE_COUNT = 3 + + +def _calibrate_offset_via_vlm( + *, + nodes: list[TitleNode], + toc_hierarchies: list[dict[str, Any]] | None, + ctx: ToolContext | None, + page_texts: dict[int, str], + page_count: int, +) -> tuple[int | None, dict[tuple[str, ...], TitleMatch]]: + """Scan pages after the TOC to find the first leaf entry via VLM. + + Computes offset = confirmed_physical_page - printed_page. + Returns (offset, match_overrides) where match_overrides contains the + confirmed entry so downstream locate doesn't re-process it. + """ + if ctx is None: + return None, {} + + toc_physical_end = _toc_cluster_end_page(toc_hierarchies) + if toc_physical_end is None: + return None, {} + + scan_start = toc_physical_end + 1 + scan_end = min(scan_start + _CALIBRATION_WINDOW_PAGES - 1, page_count) + if scan_start > page_count: + return None, {} + + leaves = list(iter_leaf_title_nodes(nodes)) + probe_leaves = [ + (path_titles, node) + for path_titles, node in leaves + if node.printed_page is not None + ][:_CALIBRATION_LEAF_PROBE_COUNT] + + if not probe_leaves: + return None, {} + + scan_pages = list(range(scan_start, scan_end + 1)) + candidates = [ + TitleMatch( + page=page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=scan_pages, + evidence={"calibration_probe": True}, + ) + for page in scan_pages + ] + + for path_titles, node in probe_leaves: + result = verify_section_page_choice( + ctx=ctx, + title=node.title, + candidate_matches=candidates, + candidate_page_cap=len(scan_pages), + ) + selected = result.get("selected_page") + if selected is not None and result.get("confidence", 0) >= 0.6: + offset = selected - node.printed_page + match = TitleMatch( + page=selected, + confidence=result.get("confidence", 0.75), + source="agent_vlm", + matched_line="", + score=result.get("confidence", 0.75), + candidates=[selected], + evidence={ + "calibration": True, + "printed_page": node.printed_page, + "reason": result.get("reason", ""), + }, + ) + logger.info( + "[page_memory.skeleton] calibration confirmed: title={!r} " + "printed_page={} physical_page={} offset={}", + node.title, + node.printed_page, + selected, + offset, + ) + return offset, {path_titles: match} + + logger.info("[page_memory.skeleton] calibration: no leaf confirmed in scan window") + return None, {} + + +def _toc_cluster_end_page(toc_hierarchies: list[dict[str, Any]] | None) -> int | None: + """Get the last physical page of the primary TOC cluster.""" + if not toc_hierarchies: + return None + end_pages: list[int] = [] + for hierarchy in toc_hierarchies: + end = _toc_range_end(hierarchy) + if end is not None: + end_pages.append(end) + return max(end_pages) if end_pages else None + + +# ── Offset-guided bulk anchoring with recursive recalibrate (Phase A3) ─────── + +_TAIL_VERIFY_CONFIDENCE_THRESHOLD = 0.6 +_MAX_RECALIBRATE_DEPTH = 5 +_MAX_RECALIBRATE_DELTA = 5 + + +def _verify_offset_tail( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, +) -> bool: + """VLM-verify that the offset holds for the last leaf entry (Theorem 1). + + If head offset == tail offset, monotonicity guarantees all intermediate + entries share the same offset. + + Prefers a tail leaf whose expected page is strictly less than page_count + (boundary pages are unreliable for VLM verification). + """ + tail_leaves = [ + (path, node) for path, node in reversed(leaves) if node.printed_page is not None + ] + if not tail_leaves: + return True + + # Prefer non-boundary: printed_page + offset < page_count + selected = None + for path, node in tail_leaves: + pp = node.printed_page + if pp is None: + continue + expected = pp + offset + if 1 <= expected < page_count: + selected = (path, node) + break + if selected is None: + # All leaves are at the boundary; fall back to the last one + selected = tail_leaves[0] + + path, node = selected + printed_page = node.printed_page + if printed_page is None: + return True + expected_page = printed_page + offset + if expected_page < 1 or expected_page > page_count: + return False + + candidate = TitleMatch( + page=expected_page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[expected_page], + evidence={"tail_verify_probe": True}, + ) + result = verify_section_page_choice( + ctx=ctx, + title=node.title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + confirmed = ( + result.get("selected_page") == expected_page + and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD + ) + logger.info( + "[page_memory.skeleton] tail verify: title={!r} expected_page={} confirmed={} confidence={}", + node.title, + expected_page, + confirmed, + result.get("confidence", 0), + ) + return confirmed + + +def _vlm_confirm_single_page( + *, + ctx: ToolContext, + title: str, + expected_page: int, + page_count: int, +) -> bool: + """Single-page VLM confirmation for binary search steps.""" + if expected_page < 1 or expected_page > page_count: + return False + candidate = TitleMatch( + page=expected_page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[expected_page], + evidence={"bisect_probe": True}, + ) + result = verify_section_page_choice( + ctx=ctx, + title=title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + return ( + result.get("selected_page") == expected_page + and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD + ) + + +def _bisect_offset_breakpoint( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, +) -> int: + """Binary search for the last leaf index where offset is valid. O(log n) VLM calls.""" + lo, hi = 0, len(leaves) - 1 + while lo < hi: + mid = (lo + hi + 1) // 2 + _, node = leaves[mid] if node.printed_page is None: + hi = mid - 1 continue - printed_by_title[_title_key(node.title)] = node.printed_page + expected = node.printed_page + offset + if _vlm_confirm_single_page( + ctx=ctx, title=node.title, expected_page=expected, page_count=page_count + ): + lo = mid + else: + hi = mid - 1 + logger.info( + "[page_memory.skeleton] bisect breakpoint: last_valid_index={} / total={}", + lo, + len(leaves), + ) + return lo - offsets: list[int] = [] - h1_result = getattr(anatomy, "h1_result", None) - for candidate in getattr(h1_result, "h1_candidates", []) or []: - printed_page = printed_by_title.get(_title_key(candidate.title)) - if printed_page is not None: - offsets.append(int(candidate.page) - printed_page) - if not offsets: + +def _bulk_offset_matches( + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, +) -> dict[tuple[str, ...], TitleMatch]: + """Generate TitleMatch overrides for all leaves using offset. No VLM calls.""" + matches: dict[tuple[str, ...], TitleMatch] = {} + for path_titles, node in leaves: + if node.printed_page is None: + continue + page = node.printed_page + offset + matches[path_titles] = TitleMatch( + page=page, + confidence=0.88, + source="agent_vlm", + matched_line="", + score=0.88, + candidates=[page], + evidence={ + "bulk_offset": True, + "offset": offset, + "printed_page": node.printed_page, + }, + ) + return matches + + +def _recalibrate_after_breakpoint( + *, + entry_node: TitleNode, + old_offset: int, + ctx: ToolContext, + page_count: int, +) -> int | None: + """Probe offsets old_offset+1, +2, ... to find new offset after breakpoint. + + Monotonicity guarantees new offset > old offset, so search space is tiny. + """ + entry_printed_page = entry_node.printed_page + if entry_printed_page is None: return None - offsets.sort() - return offsets[len(offsets) // 2] + for delta in range(1, _MAX_RECALIBRATE_DELTA + 1): + new_offset = old_offset + delta + if _vlm_confirm_single_page( + ctx=ctx, + title=entry_node.title, + expected_page=entry_printed_page + new_offset, + page_count=page_count, + ): + logger.info( + "[page_memory.skeleton] recalibrate: title={!r} new_offset={} (delta=+{})", + entry_node.title, + new_offset, + delta, + ) + return new_offset + return None -def _walk_nodes(nodes: list[TitleNode]) -> list[TitleNode]: - walked: list[TitleNode] = [] - for node in nodes: - walked.append(node) - walked.extend(_walk_nodes(node.children)) - return walked +def _offset_guided_anchoring( + *, + nodes: list[TitleNode], + offset: int, + ctx: ToolContext, + page_count: int, + calibration_overrides: dict[tuple[str, ...], TitleMatch], +) -> dict[tuple[str, ...], TitleMatch] | None: + """Offset-guided bulk anchoring with recursive recalibrate on breakpoints. + + Strategy: + 1. Tail verify last leaf with current offset + 2. If pass → bulk apply all leaves (Theorem 1) + 3. If fail → binary search for breakpoint + 4. Bulk apply leaves before breakpoint + 5. Recalibrate: probe remaining[0] with offset+1, +2, ... (monotonicity) + 6. Recurse on remaining segment with new offset + 7. If recalibrate fails → return partial (caller falls back for remainder) + + Returns match_overrides for all anchored leaves, or None for full fallback. + """ + leaves = [ + (path, node) + for path, node in iter_leaf_title_nodes(nodes) + if node.printed_page is not None + ] + if len(leaves) < 2: + return None + + all_matches: dict[tuple[str, ...], TitleMatch] = {} + all_matches.update(calibration_overrides) + + _anchor_segment_recursive( + leaves=leaves, + offset=offset, + ctx=ctx, + page_count=page_count, + matches=all_matches, + depth=0, + ) + + if not all_matches: + return None + + logger.info( + "[page_memory.skeleton] offset bulk anchoring: {} / {} leaves anchored", + len(all_matches), + len(leaves), + ) + return all_matches + + +def _anchor_segment_recursive( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, + matches: dict[tuple[str, ...], TitleMatch], + depth: int, +) -> None: + """Recursively anchor a segment of leaves, handling multiple breakpoints.""" + if not leaves or depth >= _MAX_RECALIBRATE_DEPTH: + return + + if _verify_offset_tail(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count): + bulk = _bulk_offset_matches(leaves, offset) + matches.update(bulk) + return + + bp = _bisect_offset_breakpoint(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count) + confirmed_leaves = leaves[: bp + 1] + if confirmed_leaves: + bulk = _bulk_offset_matches(confirmed_leaves, offset) + matches.update(bulk) + + remaining = leaves[bp + 1:] + if not remaining: + return + + _, first_remaining_node = remaining[0] + new_offset = _recalibrate_after_breakpoint( + entry_node=first_remaining_node, + old_offset=offset, + ctx=ctx, + page_count=page_count, + ) + if new_offset is None: + return + + _anchor_segment_recursive( + leaves=remaining, + offset=new_offset, + ctx=ctx, + page_count=page_count, + matches=matches, + depth=depth + 1, + ) + + +# ── Multi-TOC grafting (Track B) ───────────────────────────────────────────── + + +def _resolve_pending_tocs( + *, + pending_tocs: list[dict[str, Any]], + primary_ranges: list[ResolvedHierarchyRange], + ctx: ToolContext | None, + page_texts: dict[int, str], + page_count: int, + filename: str, + body_pages: list[int], +) -> list[SectionSkeleton]: + """Independently calibrate and anchor each pending TOC, then graft results. + + Each pending TOC gets its own offset via VLM calibration + tail verify, + then entries are bulk-anchored (or fallback to residual agent). + Classification is PARALLEL (append at root level) or CONTAINED (skip). + """ + if not pending_tocs or ctx is None: + return [] + + all_secondary_skeletons: list[SectionSkeleton] = [] + + for i, pending_toc in enumerate(pending_tocs): + toc_range = pending_toc.get("toc_range") + nodes = extract_toc_nodes([pending_toc]) + if not nodes: + continue + nodes = _collapse_intermediate_single_child_chains(nodes) + + # Each TOC's content scope: [toc_range_end + 1, next_toc_start - 1] + toc_end = _toc_range_end(pending_toc) + toc_scope_start = (toc_end + 1) if toc_end is not None else None + next_starts: list[int] = [] + for j in range(i + 1, len(pending_tocs)): + start = _toc_range_start(pending_tocs[j]) + if start is not None: + next_starts.append(start) + toc_scope_end = (min(next_starts) - 1) if next_starts else page_count + toc_body_pages = [ + p for p in body_pages + if p <= toc_scope_end and (toc_scope_start is None or p >= toc_scope_start) + ] + + offset, cal_overrides = _calibrate_offset_via_vlm( + nodes=nodes, + toc_hierarchies=[pending_toc], + ctx=ctx, + page_texts=page_texts, + page_count=toc_scope_end, + ) + + if offset is None: + logger.info( + "[page_memory.skeleton] pending TOC toc_range={}: calibration failed, skipping", + toc_range, + ) + continue + + relationship = _classify_toc_relationship( + offset=offset, + nodes=nodes, + primary_ranges=primary_ranges, + page_count=page_count, + ) + if relationship == "unresolvable": + logger.info( + "[page_memory.skeleton] pending TOC toc_range={}: unresolvable, skipping", + toc_range, + ) + continue + + offset_matches = _offset_guided_anchoring( + nodes=nodes, + offset=offset, + ctx=ctx, + page_count=toc_scope_end, + calibration_overrides=cal_overrides, + ) + + if offset_matches is not None: + match_overrides = offset_matches + locate_summary: dict[str, Any] = { + "agent": "offset_guided_bulk", + "offset": offset, + "bulk_count": len(offset_matches), + "toc_relationship": relationship, + } + else: + match_overrides = cal_overrides + locate_summary = { + "agent": "offset_only", + "offset": offset, + "toc_relationship": relationship, + "reason": "offset_guided_anchoring_skipped_or_empty", + } + + 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, + ) + + toc_selection_info: dict[str, Any] = { + "toc_range": toc_range, + "offset": offset, + "relationship": relationship, + } + for item in ranges: + skeleton = _range_to_skeleton( + item, + filename=filename, + page_count=toc_scope_end, + locate_summary=locate_summary, + toc_selection=toc_selection_info, + ) + all_secondary_skeletons.append(skeleton) + + logger.info( + "[page_memory.skeleton] pending TOC toc_range={}: " + "relationship={} offset={} skeletons={}", + toc_range, + relationship, + offset, + len(ranges), + ) + + return all_secondary_skeletons + + +def _classify_toc_relationship( + *, + offset: int, + nodes: list[TitleNode], + primary_ranges: list[ResolvedHierarchyRange], + page_count: int, +) -> str: + """Classify a pending TOC as parallel or contained vs primary ranges. + + parallel: the pending TOC covers pages beyond the primary tree's *anchored* + content (i.e. the last explicitly-located section start page). + contained: the pending TOC's content falls strictly within a primary + section's explicitly-anchored range. + """ + leaves = [ + node for _, node in iter_leaf_title_nodes(nodes) if node.printed_page is not None + ] + if not leaves: + return "unresolvable" + + first_printed = leaves[0].printed_page + last_printed = leaves[-1].printed_page + if first_printed is None or last_printed is None: + return "unresolvable" + first_physical = first_printed + offset + last_physical = last_printed + offset + + if first_physical < 1 or first_physical > page_count: + return "unresolvable" + + if not primary_ranges: + return "parallel" + + # Use the last *start_page* among primary ranges as the boundary of + # explicitly-anchored content. The end_page of the last section is often + # extended to page_count by default and doesn't reflect real content coverage. + last_anchored_start = max( + (r.start_page for r in primary_ranges if r.start_page is not None), default=0 + ) + + if first_physical > last_anchored_start: + return "parallel" + min_level = min(r.level for r in primary_ranges) + top_level_ranges = [r for r in primary_ranges if r.level == min_level] + for r in top_level_ranges: + if r.start_page and r.end_page: + if r.start_page <= first_physical and last_physical <= r.end_page: + return "contained" -def _title_key(title: str) -> str: - return normalize_heading_text(clean_toc_title(title) or title).casefold() + return "parallel" def _clamp_page(page: int, page_count: int) -> int: diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py index 3b151378..8858056f 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py @@ -6,6 +6,9 @@ from shared.services.retrieval.agentic.core.budget import budget_status_from_snapshot from shared.services.retrieval.agentic.navigation.path_ledger import PathLedger +from shared.services.retrieval.agentic.navigation.state import ( + RejectionRecord, +) from shared.services.retrieval.search.lexical_text import normalize_section_path from shared.utils.text_utils import truncate_content_preview @@ -68,8 +71,7 @@ def build_legal_actions( collected_paths: list[dict[str, Any]], expanded_scopes: set[str], discovery_hints: list[dict[str, Any]] | None = None, - rejected_paths: set[str] | None = None, - rejected_collect_paths: set[str] | None = None, + rejected: dict[str, RejectionRecord] | None = None, total_images: int, total_tables: int, disabled_asset_types: set[str] | None = None, @@ -79,11 +81,14 @@ def build_legal_actions( covered_paths = _covered_paths(collected_paths) outline_paths = _outline_paths(collected_paths) budget_mode = budget_status_from_snapshot(budget_snapshot) - rejected = {PathLedger.normalize(path) for path in rejected_paths or set()} - rejected_collects = { - normalized - for path in rejected_collect_paths or set() - if (normalized := PathLedger.normalize(path)) + rejection_ledger = rejected or {} + tool_adjudicated_paths = { + path for path, record in rejection_ledger.items() + if record.reason == "tool_adjudicated" + } + navigational_abandoned_paths = { + path for path, record in rejection_ledger.items() + if record.reason == "navigational_abandon" } discovery_scores = _discovery_scores_by_path(discovery_hints or []) scored_items = _score_items(items, discovery_scores) @@ -100,9 +105,12 @@ def build_legal_actions( path = str(item.get("path") or "").strip() if not path or path == "Root": continue - if PathLedger.is_covered(path, covered_paths): + normalized_path = PathLedger.normalize(path) + if PathLedger.is_covered(normalized_path, covered_paths): continue - if PathLedger.is_covered(path, rejected_collects): + # tool_adjudicated rejections are content-level negative and are not + # revived this round (TODO: future strong-signal revival). + if PathLedger.is_covered(normalized_path, tool_adjudicated_paths): continue action_set.add(LegalAction( @@ -112,7 +120,7 @@ def build_legal_actions( target_scope=path, note=( "upgrade outline to full evidence" - if path in outline_paths + if normalized_path in outline_paths else _item_note(item) ), score=float(item.get("relevance_score") or 0.0), @@ -128,18 +136,22 @@ def build_legal_actions( critical_expand = True if item.get("is_leaf"): continue - if path == current_scope: + if normalized_path == current_scope: continue - if current_scope and PathLedger.is_ancestor(path, current_scope): + if current_scope and PathLedger.is_ancestor(normalized_path, current_scope): continue - if path in expanded_scopes: + if normalized_path in expanded_scopes: continue if ( budget_mode == "TIGHT" - and path not in expand_allowlist + and normalized_path not in expand_allowlist ): continue - if path in rejected and not _path_has_discovery_signal(path, discovery_scores): + # EXPAND suppression: navigational_abandon (weak) suppresses unless a + # discovery / lexical signal revives the path. + if normalized_path in navigational_abandoned_paths and not _has_discovery_signal( + normalized_path, discovery_scores + ): continue action_set.add(LegalAction( @@ -166,9 +178,10 @@ def build_legal_actions( seen_discovery_paths.add(path) if PathLedger.is_covered(path, covered_paths): continue + # tool_adjudicated rejections are not revived by discovery this round. # TODO: allow tool-specific LLM adjudicators to revive rejected # collects when validity cannot be determined structurally. - if PathLedger.is_covered(path, rejected_collects): + if PathLedger.is_covered(path, tool_adjudicated_paths): continue if any(action.path == path for action in action_set.collect): continue @@ -236,9 +249,8 @@ def format_agent_state_block( current_scope: str | None, query_intent: str, expanded_scopes: set[str], - rejected_paths: set[str], + rejected: dict[str, RejectionRecord], collected_paths: list[dict[str, Any]], - rejected_collect_paths: set[str] | None = None, prior_tool_result: dict[str, Any] | None, search_context: str, budget_snapshot: dict[str, Any] | None, @@ -267,15 +279,26 @@ def format_agent_state_block( lines.append(f' - "{path}"') else: lines.append("Expanded scopes: none") - rejected_collects = set(rejected_collect_paths or set()) - low_value_rejected = set(rejected_paths) - rejected_collects - if low_value_rejected: - lines.append("Low-value scopes avoided unless revived by discovery:") - for path in sorted(low_value_rejected): + + navigational_abandoned = sorted( + path for path, record in rejected.items() + if record.reason == "navigational_abandon" + ) + tool_adjudicated = sorted( + path for path, record in rejected.items() + if record.reason == "tool_adjudicated" + ) + if navigational_abandoned: + lines.append( + "Scopes avoided (soft; revived by discovery):" + ) + for path in navigational_abandoned: lines.append(f' - "{path}"') - if rejected_collect_paths: - lines.append("Collects rejected by tool reconciliation:") - for path in sorted(rejected_collect_paths): + if tool_adjudicated: + lines.append( + "Collects rejected by tool reconciliation (content-level; not revived):" + ) + for path in tool_adjudicated: lines.append(f' - "{path}"') full_paths, outline_paths = _dedupe_collection_modes(collected_paths) @@ -648,10 +671,11 @@ def _expand_allowlist( return set(candidates[:limit]) -def _path_has_discovery_signal( +def _has_discovery_signal( path: str, discovery_scores: dict[str, float], ) -> bool: + """A path has a discovery signal if it, an ancestor, or a descendant appears in discovery.""" return any( candidate == path or PathLedger.is_ancestor(path, candidate) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py index bc6c38fa..22d9c273 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py @@ -169,7 +169,6 @@ async def _navigate_collector( Returns (doc_pending_assets, collected_paths). """ - doc_exclude: set[str] = set() doc_discovery_hints = self._discovery_by_doc.get(doc.document_id, []) doc_pending_assets: list[dict[str, Any]] = [] nav_state = NavigationState( @@ -201,8 +200,7 @@ async def _navigate_collector( nav_state.step_count += 1 before_scope = nav_state.current_scope expanded_before = set(nav_state.expanded_scopes) - rejected_before = set(nav_state.rejected_paths) - rejected_collect_before = set(nav_state.rejected_collect_paths) + rejected_before = dict(nav_state.rejected) collected_before_count = len(nav_state.collected_paths) doc_llm_fn = self._llm_budget.for_document( @@ -226,13 +224,11 @@ async def _navigate_collector( namespace=self._namespace, doc_name=doc_name, scope_path=nav_state.current_scope, - exclude_paths=doc_exclude, budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None, nav_trace=nav_state.nav_trace if nav_state.nav_trace else None, collected_paths=nav_state.collected_paths, expanded_scopes=nav_state.expanded_scopes, - rejected_paths=nav_state.rejected_paths, - rejected_collect_paths=nav_state.rejected_collect_paths, + rejected=nav_state.rejected, disabled_asset_types=self._disabled_asset_types | nav_state.blocked_asset_types_for_scope( nav_state.current_scope ), @@ -292,8 +288,11 @@ async def _navigate_collector( if rejected_collects: nav_result.collect = collect_reconcile["accepted_collects"] for path in rejected_collects: - nav_state.mark_rejected_collect(path) - doc_exclude.add(path) + nav_state.mark_rejected_collect( + path, + step=nav_state.step_count, + detail=collect_reconcile.get("reason", ""), + ) logger.info( " agentic: tool reconciliation rejected collects: " f"{rejected_collects}" @@ -309,10 +308,10 @@ async def _navigate_collector( scope_context=nav_state.current_scope, ) collected_in_step.append(path) - # Outline collections should NOT exclude children — the intent - # is "see structure, then drill deeper for full content". - if coll_item.get("hydrate_mode") != "outline": - doc_exclude.add(path) + # Outline collections keep children visible — the intent is + # "see structure, then drill deeper for full content". + # Coverage / action filtering now derives from collected_paths + # via the navigation state ledger (no physical exclusion). # ── Process navigation action ──────────────────────────────── should_break = False @@ -336,7 +335,11 @@ async def _navigate_collector( else: back_target = nav_result.back_to # None = root if PathLedger.valid_back_target(nav_state.current_scope, back_target): - nav_state.mark_rejected_if_unproductive(nav_state.current_scope) + nav_state.mark_rejected_if_unproductive( + nav_state.current_scope, + step=nav_state.step_count, + detail="back_from_unproductive_scope", + ) nav_state.current_scope = back_target else: logger.warning( @@ -363,7 +366,6 @@ async def _navigate_collector( before_scope=before_scope, expanded_before=expanded_before, rejected_before=rejected_before, - rejected_collect_before=rejected_collect_before, collected_before_count=collected_before_count, ) trace_entry: dict[str, Any] = { diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py index 841059e5..83b6ac57 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py @@ -2,45 +2,111 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal from shared.services.retrieval.agentic.navigation.path_ledger import PathLedger +RejectReason = Literal["tool_adjudicated", "navigational_abandon"] + +# Strength ordering: a stronger reason overrides a weaker one so we keep the +# most informative record for a given path. +_REASON_STRENGTH: dict[RejectReason, int] = { + "tool_adjudicated": 2, + "navigational_abandon": 1, +} + + +@dataclass +class RejectionRecord: + """A single rejection entry in the navigation state ledger. + + Two reasons are distinguished: + + - ``tool_adjudicated``: a SEARCH_* tool reconciliation proved the path has + no matching assets. Strong (content-level) negative signal. Not revived + this round; future strong-signal revival is a TODO. + - ``navigational_abandon``: BACK left an unproductive scope. Weak (soft) + signal; any discovery / lexical hit revives the path. + """ + + path: str + reason: RejectReason + step: int + detail: str = "" + @dataclass class NavigationState: - """Mutable state for one document navigation loop.""" + """Mutable state for one document navigation loop. + + Path state is tracked in two orthogonal dimensions: + + - **Coverage** (positive "already taken as evidence"): derived from + ``collected_paths`` via :meth:`covered_paths` / :meth:`outline_paths`. + - **Rejection** (negative "evaluated, not taken"): a single labelled + ledger in :attr:`rejected` keyed by normalized path. Replaces the + former ``rejected_paths`` / ``rejected_collect_paths`` dual sets. + """ document_id: str document_name: str job_result_id: str current_scope: str | None = None expanded_scopes: set[str] = field(default_factory=set) - rejected_paths: set[str] = field(default_factory=set) - rejected_collect_paths: set[str] = field(default_factory=set) + rejected: dict[str, RejectionRecord] = field(default_factory=dict) collected_paths: list[dict[str, Any]] = field(default_factory=list) nav_trace: list[dict[str, Any]] = field(default_factory=list) tool_history: list[dict[str, Any]] = field(default_factory=list) blocked_asset_searches: set[str] = field(default_factory=set) step_count: int = 0 + # ── Coverage helpers (single source: collected_paths) ──────────────── + + def covered_paths(self) -> set[str]: + """Full-evidence paths (hydrate_mode != 'outline').""" + return { + PathLedger.normalize(str(item.get("path") or "")) + for item in self.collected_paths + if item.get("path") and item.get("hydrate_mode") != "outline" + } + + def outline_paths(self) -> set[str]: + """Outline-only paths; excludes any path also collected as full.""" + full = self.covered_paths() + return { + PathLedger.normalize(str(item.get("path") or "")) + for item in self.collected_paths + if item.get("path") + and item.get("hydrate_mode") == "outline" + and PathLedger.normalize(str(item.get("path") or "")) not in full + } + + # ── Snapshot / delta for replayable traces ────────────────────────── + def snapshot_delta( self, *, before_scope: str | None, expanded_before: set[str], - rejected_before: set[str], - rejected_collect_before: set[str], + rejected_before: dict[str, RejectionRecord], collected_before_count: int, ) -> dict[str, Any]: + rejected_added: list[dict[str, Any]] = [] + for path, record in self.rejected.items(): + if path in rejected_before: + continue + rejected_added.append({ + "path": record.path, + "reason": record.reason, + "step": record.step, + "detail": record.detail, + }) + rejected_added.sort(key=lambda item: (item["step"], item["path"])) return { "current_scope_before": before_scope or "root", "current_scope_after": self.current_scope or "root", "expanded_added": sorted(self.expanded_scopes - expanded_before), - "rejected_added": sorted(self.rejected_paths - rejected_before), - "rejected_collect_added": sorted( - self.rejected_collect_paths - rejected_collect_before - ), + "rejected_added": rejected_added, "collected_added": [ item.get("path", "") for item in self.collected_paths[collected_before_count:] @@ -48,6 +114,8 @@ def snapshot_delta( ], } + # ── Mutation helpers ──────────────────────────────────────────────── + def add_collected( self, item: dict[str, Any], @@ -66,13 +134,35 @@ def mark_expanded(self, path: str | None) -> None: if normalized: self.expanded_scopes.add(normalized) - def mark_rejected_collect(self, path: str | None) -> None: + def mark_rejected_collect( + self, + path: str | None, + *, + step: int, + detail: str = "", + ) -> None: + """Record a tool-adjudicated rejection (strong, content-level).""" normalized = PathLedger.normalize(path) - if normalized: - self.rejected_paths.add(normalized) - self.rejected_collect_paths.add(normalized) + if not normalized: + return + self._upsert_rejection( + normalized, + reason="tool_adjudicated", + step=step, + detail=detail, + ) - def mark_rejected_if_unproductive(self, path: str | None) -> None: + def mark_rejected_if_unproductive( + self, + path: str | None, + *, + step: int, + detail: str = "", + ) -> None: + """Record a soft navigational abandon when leaving an unproductive scope. + + Only written when no stronger record exists for the path. + """ normalized = PathLedger.normalize(path) if not normalized: return @@ -81,8 +171,40 @@ def mark_rejected_if_unproductive(self, path: str | None) -> None: and PathLedger.is_same_or_descendant(item.get("path"), normalized) for item in self.collected_paths ) - if not has_full_collect: - self.rejected_paths.add(normalized) + if has_full_collect: + return + self._upsert_rejection( + normalized, + reason="navigational_abandon", + step=step, + detail=detail, + ) + + def _upsert_rejection( + self, + normalized_path: str, + *, + reason: RejectReason, + step: int, + detail: str, + ) -> None: + existing = self.rejected.get(normalized_path) + if existing is not None and _REASON_STRENGTH[existing.reason] >= _REASON_STRENGTH[reason]: + # Keep the stronger prior record. + return + self.rejected[normalized_path] = RejectionRecord( + path=normalized_path, + reason=reason, + step=step, + detail=detail, + ) + + def rejected_paths_with_reason(self, reason: RejectReason) -> set[str]: + """All paths rejected with a specific reason label.""" + return { + path for path, record in self.rejected.items() + if record.reason == reason + } def blocked_asset_types_for_scope(self, scope: str | None) -> set[str]: prefix = f"{PathLedger.normalize(scope) or 'root'}:" diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py index 2873a2e2..066fa130 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py @@ -27,6 +27,7 @@ format_agent_state_block, ) from shared.services.retrieval.agentic.navigation.section_tree import load_child_sections +from shared.services.retrieval.agentic.navigation.state import RejectionRecord from shared.services.retrieval.agentic.core.types import DocTreeNode, NavigateStepResult from shared.services.retrieval.llm_adapter import LLMFn @@ -48,8 +49,7 @@ async def navigate_step( nav_trace: list[dict[str, Any]] | None = None, collected_paths: list[dict[str, Any]] | None = None, expanded_scopes: set[str] | None = None, - rejected_paths: set[str] | None = None, - rejected_collect_paths: set[str] | None = None, + rejected: dict[str, RejectionRecord] | None = None, disabled_asset_types: set[str] | None = None, discovery_hints: list[dict[str, Any]] | None = None, section_rows: list | None = None, @@ -89,14 +89,14 @@ async def navigate_step( expanded_path_set = set(expanded_scopes or _expanded_paths_from_trace(nav_trace or [])) if scope_path: expanded_path_set.add(scope_path) + rejection_ledger = rejected or {} provisional_action_set = build_legal_actions( items=items, current_scope=scope_path, collected_paths=collected_paths or [], expanded_scopes=expanded_path_set, discovery_hints=discovery_hints, - rejected_paths=rejected_paths or set(), - rejected_collect_paths=rejected_collect_paths or set(), + rejected=rejection_ledger, total_images=total_images, total_tables=total_tables, disabled_asset_types=disabled_asset_types or set(), @@ -137,8 +137,7 @@ async def navigate_step( collected_paths=collected_paths or [], expanded_scopes=expanded_path_set, discovery_hints=discovery_hints, - rejected_paths=rejected_paths or set(), - rejected_collect_paths=rejected_collect_paths or set(), + rejected=rejection_ledger, total_images=total_images, total_tables=total_tables, disabled_asset_types=disabled_asset_types or set(), @@ -166,16 +165,17 @@ async def navigate_step( "search": [item.id for item in action_set.search], "finish": [action_set.finish.id] if action_set.finish else [], }, - "rejected_paths": sorted(rejected_paths or set()), - "rejected_collect_paths": sorted(rejected_collect_paths or set()), + "rejected": { + path: {"reason": record.reason, "step": record.step} + for path, record in rejection_ledger.items() + }, } agent_state_block = format_agent_state_block( current_scope=scope_path, query_intent=query_intent, expanded_scopes=expanded_path_set, - rejected_paths=rejected_paths or set(), + rejected=rejection_ledger, collected_paths=collected_paths or [], - rejected_collect_paths=rejected_collect_paths or set(), prior_tool_result=prior_tool_result, search_context=search_context, budget_snapshot=adjusted_snapshot, diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index d9d6b331..65e16187 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -90,8 +90,7 @@ async def navigate_step( nav_trace: list[dict[str, Any]] | None = None, collected_paths: list[dict[str, Any]] | None = None, expanded_scopes: set[str] | None = None, - rejected_paths: set[str] | None = None, - rejected_collect_paths: set[str] | None = None, + rejected: dict[str, Any] | None = None, disabled_asset_types: set[str] | None = None, discovery_hints: list[dict[str, Any]] | None = None, section_rows: list | None = None, @@ -114,8 +113,7 @@ async def navigate_step( nav_trace=nav_trace, collected_paths=collected_paths, expanded_scopes=expanded_scopes, - rejected_paths=rejected_paths, - rejected_collect_paths=rejected_collect_paths, + rejected=rejected, disabled_asset_types=disabled_asset_types, discovery_hints=discovery_hints, section_rows=section_rows,