diff --git a/apps/worker/app/services/document_agent/agents/calibration/loop.py b/apps/worker/app/services/document_agent/agents/calibration/loop.py index 852c9991..0ebb2603 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/loop.py +++ b/apps/worker/app/services/document_agent/agents/calibration/loop.py @@ -25,7 +25,7 @@ from app.services.document_agent.budget import BudgetTracker, StageEnvelope from app.services.document_agent.manifest import ToolContext, ToolResult from app.services.document_agent.state import AgentBlackboard -from app.services.document_agent.structure.structure_anchoring import ( +from app.services.document_agent.structure.anchoring_primitives import ( deserialize_skeleton_anchor, serialize_skeleton_anchor, ) diff --git a/apps/worker/app/services/document_agent/agents/calibration/orchestrator.py b/apps/worker/app/services/document_agent/agents/calibration/orchestrator.py new file mode 100644 index 00000000..88736bf5 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/orchestrator.py @@ -0,0 +1,57 @@ +"""Calibration orchestration across Phase-1 and structure Phase-2.""" + +from __future__ import annotations + +from typing import Any + +from app.services.document_agent.agents.calibration.procedure import ( + finalize_calibration_result, + flat_toc_entries, +) +from app.services.document_agent.agents.calibration import service +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.structure.hierarchy_locator import TitleNode +from app.services.document_agent.structure.anchoring_primitives import ( + SkeletonAnchor, + anchor_hierarchy_from_offset, +) + + +def anchor_hierarchy( + *, + nodes: list[TitleNode], + toc_hierarchies: list[dict[str, Any]] | None, + page_texts: dict[int, str], + body_pages: list[int], + page_count: int, + ctx: ToolContext | None, +) -> tuple[list[TitleNode], SkeletonAnchor]: + """Run calibration Phase-1 and the production Phase-2 completion.""" + phase1 = service.calibrate_offset( + nodes=nodes, + toc_hierarchies=toc_hierarchies, + ctx=ctx, + page_texts=page_texts, + page_count=page_count, + ) + if phase1.status == "failed" and not phase1.regimes and phase1.offset is None: + return anchor_hierarchy_from_offset( + nodes=nodes, + offset_hint=None, + calibration_overrides={}, + page_texts=page_texts, + body_pages=body_pages, + page_count=page_count, + ctx=ctx, + ) + working, anchor, _finalized = finalize_calibration_result( + result=phase1, + entries=flat_toc_entries(toc_hierarchies), + toc_hierarchies=list(toc_hierarchies or []), + ctx=ctx, + page_count=page_count, + page_texts=page_texts, + body_pages=body_pages, + nodes=nodes, + ) + return working, anchor diff --git a/apps/worker/app/services/document_agent/agents/calibration/procedure.py b/apps/worker/app/services/document_agent/agents/calibration/procedure.py index 40642bf7..a3ed401d 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/procedure.py +++ b/apps/worker/app/services/document_agent/agents/calibration/procedure.py @@ -31,18 +31,44 @@ normalize_page_kind, parse_printed_page, ) -from app.services.document_agent.structure.structure_anchoring import ( +from app.services.document_agent.structure.anchoring_primitives import ( SkeletonAnchor, locate_null_page_parent_overrides, - offset_guided_anchoring, prune_unanchored_toc_leaves, serialize_skeleton_anchor, ) +from app.services.document_agent.structure import anchoring_primitives as _anchoring +from app.services.document_agent.structure.page_locate_agent import ( + verify_section_page_choice, +) # Re-export under prior names so existing imports keep working. normalize_kind = normalize_page_kind +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: + """Forward phase-2 anchoring while preserving the historical patch seam.""" + original = _anchoring.verify_section_page_choice + _anchoring.verify_section_page_choice = verify_section_page_choice + try: + return _anchoring.offset_guided_anchoring( + nodes=nodes, + offset=offset, + ctx=ctx, + page_count=page_count, + calibration_overrides=calibration_overrides, + ) + finally: + _anchoring.verify_section_page_choice = original + + def pick_primary_offset(result: CalibrationResult) -> int | None: """Prefer decimal-regime candidate offset; else first regime with an offset.""" for regime in result.regimes: @@ -298,7 +324,7 @@ def anchor_hierarchy_from_regimes( if ctx is None: # Offline: still apply deterministic printed+offset for this regime. - from app.services.document_agent.structure.structure_anchoring import ( + from app.services.document_agent.structure.anchoring_primitives import ( bulk_offset_matches, ) diff --git a/apps/worker/app/services/document_agent/structure/anchoring_primitives.py b/apps/worker/app/services/document_agent/structure/anchoring_primitives.py new file mode 100644 index 00000000..390e8ba3 --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/anchoring_primitives.py @@ -0,0 +1,795 @@ +"""Shared hierarchy anchoring: Phase-2 bulk/bisect/null-page + SkeletonAnchor. +Phase-1 offset discovery lives in ``document_agent.agents.calibration``. +``anchor_hierarchy`` composes Phase-1 + Phase-2 for production callers. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any + +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.structure.hierarchy_locator import ( + TitleMatch, + TitleNode, + first_leaf_start_under, + iter_leaf_title_nodes, + last_leaf_start_under, + locate_title_compact_strict, +) +from app.services.document_agent.structure.page_locate_agent import ( + verify_section_page_choice, +) +from loguru import logger + + +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). + """ + 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( + "[structure_anchoring] pruned {} out-of-scope TOC nodes " + "(printed_page + offset={} exceeds page_count={})", + removed, + offset, + page_count, + ) + + return pruned, removed + + +def prune_unanchored_toc_leaves( + nodes: list[TitleNode], + *, + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> tuple[list[TitleNode], int]: + """Remove TOC leaves that have no physical ``match_overrides`` entry. + + Implements Phase-2 ``suffix = no TOC``: after bulk/bisect/recalibrate, any + leaf that was not successfully anchored is dropped from the coarse tree + instead of sticky ``inherited_unlocated`` ranges. Childless parents are + removed unless they themselves have an override. + """ + removed = 0 + + def _prune( + node: TitleNode, parent_titles: tuple[str, ...] + ) -> TitleNode | None: + nonlocal removed + path = (*parent_titles, node.title) + if node.children: + children: list[TitleNode] = [] + for child in node.children: + kept = _prune(child, path) + if kept is not None: + children.append(kept) + if children: + return replace(node, children=children) + if path in match_overrides: + return replace(node, children=[]) + removed += 1 + return None + if path in match_overrides: + return node + removed += 1 + return None + + out: list[TitleNode] = [] + for node in nodes: + kept = _prune(node, ()) + if kept is not None: + out.append(kept) + + if removed: + logger.info( + "[structure_anchoring] pruned {} unanchored TOC nodes " + "(suffix / no match_overrides → no TOC)", + removed, + ) + return out, removed + + +def toc_range_start(hierarchy: dict[str, Any]) -> int | None: + toc_range = hierarchy.get("toc_range") + if not isinstance(toc_range, (list, tuple)) or not toc_range: + return None + try: + return int(toc_range[0]) + except (TypeError, ValueError): + return None + + +def toc_range_end(hierarchy: dict[str, Any]) -> int | None: + toc_range = hierarchy.get("toc_range") + if not isinstance(toc_range, (list, tuple)) or not toc_range: + return None + try: + return int(toc_range[-1]) + except (TypeError, ValueError): + return None + + +# ── 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( + "[structure_anchoring] 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( + "[structure_anchoring] 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( + "[structure_anchoring] 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( + "[structure_anchoring] 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 + + +# ── Offset-guided bulk anchoring with recursive recalibrate (Phase-2) ─────── + +_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( + "[structure_anchoring] 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 + 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( + "[structure_anchoring] bisect breakpoint: last_valid_index={} / total={}", + lo, + len(leaves), + ) + return lo + + +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 + 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( + "[structure_anchoring] recalibrate: title={!r} new_offset={} (delta=+{})", + entry_node.title, + new_offset, + delta, + ) + return new_offset + return None + + +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 not leaves: + return dict(calibration_overrides) or None + + all_matches: dict[tuple[str, ...], TitleMatch] = {} + all_matches.update(calibration_overrides) + + # Single-leaf regimes (roman front-matter, F-1 appendix, …) still get a + # deterministic printed→physical override; Phase-1 already calibrated them. + if len(leaves) == 1: + all_matches.update(bulk_offset_matches(leaves, offset)) + else: + _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( + "[structure_anchoring] 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, + ) + + +@dataclass +class SkeletonAnchor: + offset: int | None + offset_status: str + match_overrides: dict[tuple[str, ...], TitleMatch] + null_page_report: list[dict[str, Any]] + bulk_count: int + pruned_count: int = 0 + locate_agent: str = "offset_only" + + +def serialize_title_match(match: TitleMatch) -> dict[str, Any]: + return { + "page": match.page, + "confidence": match.confidence, + "source": match.source, + "matched_line": match.matched_line, + "score": match.score, + "candidates": list(match.candidates), + "evidence": dict(match.evidence or {}), + } + + +def serialize_skeleton_anchor(anchor: SkeletonAnchor) -> dict[str, Any]: + """JSON-friendly SkeletonAnchor (path tuples joined by ' / ').""" + overrides: dict[str, Any] = {} + for path, match in (anchor.match_overrides or {}).items(): + key = " / ".join(str(part) for part in path) + overrides[key] = serialize_title_match(match) + return { + "offset": anchor.offset, + "offset_status": anchor.offset_status, + "match_overrides": overrides, + "null_page_report": list(anchor.null_page_report or []), + "bulk_count": int(anchor.bulk_count or 0), + "pruned_count": int(anchor.pruned_count or 0), + "locate_agent": anchor.locate_agent, + } + + +def deserialize_title_match(data: dict[str, Any]) -> TitleMatch: + return TitleMatch( + page=int(data["page"]), + confidence=float(data.get("confidence") or 0.0), + source=data.get("source") or "agent_vlm", # type: ignore[arg-type] + matched_line=str(data.get("matched_line") or ""), + score=float(data.get("score") or 0.0), + candidates=[int(p) for p in (data.get("candidates") or [])], + evidence=dict(data.get("evidence") or {}), + ) + + +def deserialize_skeleton_anchor(data: dict[str, Any]) -> SkeletonAnchor: + raw_overrides = data.get("match_overrides") or {} + overrides: dict[tuple[str, ...], TitleMatch] = {} + if isinstance(raw_overrides, dict): + for key, value in raw_overrides.items(): + if not isinstance(value, dict): + continue + if isinstance(key, str): + path = tuple(part.strip() for part in key.split(" / ") if part.strip()) + elif isinstance(key, (list, tuple)): + path = tuple(str(part) for part in key) + else: + continue + if path: + overrides[path] = deserialize_title_match(value) + return SkeletonAnchor( + offset=data.get("offset") if data.get("offset") is None else int(data["offset"]), + offset_status=str(data.get("offset_status") or "failed"), + match_overrides=overrides, + null_page_report=list(data.get("null_page_report") or []), + bulk_count=int(data.get("bulk_count") or 0), + pruned_count=int(data.get("pruned_count") or 0), + locate_agent=str(data.get("locate_agent") or "offset_only"), + ) + + +def anchor_hierarchy_from_offset( + *, + nodes: list[TitleNode], + offset_hint: int | None, + calibration_overrides: dict[tuple[str, ...], TitleMatch] | None = None, + page_texts: dict[int, str], + body_pages: list[int], + page_count: int, + ctx: ToolContext | None, +) -> tuple[list[TitleNode], SkeletonAnchor]: + """Production prune → bulk → null-page given a precomputed offset. + + Phase-2 entry after Agent ``calibrate_offset`` (Phase-1). + """ + seed_overrides = dict(calibration_overrides or {}) + pruned_count = 0 + working = nodes + if offset_hint is not None: + working, pruned_count = prune_out_of_scope_nodes( + working, offset=offset_hint, page_count=page_count + ) + + offset_matches: dict[tuple[str, ...], TitleMatch] | None = None + if offset_hint is not None and ctx is not None and working: + offset_matches = offset_guided_anchoring( + nodes=working, + offset=offset_hint, + ctx=ctx, + page_count=page_count, + calibration_overrides=seed_overrides, + ) + + if offset_matches is not None: + match_overrides = offset_matches + locate_agent = "offset_guided_bulk" + bulk_count = len(offset_matches) + else: + match_overrides = seed_overrides + locate_agent = "offset_only" + bulk_count = 0 + + working, unanchored_removed = prune_unanchored_toc_leaves( + working, match_overrides=match_overrides + ) + pruned_count += unanchored_removed + + match_overrides, null_page_report = locate_null_page_parent_overrides( + nodes=working, + match_overrides=match_overrides, + page_texts=page_texts, + body_pages=body_pages, + ctx=ctx, + ) + + if offset_hint is None: + offset_status = "failed" if ctx is not None else "skipped" + else: + offset_status = "ok" + + return working, SkeletonAnchor( + offset=offset_hint, + offset_status=offset_status, + match_overrides=match_overrides, + null_page_report=null_page_report, + bulk_count=bulk_count, + pruned_count=pruned_count, + locate_agent=locate_agent, + ) diff --git a/apps/worker/app/services/document_agent/structure/structure_anchoring.py b/apps/worker/app/services/document_agent/structure/structure_anchoring.py index 24fb2d22..b32a6849 100644 --- a/apps/worker/app/services/document_agent/structure/structure_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/structure_anchoring.py @@ -1,546 +1,55 @@ -"""Shared hierarchy anchoring: Phase-2 bulk/bisect/null-page + SkeletonAnchor. +"""Compatibility exports for hierarchy anchoring. -Phase-1 offset discovery lives in ``document_agent.agents.calibration``. -``anchor_hierarchy`` composes Phase-1 + Phase-2 for production callers. +Low-level anchoring primitives live in :mod:`anchoring_primitives`; the +calibration-owned orchestrator is resolved lazily to keep imports acyclic. """ from __future__ import annotations -from dataclasses import dataclass, replace +from importlib import import_module from typing import Any -from app.services.document_agent.manifest import ToolContext -from app.services.document_agent.structure.hierarchy_locator import ( +from app.services.document_agent.structure import anchoring_primitives as _anchoring +from app.services.document_agent.structure.anchoring_primitives import ( + SkeletonAnchor, TitleMatch, TitleNode, - first_leaf_start_under, - iter_leaf_title_nodes, - last_leaf_start_under, - locate_title_compact_strict, ) +from app.services.document_agent.manifest import ToolContext from app.services.document_agent.structure.page_locate_agent import ( verify_section_page_choice, ) -from loguru import logger - - -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). - """ - 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( - "[structure_anchoring] pruned {} out-of-scope TOC nodes " - "(printed_page + offset={} exceeds page_count={})", - removed, - offset, - page_count, - ) - - return pruned, removed - - -def prune_unanchored_toc_leaves( - nodes: list[TitleNode], - *, - match_overrides: dict[tuple[str, ...], TitleMatch], -) -> tuple[list[TitleNode], int]: - """Remove TOC leaves that have no physical ``match_overrides`` entry. - - Implements Phase-2 ``suffix = no TOC``: after bulk/bisect/recalibrate, any - leaf that was not successfully anchored is dropped from the coarse tree - instead of sticky ``inherited_unlocated`` ranges. Childless parents are - removed unless they themselves have an override. - """ - removed = 0 - - def _prune( - node: TitleNode, parent_titles: tuple[str, ...] - ) -> TitleNode | None: - nonlocal removed - path = (*parent_titles, node.title) - if node.children: - children: list[TitleNode] = [] - for child in node.children: - kept = _prune(child, path) - if kept is not None: - children.append(kept) - if children: - return replace(node, children=children) - if path in match_overrides: - return replace(node, children=[]) - removed += 1 - return None - if path in match_overrides: - return node - removed += 1 - return None - - out: list[TitleNode] = [] - for node in nodes: - kept = _prune(node, ()) - if kept is not None: - out.append(kept) - - if removed: - logger.info( - "[structure_anchoring] pruned {} unanchored TOC nodes " - "(suffix / no match_overrides → no TOC)", - removed, - ) - return out, removed - - -def toc_range_start(hierarchy: dict[str, Any]) -> int | None: - toc_range = hierarchy.get("toc_range") - if not isinstance(toc_range, (list, tuple)) or not toc_range: - return None - try: - return int(toc_range[0]) - except (TypeError, ValueError): - return None - - -def toc_range_end(hierarchy: dict[str, Any]) -> int | None: - toc_range = hierarchy.get("toc_range") - if not isinstance(toc_range, (list, tuple)) or not toc_range: - return None - try: - return int(toc_range[-1]) - except (TypeError, ValueError): - return None - - -# ── 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( - "[structure_anchoring] 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( - "[structure_anchoring] 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( - "[structure_anchoring] 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( - "[structure_anchoring] 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 - - -# ── Offset-guided bulk anchoring with recursive recalibrate (Phase-2) ─────── - -_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( - "[structure_anchoring] 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 - 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( - "[structure_anchoring] bisect breakpoint: last_valid_index={} / total={}", - lo, - len(leaves), - ) - return lo - -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 - 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( - "[structure_anchoring] recalibrate: title={!r} new_offset={} (delta=+{})", - entry_node.title, - new_offset, - delta, - ) - return new_offset - return None +__all__ = [ + "SkeletonAnchor", + "TitleMatch", + "TitleNode", + "anchor_hierarchy", + "anchor_hierarchy_from_offset", + "bulk_offset_matches", + "deserialize_skeleton_anchor", + "deserialize_title_match", + "locate_null_page_parent_overrides", + "offset_guided_anchoring", + "prune_out_of_scope_nodes", + "prune_unanchored_toc_leaves", + "serialize_skeleton_anchor", + "serialize_title_match", + "toc_range_end", + "toc_range_start", +] + +anchor_hierarchy_from_offset = _anchoring.anchor_hierarchy_from_offset +bulk_offset_matches = _anchoring.bulk_offset_matches +deserialize_skeleton_anchor = _anchoring.deserialize_skeleton_anchor +deserialize_title_match = _anchoring.deserialize_title_match +locate_null_page_parent_overrides = _anchoring.locate_null_page_parent_overrides +prune_out_of_scope_nodes = _anchoring.prune_out_of_scope_nodes +prune_unanchored_toc_leaves = _anchoring.prune_unanchored_toc_leaves +serialize_skeleton_anchor = _anchoring.serialize_skeleton_anchor +serialize_title_match = _anchoring.serialize_title_match +toc_range_end = _anchoring.toc_range_end +toc_range_start = _anchoring.toc_range_start def offset_guided_anchoring( @@ -551,249 +60,19 @@ def offset_guided_anchoring( 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 not leaves: - return dict(calibration_overrides) or None - - all_matches: dict[tuple[str, ...], TitleMatch] = {} - all_matches.update(calibration_overrides) - - # Single-leaf regimes (roman front-matter, F-1 appendix, …) still get a - # deterministic printed→physical override; Phase-1 already calibrated them. - if len(leaves) == 1: - all_matches.update(bulk_offset_matches(leaves, offset)) - else: - _anchor_segment_recursive( - leaves=leaves, + """Forward phase-2 anchoring while preserving the historical patch seam.""" + original = _anchoring.verify_section_page_choice + _anchoring.verify_section_page_choice = verify_section_page_choice + try: + return _anchoring.offset_guided_anchoring( + nodes=nodes, offset=offset, ctx=ctx, page_count=page_count, - matches=all_matches, - depth=0, + calibration_overrides=calibration_overrides, ) - - if not all_matches: - return None - - logger.info( - "[structure_anchoring] 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, - ) - - -@dataclass -class SkeletonAnchor: - offset: int | None - offset_status: str - match_overrides: dict[tuple[str, ...], TitleMatch] - null_page_report: list[dict[str, Any]] - bulk_count: int - pruned_count: int = 0 - locate_agent: str = "offset_only" - - -def serialize_title_match(match: TitleMatch) -> dict[str, Any]: - return { - "page": match.page, - "confidence": match.confidence, - "source": match.source, - "matched_line": match.matched_line, - "score": match.score, - "candidates": list(match.candidates), - "evidence": dict(match.evidence or {}), - } - - -def serialize_skeleton_anchor(anchor: SkeletonAnchor) -> dict[str, Any]: - """JSON-friendly SkeletonAnchor (path tuples joined by ' / ').""" - overrides: dict[str, Any] = {} - for path, match in (anchor.match_overrides or {}).items(): - key = " / ".join(str(part) for part in path) - overrides[key] = serialize_title_match(match) - return { - "offset": anchor.offset, - "offset_status": anchor.offset_status, - "match_overrides": overrides, - "null_page_report": list(anchor.null_page_report or []), - "bulk_count": int(anchor.bulk_count or 0), - "pruned_count": int(anchor.pruned_count or 0), - "locate_agent": anchor.locate_agent, - } - - -def deserialize_title_match(data: dict[str, Any]) -> TitleMatch: - return TitleMatch( - page=int(data["page"]), - confidence=float(data.get("confidence") or 0.0), - source=data.get("source") or "agent_vlm", # type: ignore[arg-type] - matched_line=str(data.get("matched_line") or ""), - score=float(data.get("score") or 0.0), - candidates=[int(p) for p in (data.get("candidates") or [])], - evidence=dict(data.get("evidence") or {}), - ) - - -def deserialize_skeleton_anchor(data: dict[str, Any]) -> SkeletonAnchor: - raw_overrides = data.get("match_overrides") or {} - overrides: dict[tuple[str, ...], TitleMatch] = {} - if isinstance(raw_overrides, dict): - for key, value in raw_overrides.items(): - if not isinstance(value, dict): - continue - if isinstance(key, str): - path = tuple(part.strip() for part in key.split(" / ") if part.strip()) - elif isinstance(key, (list, tuple)): - path = tuple(str(part) for part in key) - else: - continue - if path: - overrides[path] = deserialize_title_match(value) - return SkeletonAnchor( - offset=data.get("offset") if data.get("offset") is None else int(data["offset"]), - offset_status=str(data.get("offset_status") or "failed"), - match_overrides=overrides, - null_page_report=list(data.get("null_page_report") or []), - bulk_count=int(data.get("bulk_count") or 0), - pruned_count=int(data.get("pruned_count") or 0), - locate_agent=str(data.get("locate_agent") or "offset_only"), - ) - - -def anchor_hierarchy_from_offset( - *, - nodes: list[TitleNode], - offset_hint: int | None, - calibration_overrides: dict[tuple[str, ...], TitleMatch] | None = None, - page_texts: dict[int, str], - body_pages: list[int], - page_count: int, - ctx: ToolContext | None, -) -> tuple[list[TitleNode], SkeletonAnchor]: - """Production prune → bulk → null-page given a precomputed offset. - - Phase-2 entry after Agent ``calibrate_offset`` (Phase-1). - """ - seed_overrides = dict(calibration_overrides or {}) - pruned_count = 0 - working = nodes - if offset_hint is not None: - working, pruned_count = prune_out_of_scope_nodes( - working, offset=offset_hint, page_count=page_count - ) - - offset_matches: dict[tuple[str, ...], TitleMatch] | None = None - if offset_hint is not None and ctx is not None and working: - offset_matches = offset_guided_anchoring( - nodes=working, - offset=offset_hint, - ctx=ctx, - page_count=page_count, - calibration_overrides=seed_overrides, - ) - - if offset_matches is not None: - match_overrides = offset_matches - locate_agent = "offset_guided_bulk" - bulk_count = len(offset_matches) - else: - match_overrides = seed_overrides - locate_agent = "offset_only" - bulk_count = 0 - - working, unanchored_removed = prune_unanchored_toc_leaves( - working, match_overrides=match_overrides - ) - pruned_count += unanchored_removed - - match_overrides, null_page_report = locate_null_page_parent_overrides( - nodes=working, - match_overrides=match_overrides, - page_texts=page_texts, - body_pages=body_pages, - ctx=ctx, - ) - - if offset_hint is None: - offset_status = "failed" if ctx is not None else "skipped" - else: - offset_status = "ok" - - return working, SkeletonAnchor( - offset=offset_hint, - offset_status=offset_status, - match_overrides=match_overrides, - null_page_report=null_page_report, - bulk_count=bulk_count, - pruned_count=pruned_count, - locate_agent=locate_agent, - ) + finally: + _anchoring.verify_section_page_choice = original def anchor_hierarchy( @@ -805,45 +84,15 @@ def anchor_hierarchy( page_count: int, ctx: ToolContext | None, ) -> tuple[list[TitleNode], SkeletonAnchor]: - """Run Phase-1 calibrate_offset → multi-regime Phase-2 merge. - - Returns possibly-pruned nodes and the anchor payload. Caller owns - resolve_hierarchy_page_ranges / skeleton assembly. - """ - from app.services.document_agent.agents.calibration.procedure import ( - finalize_calibration_result, - flat_toc_entries, - ) - from app.services.document_agent.agents.calibration.service import ( - calibrate_offset, + """Resolve the calibration-owned orchestration entry point on demand.""" + orchestrator = import_module( + "app.services.document_agent.agents.calibration.orchestrator" ) - - phase1 = calibrate_offset( + return orchestrator.anchor_hierarchy( nodes=nodes, toc_hierarchies=toc_hierarchies, - ctx=ctx, page_texts=page_texts, + body_pages=body_pages, page_count=page_count, - ) - if phase1.status == "failed" and not phase1.regimes and phase1.offset is None: - return anchor_hierarchy_from_offset( - nodes=nodes, - offset_hint=None, - calibration_overrides={}, - page_texts=page_texts, - body_pages=body_pages, - page_count=page_count, - ctx=ctx, - ) - - working, anchor, _finalized = finalize_calibration_result( - result=phase1, - entries=flat_toc_entries(toc_hierarchies), - toc_hierarchies=list(toc_hierarchies or []), ctx=ctx, - page_count=page_count, - page_texts=page_texts, - body_pages=body_pages, - nodes=nodes, ) - return working, anchor diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 6a35e63d..3dfd19c8 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -22,8 +22,8 @@ resolve_hierarchy_page_ranges, ) from app.services.document_agent.agents.calibration import calibrate_offset -from app.services.document_agent.structure.structure_anchoring import ( - anchor_hierarchy, +from app.services.document_agent.agents.calibration.orchestrator import anchor_hierarchy +from app.services.document_agent.structure.anchoring_primitives import ( toc_range_end, toc_range_start, )