From b8d8268c84581e0f928af25a74c9b23d11a956ae Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 25 Aug 2026 20:57:07 +0800 Subject: [PATCH 1/7] feat: enhance TOC processing and page exclusion logic - Introduced new functions for rehoming calibrated forests and handling pending records in the TOC anchoring process. - Updated the `compute_fat_leaf_pages` function to exclude TOC pages from the fat span test, ensuring accurate page calculations. - Refactored various functions to utilize the new `pages_excluding_toc` utility for consistent TOC page exclusion across multiple components. - Removed the deprecated `TocPagePolicy` class, streamlining the codebase. - Added tests to validate the exclusion of TOC pages in fat leaf calculations. --- .../document_agent/structure/toc_anchoring.py | 77 +++- .../document_agent/structure/toc_graft.py | 29 +- .../document_agent/structure/toc_rehome.py | 292 ++++++++++++++ .../services/page_memory/fine_hierarchy.py | 15 +- .../services/page_memory/memory_service.py | 53 ++- .../scripts/page_memory/_debug_pm_shared.py | 18 +- .../debug_pm_stage3_scope_fine_hierarchy.py | 275 +++++++++----- .../page_memory/debug_pm_stage4_assets.py | 31 +- .../debug_pm_stage5_tagging_finalize.py | 89 ++--- .../scripts/page_memory/toc_page_policy.py | 33 -- ...est_page_memory_fine_hierarchy_contract.py | 27 ++ .../contract/test_toc_rehome_contract.py | 355 ++++++++++++++++++ 12 files changed, 1070 insertions(+), 224 deletions(-) create mode 100644 apps/worker/app/services/document_agent/structure/toc_rehome.py delete mode 100644 apps/worker/scripts/page_memory/toc_page_policy.py create mode 100644 apps/worker/tests/contract/test_toc_rehome_contract.py diff --git a/apps/worker/app/services/document_agent/structure/toc_anchoring.py b/apps/worker/app/services/document_agent/structure/toc_anchoring.py index 1e9ccae0b..dd183d5ef 100644 --- a/apps/worker/app/services/document_agent/structure/toc_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/toc_anchoring.py @@ -55,6 +55,10 @@ def run_toc_anchoring(ctx: ToolContext) -> None: beat the confirmed printed TOC pages on coverage, anchor from outline with physical overrides (no calibrate VLM). Otherwise keep the extracted TOC tree and run the existing VLM calibration path. + + After calibration (primary + pending), each forest runs same-tree monotonic + rehome, then pending forests are classified/grafted. Blackboard writes happen + only after that sequence. """ from app.services.document_agent.calibration.orchestrator import ( anchor_hierarchy, @@ -115,6 +119,14 @@ def run_toc_anchoring(ctx: ToolContext) -> None: page_count=page_count, body_pages=body_pages, ) + + # scope-local TOC pre-pass: future hook (no-op this period) + resolve_nodes, skeleton_anchor = _rehome_calibrated_forest( + resolve_nodes, skeleton_anchor + ) + _rehome_pending_records(pending_records) + + if pending_records: _assign_toc_relationships( root_anchor=skeleton_anchor, pending_records=pending_records, @@ -241,6 +253,10 @@ def _write_outline_skeleton( ctx=ctx, ) skeleton_anchor = replace(skeleton_anchor, source="pdf_outline") + # scope-local TOC pre-pass: future hook (no-op this period) + resolve_nodes, skeleton_anchor = _rehome_calibrated_forest( + resolve_nodes, skeleton_anchor + ) ctx.blackboard.skeleton_anchor = serialize_skeleton_anchor(skeleton_anchor) ctx.blackboard.skeleton_nodes = [ serialize_title_node(node) for node in resolve_nodes @@ -256,6 +272,41 @@ def _write_outline_skeleton( return True +def _rehome_calibrated_forest( + nodes: list[TitleNode], + skeleton_anchor: SkeletonAnchor, +) -> tuple[list[TitleNode], SkeletonAnchor]: + from app.services.document_agent.structure.toc_rehome import rehome_skeleton_forest + + rehomed_nodes, rehomed_anchor, _events = rehome_skeleton_forest( + nodes, skeleton_anchor + ) + return rehomed_nodes, rehomed_anchor + + +def _rehome_pending_records(pending_records: list[dict[str, Any]]) -> None: + from app.services.document_agent.structure.toc_rehome import rehome_skeleton_forest + + for record in pending_records: + nodes_raw = record.get("nodes") or [] + anchor_raw = record.get("skeleton_anchor") + if not isinstance(anchor_raw, dict) or not nodes_raw: + continue + nodes = [ + deserialize_title_node(node) + for node in nodes_raw + if isinstance(node, dict) + ] + if not nodes: + continue + rehomed_nodes, rehomed_anchor, _events = rehome_skeleton_forest( + nodes, + deserialize_skeleton_anchor(anchor_raw), + ) + record["nodes"] = [serialize_title_node(node) for node in rehomed_nodes] + record["skeleton_anchor"] = serialize_skeleton_anchor(rehomed_anchor) + + def outline_physical_overrides( hierarchies: list[dict[str, Any]], ) -> dict[tuple[str, ...], TitleMatch]: @@ -348,9 +399,31 @@ def select_global_toc_hierarchies( return selected, pending, summary +def pages_excluding_toc(pages: Any, toc_pages: Any) -> list[int]: + """Return ``pages`` with TOC pages removed (stable order). + + Same exclusion source as TEXT-track MinerU ``exclude_pages`` and anchoring + ``body_pages``: ``toc_result.toc_pages``. + """ + excluded: set[int] = set() + for raw in toc_pages or []: + try: + excluded.add(int(raw)) + except (TypeError, ValueError): + continue + result: list[int] = [] + for raw in pages or []: + try: + page = int(raw) + except (TypeError, ValueError): + continue + if page not in excluded: + result.append(page) + return result + + def body_pages_excluding_toc(toc_pages: Any, page_count: int) -> list[int]: - excluded = {int(page) for page in (toc_pages or [])} - return [page for page in range(1, page_count + 1) if page not in excluded] + return pages_excluding_toc(range(1, page_count + 1), toc_pages) def pending_toc_body_scope( diff --git a/apps/worker/app/services/document_agent/structure/toc_graft.py b/apps/worker/app/services/document_agent/structure/toc_graft.py index be8991ce2..7e094aa77 100644 --- a/apps/worker/app/services/document_agent/structure/toc_graft.py +++ b/apps/worker/app/services/document_agent/structure/toc_graft.py @@ -199,7 +199,7 @@ def _graft_one_child( reason="no_covering_path", ) return primary_children - covering_node = _node_at_path(primary_children, covering) + covering_node = node_at_path(primary_children, covering) if covering_node is None: _record_skip( events=events, @@ -223,7 +223,7 @@ def _graft_one_child( events=events, ), ) - return _replace_node_at_path(primary_children, covering, updated_parent) + return replace_node_at_path(primary_children, covering, updated_parent) def _attach_new_child( @@ -238,7 +238,7 @@ def _attach_new_child( contained_overrides: dict[tuple[str, ...], TitleMatch], events: list[dict[str, Any]], ) -> list[TitleNode]: - grafted = _rebase_levels(child, (parent.level + 1) - child.level) + grafted = rebase_levels(child, (parent.level + 1) - child.level) new_path = (*parent_path, grafted.title) events.append( { @@ -248,14 +248,14 @@ def _attach_new_child( "start": start, } ) - _remap_overrides( + remap_overrides( contained_overrides=contained_overrides, primary_overrides=primary_overrides, old_prefix=contained_path, new_prefix=new_path, drop_root=False, ) - return _insert_by_start( + return insert_by_start( siblings=primary_children, new_node=grafted, start=start, @@ -301,7 +301,7 @@ def _record_skip( events.append(event) -def _remap_overrides( +def remap_overrides( *, contained_overrides: dict[tuple[str, ...], TitleMatch], primary_overrides: dict[tuple[str, ...], TitleMatch], @@ -309,6 +309,7 @@ def _remap_overrides( new_prefix: tuple[str, ...], drop_root: bool, ) -> None: + """Copy override entries under ``old_prefix`` onto ``new_prefix`` paths.""" prefix_len = len(old_prefix) for path, match in contained_overrides.items(): if path[:prefix_len] != old_prefix: @@ -320,15 +321,16 @@ def _remap_overrides( primary_overrides[new_path] = match -def _rebase_levels(node: TitleNode, delta: int) -> TitleNode: +def rebase_levels(node: TitleNode, delta: int) -> TitleNode: + """Shift ``node.level`` (and descendants) by ``delta``.""" return replace( node, level=node.level + delta, - children=[_rebase_levels(child, delta) for child in node.children], + children=[rebase_levels(child, delta) for child in node.children], ) -def _insert_by_start( +def insert_by_start( *, siblings: list[TitleNode], new_node: TitleNode, @@ -336,6 +338,7 @@ def _insert_by_start( parent_path: tuple[str, ...], primary_overrides: dict[tuple[str, ...], TitleMatch], ) -> list[TitleNode]: + """Insert ``new_node`` before the first sibling whose own page is ``> start``.""" insert_at = len(siblings) for index, sibling in enumerate(siblings): match = primary_overrides.get((*parent_path, sibling.title)) @@ -347,7 +350,8 @@ def _insert_by_start( return updated -def _node_at_path(nodes: list[TitleNode], path: tuple[str, ...]) -> TitleNode | None: +def node_at_path(nodes: list[TitleNode], path: tuple[str, ...]) -> TitleNode | None: + """Return the node at ``path`` under ``nodes``, or None if missing.""" current: list[TitleNode] = nodes node: TitleNode | None = None for title in path: @@ -358,11 +362,12 @@ def _node_at_path(nodes: list[TitleNode], path: tuple[str, ...]) -> TitleNode | return node -def _replace_node_at_path( +def replace_node_at_path( nodes: list[TitleNode], path: tuple[str, ...], new_node: TitleNode, ) -> list[TitleNode]: + """Replace the node at ``path`` with ``new_node`` (immutable list update).""" title, *rest = path updated: list[TitleNode] = [] for node in nodes: @@ -375,7 +380,7 @@ def _replace_node_at_path( updated.append( replace( node, - children=_replace_node_at_path(node.children, tuple(rest), new_node), + children=replace_node_at_path(node.children, tuple(rest), new_node), ) ) return updated diff --git a/apps/worker/app/services/document_agent/structure/toc_rehome.py b/apps/worker/app/services/document_agent/structure/toc_rehome.py new file mode 100644 index 000000000..f459aa76c --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/toc_rehome.py @@ -0,0 +1,292 @@ +"""Same-forest TOC rehome: layer-wise own-page monotonic repair before classify.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any + +from loguru import logger + +from app.services.document_agent.structure.anchoring_primitives import SkeletonAnchor +from app.services.document_agent.structure.hierarchy_locator import TitleMatch, TitleNode +from app.services.document_agent.structure.toc_graft import ( + insert_by_start, + node_at_path, + rebase_levels, + replace_node_at_path, +) + +_LOG_PREFIX = "[profile.toc_rehome]" + + +@dataclass(frozen=True) +class RehomePlan: + source_path: tuple[str, ...] + dest_parent_path: tuple[str, ...] + physical_page: int + + +@dataclass(frozen=True) +class RehomeResult: + nodes: list[TitleNode] + match_overrides: dict[tuple[str, ...], TitleMatch] + events: list[dict[str, Any]] + + +def own_page( + path: tuple[str, ...], + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> int | None: + """Node's own calibrated physical page; never inferred from descendants.""" + match = match_overrides.get(path) + if match is None or match.page is None: + return None + return int(match.page) + + +def rehome_forest( + nodes: list[TitleNode], + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> RehomeResult: + """Repair non-monotonic own-page order within one calibrated TOC forest.""" + overrides = dict(match_overrides) + plans = _collect_plans(siblings=list(nodes), parent_path=(), overrides=overrides) + if not plans: + return RehomeResult(nodes=list(nodes), match_overrides=overrides, events=[]) + + events: list[dict[str, Any]] = [ + { + "action": "plan", + "source_path": list(plan.source_path), + "dest_parent_path": list(plan.dest_parent_path), + "physical_page": plan.physical_page, + } + for plan in plans + ] + logger.info("{} plans={}", _LOG_PREFIX, len(plans)) + new_nodes, overrides = _apply_plans(list(nodes), overrides, plans) + return RehomeResult(nodes=new_nodes, match_overrides=overrides, events=events) + + +def rehome_skeleton_forest( + nodes: list[TitleNode], + anchor: SkeletonAnchor, +) -> tuple[list[TitleNode], SkeletonAnchor, list[dict[str, Any]]]: + """Run ``rehome_forest`` and return an updated ``SkeletonAnchor``.""" + result = rehome_forest(nodes, anchor.match_overrides) + return ( + result.nodes, + replace(anchor, match_overrides=result.match_overrides), + result.events, + ) + + +def _collect_plans( + *, + siblings: list[TitleNode], + parent_path: tuple[str, ...], + overrides: dict[tuple[str, ...], TitleMatch], +) -> list[RehomePlan]: + plans: list[RehomePlan] = [] + effective = _expand_effective( + siblings=siblings, + parent_path=parent_path, + overrides=overrides, + ) + cursor: int | None = None + for _node, source_path, page in effective: + if cursor is not None and page < cursor: + plans.append( + RehomePlan( + source_path=source_path, + dest_parent_path=parent_path, + physical_page=page, + ) + ) + continue + if cursor is None or page > cursor: + cursor = page + + # Drill into every own-paged effective node (includes transparent promotions). + # Unpaged shells are not parents here — their children already appear in effective. + for node, source_path, _page in effective: + if node.children: + plans.extend( + _collect_plans( + siblings=list(node.children), + parent_path=source_path, + overrides=overrides, + ) + ) + return plans + + +def _expand_effective( + *, + siblings: list[TitleNode], + parent_path: tuple[str, ...], + overrides: dict[tuple[str, ...], TitleMatch], +) -> list[tuple[TitleNode, tuple[str, ...], int]]: + """Build this level's ordered sequence of own-paged nodes (unpaged non-leaves transparent).""" + effective: list[tuple[TitleNode, tuple[str, ...], int]] = [] + for sibling in siblings: + path = (*parent_path, sibling.title) + page = own_page(path, overrides) + if page is not None: + effective.append((sibling, path, page)) + continue + if sibling.children: + effective.extend( + _expand_effective( + siblings=list(sibling.children), + parent_path=path, + overrides=overrides, + ) + ) + return effective + + +def _apply_plans( + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], TitleMatch], + plans: list[RehomePlan], +) -> tuple[list[TitleNode], dict[tuple[str, ...], TitleMatch]]: + ordered = sorted( + plans, + key=lambda plan: (-len(plan.source_path), plan.physical_page, plan.source_path), + ) + working = list(nodes) + for plan in ordered: + working, detached = _detach_node(working, plan.source_path) + if detached is None: + raise ValueError( + f"{_LOG_PREFIX} detach miss source_path={plan.source_path!r}" + ) + target_level = len(plan.dest_parent_path) + 1 + detached = rebase_levels(detached, target_level - detached.level) + new_path = (*plan.dest_parent_path, detached.title) + working = _insert_under_parent( + nodes=working, + parent_path=plan.dest_parent_path, + new_node=detached, + start=plan.physical_page, + overrides=overrides, + ) + _move_override_prefix(overrides, plan.source_path, new_path) + working, overrides = _drop_empty_ancestors( + working, + overrides, + plan.source_path[:-1], + ) + logger.info( + "{} moved {} -> {} page={}", + _LOG_PREFIX, + plan.source_path, + new_path, + plan.physical_page, + ) + return working, overrides + + +def _detach_node( + nodes: list[TitleNode], + path: tuple[str, ...], +) -> tuple[list[TitleNode], TitleNode | None]: + if not path: + return nodes, None + title, *rest = path + if not rest: + detached: TitleNode | None = None + kept: list[TitleNode] = [] + for node in nodes: + if detached is None and node.title == title: + detached = node + continue + kept.append(node) + return kept, detached + + updated: list[TitleNode] = [] + detached_node: TitleNode | None = None + for node in nodes: + if node.title != title: + updated.append(node) + continue + new_children, detached_node = _detach_node(list(node.children), tuple(rest)) + updated.append(replace(node, children=new_children)) + return updated, detached_node + + +def _drop_empty_ancestors( + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], TitleMatch], + ancestor_path: tuple[str, ...], +) -> tuple[list[TitleNode], dict[tuple[str, ...], TitleMatch]]: + """Delete parents that became empty after children moved out.""" + if not ancestor_path: + return nodes, overrides + working = nodes + # Drop from deepest ancestor upward. + for depth in range(len(ancestor_path), 0, -1): + path = ancestor_path[:depth] + node = node_at_path(working, path) + if node is None: + continue + if node.children: + break + working, removed = _detach_node(working, path) + if removed is not None: + overrides.pop(path, None) + logger.info("{} drop empty shell path={}", _LOG_PREFIX, path) + return working, overrides + + +def _move_override_prefix( + overrides: dict[tuple[str, ...], TitleMatch], + old_prefix: tuple[str, ...], + new_prefix: tuple[str, ...], +) -> None: + if old_prefix == new_prefix: + return + prefix_len = len(old_prefix) + moved: dict[tuple[str, ...], TitleMatch] = {} + for path in list(overrides.keys()): + if path[:prefix_len] != old_prefix: + continue + match = overrides.pop(path) + moved[new_prefix + path[prefix_len:]] = match + overrides.update(moved) + + +def _insert_under_parent( + *, + nodes: list[TitleNode], + parent_path: tuple[str, ...], + new_node: TitleNode, + start: int, + overrides: dict[tuple[str, ...], TitleMatch], +) -> list[TitleNode]: + if not parent_path: + return insert_by_start( + siblings=nodes, + new_node=new_node, + start=start, + parent_path=(), + primary_overrides=overrides, + ) + parent = node_at_path(nodes, parent_path) + if parent is None: + raise ValueError( + f"{_LOG_PREFIX} insert parent missing path={parent_path!r}" + ) + new_children = insert_by_start( + siblings=list(parent.children), + new_node=new_node, + start=start, + parent_path=parent_path, + primary_overrides=overrides, + ) + return replace_node_at_path( + nodes, + parent_path, + replace(parent, children=new_children), + ) diff --git a/apps/worker/app/services/page_memory/fine_hierarchy.py b/apps/worker/app/services/page_memory/fine_hierarchy.py index 57db6362b..287e8d542 100644 --- a/apps/worker/app/services/page_memory/fine_hierarchy.py +++ b/apps/worker/app/services/page_memory/fine_hierarchy.py @@ -107,6 +107,7 @@ def refine_fat_leaf_skeletons( def compute_fat_leaf_pages( skeletons: list[SectionSkeleton], min_pages: int, + toc_pages: list[int] | None = None, ) -> set[int]: """Compute the set of page indices belonging to fat-leaf sections. @@ -116,13 +117,21 @@ def compute_fat_leaf_pages( Uses exclusive end boundaries when sibling starts are visible in ``skeletons``. For single-leaf scopes the closed ``end_page`` is used, which includes the shared boundary page with the next leaf. + + ``toc_pages`` are excluded from both the fat span test and the returned + page set (same body-page rule as scope processing). """ + excluded = {int(page) for page in (toc_pages or [])} fat_pages: set[int] = set() for idx, skel in enumerate(skeletons): exclusive_end = _exclusive_end(skeletons, idx) - page_span = exclusive_end - skel.start_page + 1 - if page_span > min_pages: - fat_pages.update(range(skel.start_page, exclusive_end + 1)) + body_pages = [ + page + for page in range(skel.start_page, exclusive_end + 1) + if page not in excluded + ] + if len(body_pages) > min_pages: + fat_pages.update(body_pages) return fat_pages diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index abef06c4b..d08d4c501 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -264,6 +264,7 @@ def _build_page_dataframe( ) ] + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.fine_hierarchy import build_next_title_by_path next_title_by_path = build_next_title_by_path(skeletons) @@ -272,9 +273,14 @@ def _build_page_dataframe( filename=filename, page_count=page_count, ) - coarse_pages_scope = _derive_hierarchy_page_scope( - skeletons=skeletons, - page_count=page_count, + toc_result = getattr(anatomy, "toc_result", None) if anatomy is not None else None + toc_pages = list(getattr(toc_result, "toc_pages", None) or []) + coarse_pages_scope = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=skeletons, + page_count=page_count, + ), + toc_pages, ) _record_trace_stage( trace_recorder, @@ -308,6 +314,11 @@ def _build_page_dataframe( len(coarse_scopes), scope_concurrency, ) + if toc_pages: + logger.info( + "[page_memory] excluding TOC pages from scope processing: {}", + sorted({int(page) for page in toc_pages}), + ) scope_results: list[_ScopeRunResult] = [] asset_pages_remaining = ( _resolve_asset_max_pages(page_count, page_memory_config) @@ -329,6 +340,7 @@ def _build_page_dataframe( trace_recorder=trace_recorder, page_memory_config=page_memory_config, next_title_by_path=next_title_by_path, + toc_pages=toc_pages, ) if scope_concurrency <= 1 or len(coarse_scopes) <= 1: @@ -409,9 +421,12 @@ def _build_page_dataframe( # Shared per-page lookups for node-granularity assembly. raw_text_by_page: dict[int, str] = {} image_path_by_page: dict[int, str] = {} - final_pages_scope = _derive_hierarchy_page_scope( - skeletons=skeletons, - page_count=page_count, + final_pages_scope = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=skeletons, + page_count=page_count, + ), + toc_pages, ) for page in final_pages_scope: rend = render_map.get(page) @@ -538,7 +553,9 @@ def _run_hierarchy_scope( trace_recorder: Any | None, page_memory_config: PageMemoryConfig, next_title_by_path: dict[str, str | None] | None = None, + toc_pages: list[int] | None = None, ) -> _ScopeRunResult: + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.fine_hierarchy import ( compute_fat_leaf_pages, refine_fat_leaf_skeletons, @@ -559,9 +576,12 @@ def _run_hierarchy_scope( page_count=page_count, strategy=scope.strategy, ) - coarse_pages = _derive_hierarchy_page_scope( - skeletons=scope_skeletons, - page_count=page_count, + coarse_pages = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=scope_skeletons, + page_count=page_count, + ), + toc_pages, ) logger.info( "[page_memory] scope {}/{} {} coarse ranges={}", @@ -582,7 +602,11 @@ def _run_hierarchy_scope( ) fine_min = page_memory_config.fine_min_pages - fat_leaf_pages = compute_fat_leaf_pages(scope_skeletons, min_pages=fine_min) + fat_leaf_pages = compute_fat_leaf_pages( + scope_skeletons, + min_pages=fine_min, + toc_pages=toc_pages, + ) if fat_leaf_pages: title_pages = sorted(fat_leaf_pages) with stage_timer("page_memory.title_render", page_count=len(title_pages)): @@ -656,9 +680,12 @@ def _run_hierarchy_scope( }, ) - final_pages = _derive_hierarchy_page_scope( - skeletons=scope_skeletons, - page_count=page_count, + final_pages = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=scope_skeletons, + page_count=page_count, + ), + toc_pages, ) final_scope_summary = _summarize_tag_scope( skeletons=scope_skeletons, diff --git a/apps/worker/scripts/page_memory/_debug_pm_shared.py b/apps/worker/scripts/page_memory/_debug_pm_shared.py index be93e6030..8699ee7c1 100644 --- a/apps/worker/scripts/page_memory/_debug_pm_shared.py +++ b/apps/worker/scripts/page_memory/_debug_pm_shared.py @@ -1509,10 +1509,17 @@ def build_debug_coarse_scopes( page_count: int, anatomy: Any | None = None, ) -> list[dict[str, Any]]: + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory._utils import build_hierarchy_scopes - from toc_page_policy import TocPagePolicy - policy = TocPagePolicy.from_anatomy(anatomy) + toc_result = getattr(anatomy, "toc_result", None) if anatomy is not None else None + toc_pages = list(getattr(toc_result, "toc_pages", None) or []) + toc_page_set = set() + for raw in toc_pages: + try: + toc_page_set.add(int(raw)) + except (TypeError, ValueError): + continue scopes = build_hierarchy_scopes( skeletons=skeletons, filename=filename, @@ -1525,12 +1532,13 @@ def build_debug_coarse_scopes( "start_page": scope.start_page, "end_page": scope.end_page, "strategy": scope.strategy, - "processing_pages": policy.filter_processing_pages( - list(range(scope.start_page, scope.end_page + 1)) + "processing_pages": pages_excluding_toc( + list(range(scope.start_page, scope.end_page + 1)), + toc_pages, ), "excluded_toc_pages": sorted( page - for page in policy.pure_toc_pages + for page in toc_page_set if scope.start_page <= page <= scope.end_page ), } diff --git a/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py b/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py index 81a2e7e9d..700f03219 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py @@ -1,10 +1,9 @@ #!/usr/bin/env python3 # ruff: noqa: E402 -"""Stage 3: Coarse scopes + document page tagging + per-scope fine hierarchy. +"""Stage 3: Coarse scopes + per-scope title detection / fine hierarchy / page tag. -Builds coarse hierarchy scopes, writes ``scopes//skeletons.json``, renders -and tags each selected processing page once (global concurrency), then fans -tag subsets into scopes for fine hierarchy refinement. +Mirrors production ``_run_hierarchy_scope`` (without C5 assets; that is Stage 4): +fat-leaf title render → ``tag_page_titles`` → fine hierarchy → render → ``tag_pages``. Requires Stage 2 output: _doc_agent/pipeline_state.json (with skeletons), doc_profile.json (after production ``run_toc_anchoring``). @@ -65,28 +64,29 @@ def _resolve_scope_processing_pages( scope_meta: dict[str, Any], skeletons: list[Any], page_count: int, - toc_policy: Any, + toc_pages: list[int], ) -> tuple[list[int], list[int]]: + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc + processing_pages = [ int(page) for page in (scope_meta.get("processing_pages") or []) ] if not processing_pages: - processing_pages = toc_policy.filter_processing_pages( + processing_pages = pages_excluding_toc( _derive_hierarchy_page_scope( skeletons=skeletons, page_count=page_count, - ) + ), + toc_pages, ) excluded_toc_pages = [ int(page) for page in (scope_meta.get("excluded_toc_pages") or []) ] if not excluded_toc_pages: + start = int(scope_meta.get("start_page") or 1) + end = int(scope_meta.get("end_page") or page_count) excluded_toc_pages = sorted( - set(range( - int(scope_meta.get("start_page") or 1), - int(scope_meta.get("end_page") or page_count) + 1, - )) - & toc_policy.pure_toc_pages + page for page in toc_pages if start <= int(page) <= end ) return processing_pages, excluded_toc_pages @@ -96,20 +96,35 @@ def _run_fine_hierarchy_for_scope( scope_id: str, scope_dir: Path, out_dir: Path, + pdf_path: str, page_count: int, - rendered_by_page: dict[int, Any], - tags_by_page: dict[int, Any], + page_texts: dict[int, str], + page_features: list[Any], + page_labels: list[Any], + vlm_model: str | None, next_title_by_path: dict[str, str | None], - toc_policy: Any, + toc_pages: list[int], page_memory_config: Any, token_cost_tracker: TokenCostTracker | None = None, ) -> ScopeResult: - """Consume shared page tags and refine one coarse scope.""" + """Per-scope path aligned with production ``_run_hierarchy_scope`` (no C5).""" + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.fine_hierarchy import ( compute_fat_leaf_pages, refine_fat_leaf_skeletons, ) - from app.services.page_memory.memory_service import _resolve_hierarchy_model + from app.services.page_memory.memory_service import ( + _resolve_hierarchy_model, + _summarize_tag_scope, + _summarize_tags, + ) + from app.services.page_memory.page_plan import derive_page_processing_plan + from app.services.page_memory.page_renderer import render_document_pages + from app.services.page_memory.page_tagger import ( + PageTagResult, + tag_page_titles, + tag_pages, + ) scope_stages: list[dict[str, Any]] = [] if token_cost_tracker is not None: @@ -118,12 +133,13 @@ def _run_fine_hierarchy_for_scope( skel_path = scope_dir / "skeletons.json" require_file(skel_path, hint=f"Stage 3 should have created {skel_path}") scope_meta, active_skeletons = load_scope_skeletons_artifact(skel_path) + active_skeletons = sort_skeletons(active_skeletons) strategy = str(scope_meta.get("strategy") or "coarse_scope") processing_pages, excluded_toc_pages = _resolve_scope_processing_pages( scope_meta=scope_meta, skeletons=active_skeletons, page_count=page_count, - toc_policy=toc_policy, + toc_pages=toc_pages, ) scope_manifest = _scope_manifest( @@ -131,53 +147,65 @@ def _run_fine_hierarchy_for_scope( skeletons=active_skeletons, page_count=page_count, strategy=strategy, - processing_pages=processing_pages, - excluded_toc_pages=excluded_toc_pages, ) logger.info( - "🔬 [scope {}] {} skeletons p{}-{} processing={}", + "🔬 [scope {}] {} skeletons p{}-{} processing={} excluded_toc={}", scope_id, len(active_skeletons), scope_meta.get("start_page", "?"), scope_meta.get("end_page", "?"), processing_pages, + excluded_toc_pages, ) - if not processing_pages: - logger.info(" [scope {}] no processing pages after TOC exclusion", scope_id) - return ScopeResult( - scope_id=scope_id, - skeletons=active_skeletons, - tags=[], - assets_by_page={}, - rendered=[], - final_pages=[], - scope_manifest=scope_manifest, - trace_stages=scope_stages, - ) - - rendered = [ - rendered_by_page[page] - for page in processing_pages - if page in rendered_by_page - ] - tags = [ - tags_by_page[page] - for page in processing_pages - if page in tags_by_page - ] - fine_min = page_memory_config.fine_min_pages fat_leaf_pages = compute_fat_leaf_pages( active_skeletons, min_pages=fine_min, - exclude_pages=toc_policy.pure_toc_pages, + toc_pages=toc_pages, ) if fat_leaf_pages: + title_pages = sorted(fat_leaf_pages) + title_rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=str(out_dir), + pages=title_pages, + page_features=page_features, + page_texts=page_texts, + ) + title_tags = [ + PageTagResult( + page_index=page, + summary="", + keywords=[], + strategy_used="title_detection_only", + ) + for page in title_pages + ] + title_tags = tag_page_titles( + pages=title_rendered, + tag_results=title_tags, + fat_leaf_pages=fat_leaf_pages, + vlm_model=vlm_model, + scan_direction=page_memory_config.scan_direction, + max_concurrent=page_memory_config.title_detection_concurrency, + ) + record_stage( + scope_stages, + "C3b.title_detection", + page_info=page_scope_info(title_pages), + variables={ + "scope_id": scope_id, + "tags": _summarize_tags(title_tags), + }, + ) + if token_cost_tracker is not None: + token_cost_tracker.snapshot_stage(f"C3b.title_detection:{scope_id}") active_skeletons = refine_fat_leaf_skeletons( coarse_skeletons=active_skeletons, - tag_results=tags, + tag_results=title_tags, fat_leaf_pages=fat_leaf_pages, next_title_by_path=next_title_by_path, model_name=_resolve_hierarchy_model(page_memory_config), @@ -200,8 +228,6 @@ def _run_fine_hierarchy_for_scope( skeletons=active_skeletons, page_count=page_count, strategy=f"{strategy}:refined", - processing_pages=processing_pages, - excluded_toc_pages=excluded_toc_pages, ) record_stage( scope_stages, "C4b.fine_hierarchy", @@ -215,6 +241,66 @@ def _run_fine_hierarchy_for_scope( if token_cost_tracker is not None: token_cost_tracker.snapshot_stage(f"C4b.fine_hierarchy:{scope_id}") + final_pages = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=active_skeletons, + page_count=page_count, + ), + toc_pages, + ) + final_scope_summary = _summarize_tag_scope( + skeletons=active_skeletons, + page_count=page_count, + pages=final_pages, + ) + rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=str(out_dir), + pages=final_pages, + page_features=page_features, + page_texts=page_texts, + ) + record_stage( + scope_stages, + "C1.render_pages", + page_info=page_scope_info([item.page_index for item in rendered]), + variables={ + "scope_id": scope_id, + "rendered_count": len(rendered), + "tag_scope": final_scope_summary, + }, + ) + + plans = derive_page_processing_plan( + page_count=page_count, + page_labels=page_labels, + page_features=page_features, + ) + final_page_set = set(final_pages) + plans = [plan for plan in plans if plan.page_index in final_page_set] + record_stage( + scope_stages, + "C2.page_plan", + page_info=page_scope_info([getattr(plan, "page_index", None) for plan in plans]), + variables={"scope_id": scope_id, "plan_count": len(plans)}, + ) + + tags = tag_pages( + pages=rendered, + plans=plans, + vlm_model=vlm_model, + max_concurrent=page_memory_config.tag_concurrency, + ) + record_stage( + scope_stages, + "C3.page_tagger", + page_info=page_scope_info([tag.page_index for tag in tags]), + variables={"scope_id": scope_id, "tags": _summarize_tags(tags)}, + ) + if token_cost_tracker is not None: + token_cost_tracker.snapshot_stage(f"C3.page_tagger:{scope_id}") + write_scope_artifacts( out_dir=out_dir, scope_id=scope_id, @@ -229,7 +315,7 @@ def _run_fine_hierarchy_for_scope( tags=tags, assets_by_page={}, rendered=rendered, - final_pages=processing_pages, + final_pages=final_pages, scope_manifest=scope_manifest, trace_stages=scope_stages, ) @@ -245,10 +331,9 @@ def main() -> int: args = parser.parse_args() from app.services.document_agent.pdf_text import read_page_texts + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.fine_hierarchy import build_next_title_by_path - from app.services.page_memory.memory_service import _render_and_tag_document_pages from app.services.page_memory.skeleton_extractor import SectionSkeleton - from toc_page_policy import TocPagePolicy from shared.models.schemas.page_memory_config import PageMemoryConfig pdf_path, filename, out_dir = resolve_paths(args) @@ -269,7 +354,8 @@ def main() -> int: page_count = anatomy.page_count page_features = anatomy.page_features if anatomy else [] page_labels = anatomy.page_labels if anatomy else [] - toc_policy = TocPagePolicy.from_anatomy(anatomy) + toc_result = getattr(anatomy, "toc_result", None) + toc_pages = list(getattr(toc_result, "toc_pages", None) or []) page_memory_config = PageMemoryConfig.default() skeletons = load_pipeline_skeletons(state_path) @@ -307,10 +393,13 @@ def main() -> int: "start_page": 1, "end_page": page_count, "strategy": "fallback_root", - "processing_pages": toc_policy.filter_processing_pages( - list(range(1, page_count + 1)) + "processing_pages": pages_excluding_toc( + list(range(1, page_count + 1)), + toc_pages, + ), + "excluded_toc_pages": sorted( + {int(page) for page in toc_pages} ), - "excluded_toc_pages": sorted(toc_policy.pure_toc_pages), } ] logger.info(" no skeleton hierarchy → fallback Root scope p1-{}", page_count) @@ -373,11 +462,14 @@ def main() -> int: "start_page": pr_start, "end_page": pr_end, "strategy": "manual_page_range", - "processing_pages": toc_policy.filter_processing_pages( - requested_pages + "processing_pages": pages_excluding_toc( + requested_pages, + toc_pages, ), "excluded_toc_pages": sorted( - set(requested_pages) & toc_policy.pure_toc_pages + page + for page in toc_pages + if pr_start <= int(page) <= pr_end ), } ] @@ -450,66 +542,51 @@ def main() -> int: sum(1 for title in next_title_by_path.values() if title), ) - selected_processing_pages: set[int] = set() - scope_payloads: list[tuple[str, Path, list[int]]] = [] + scope_payloads: list[tuple[str, Path]] = [] + pages_needed: set[int] = set() for sid in scope_ids: scope_dir = scopes_dir / sid - scope_meta, scope_skeletons = load_scope_skeletons_artifact( + _scope_meta, scope_skeletons = load_scope_skeletons_artifact( scope_dir / "skeletons.json" ) - processing_pages, _excluded = _resolve_scope_processing_pages( - scope_meta=scope_meta, - skeletons=scope_skeletons, - page_count=page_count, - toc_policy=toc_policy, + pages_needed.update( + pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=scope_skeletons, + page_count=page_count, + ), + toc_pages, + ) ) - selected_processing_pages.update(processing_pages) - scope_payloads.append((sid, scope_dir, processing_pages)) + scope_payloads.append((sid, scope_dir)) - processing_pages = sorted(selected_processing_pages) page_texts = read_page_texts( pdf_path, - processing_pages or list(range(1, page_count + 1)), + sorted(pages_needed) or list(range(1, page_count + 1)), timeout=300, ) logger.info( - " document page stage: {} processing pages (of {}), tag_concurrency={}", - len(processing_pages), + " page texts ready: {} pages (of {}), tag_concurrency={}", + len(page_texts), page_count, page_memory_config.tag_concurrency, ) vlm_model = getattr(args, "vlm_model", None) or os.environ.get("IMAGE_MODEL") - rendered_by_page, tags_by_page = _render_and_tag_document_pages( - pdf_path=pdf_path, - output_dir=str(out_dir), - page_count=page_count, - processing_pages=processing_pages, - page_texts=page_texts, - page_features=page_features, - page_labels=page_labels, - vlm_model=vlm_model, - toc_policy=toc_policy, - page_memory_config=page_memory_config, - trace_recorder=TraceStageAdapter(trace_stages), - ) - token_cost_tracker.snapshot_stage("C3.page_tagger") - logger.info( - " tagged {} unique pages; refining {} scopes", - len(tags_by_page), - len(scope_ids), - ) def _run_selected_scope(scope_id: str, scope_dir: Path) -> ScopeResult: return _run_fine_hierarchy_for_scope( scope_id=scope_id, scope_dir=scope_dir, out_dir=out_dir, + pdf_path=str(pdf_path), page_count=page_count, - rendered_by_page=rendered_by_page, - tags_by_page=tags_by_page, + page_texts=page_texts, + page_features=page_features, + page_labels=page_labels, + vlm_model=vlm_model, next_title_by_path=next_title_by_path, - toc_policy=toc_policy, + toc_pages=toc_pages, page_memory_config=page_memory_config, token_cost_tracker=token_cost_tracker, ) @@ -529,7 +606,7 @@ def _run_selected_scope(scope_id: str, scope_dir: Path) -> ScopeResult: sid, scope_dir, ) - for sid, scope_dir, _pages in scope_payloads + for sid, scope_dir in scope_payloads ] gevent.joinall(greenlets, raise_error=True) scope_results = [cast(ScopeResult, g.value) for g in greenlets] @@ -537,7 +614,7 @@ def _run_selected_scope(scope_id: str, scope_dir: Path) -> ScopeResult: logger.info(" serial fine hierarchy: {} scope(s)", len(scope_ids)) scope_results = [ _run_selected_scope(sid, scope_dir) - for sid, scope_dir, _pages in scope_payloads + for sid, scope_dir in scope_payloads ] for sr in scope_results: @@ -546,6 +623,10 @@ def _run_selected_scope(scope_id: str, scope_dir: Path) -> ScopeResult: merged_skeletons = sort_skeletons( [skel for sr in scope_results for skel in sr.skeletons] ) + tags_by_page: dict[int, Any] = {} + for sr in scope_results: + for tag in sr.tags: + tags_by_page[int(tag.page_index)] = tag merged_tags = [tags_by_page[page] for page in sorted(tags_by_page)] if not partial_run: write_top_level_artifacts( diff --git a/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py b/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py index 41b06b105..9f657b55b 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py @@ -64,8 +64,10 @@ def _load_scope_asset_context( scope_id: str, scope_dir: Path, page_count: int, - toc_policy: Any, + toc_pages: list[int], ) -> _ScopeAssetContext | None: + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc + fine_hierarchy_path = scope_dir / "fine_hierarchy.json" require_file( fine_hierarchy_path, @@ -83,26 +85,20 @@ def _load_scope_asset_context( final_pages = ( [int(page) for page in recorded_pages] if isinstance(recorded_pages, list) - else toc_policy.filter_processing_pages( + else pages_excluding_toc( _derive_hierarchy_page_scope( skeletons=active_skeletons, page_count=page_count, - ) + ), + toc_pages, ) ) - recorded_excluded = prior_scope.get("excluded_toc_pages") - excluded_toc_pages = ( - [int(page) for page in recorded_excluded] - if isinstance(recorded_excluded, list) - else sorted(toc_policy.pure_toc_pages) - ) + final_pages = pages_excluding_toc(final_pages, toc_pages) scope_manifest = _scope_manifest( scope_id=scope_id, skeletons=active_skeletons, page_count=page_count, strategy="fine:assets", - processing_pages=final_pages, - excluded_toc_pages=excluded_toc_pages, ) return _ScopeAssetContext( scope_id=scope_id, @@ -135,7 +131,6 @@ def main() -> int: page_asset_summary_enabled, ) from app.services.page_memory.page_renderer import render_document_pages - from toc_page_policy import TocPagePolicy from shared.models.schemas.page_memory_config import PageMemoryConfig pdf_path, filename, out_dir = resolve_paths(args) @@ -151,7 +146,8 @@ def main() -> int: anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) page_count = anatomy.page_count page_features = anatomy.page_features if anatomy else [] - toc_policy = TocPagePolicy.from_anatomy(anatomy) + toc_result = getattr(anatomy, "toc_result", None) + toc_pages = list(getattr(toc_result, "toc_pages", None) or []) scope_ids = resolve_debug_scope_ids( scopes_dir=scopes_dir, @@ -183,13 +179,16 @@ def main() -> int: scope_id=scope_id, scope_dir=scopes_dir / scope_id, page_count=page_count, - toc_policy=toc_policy, + toc_pages=toc_pages, ) if context is not None: scope_contexts.append(context) - union_pages = sorted( - {page for context in scope_contexts for page in context.pages} + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc + + union_pages = pages_excluding_toc( + sorted({page for context in scope_contexts for page in context.pages}), + toc_pages, ) logger.info( "🔬 document C5: {} unique pages across {} scopes", diff --git a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py index 281040759..5321decf1 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py @@ -71,11 +71,12 @@ def _run_tagging_for_scope( page_count: int, page_texts: dict[int, str], page_features: list[Any], - toc_policy: Any, + toc_pages: list[int], args: Any, token_cost_tracker: TokenCostTracker | None = None, ) -> ScopeResult: """Load Stage-3 combined tags and rehydrate renders for final assembly.""" + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.page_renderer import render_document_pages scope_stages: list[dict[str, Any]] = [] @@ -114,26 +115,20 @@ def _run_tagging_for_scope( final_pages = ( [int(page) for page in recorded_pages] if isinstance(recorded_pages, list) - else toc_policy.filter_processing_pages( + else pages_excluding_toc( _derive_hierarchy_page_scope( skeletons=active_skeletons, page_count=page_count, - ) + ), + toc_pages, ) ) - recorded_excluded = prior_scope.get("excluded_toc_pages") - excluded_toc_pages = ( - [int(page) for page in recorded_excluded] - if isinstance(recorded_excluded, list) - else sorted(toc_policy.pure_toc_pages) - ) + final_pages = pages_excluding_toc(final_pages, toc_pages) scope_manifest = _scope_manifest( scope_id=scope_id, skeletons=active_skeletons, page_count=page_count, strategy="fine:finalize", - processing_pages=final_pages, - excluded_toc_pages=excluded_toc_pages, ) logger.info( "🔬 [scope {}] loaded {} combined tags for {} processing pages", @@ -274,16 +269,7 @@ def main() -> int: args.finalize = True from app.services.document_agent.pdf_text import read_page_texts - from app.services.document_agent.structure.hierarchy_locator import extract_toc_nodes - from app.services.page_memory.memory_service import ( - _append_toc_nav_skeletons, - _merge_static_toc_tags, - ) - from app.services.page_memory.node_assembler import ( - build_node_chunks, - build_toc_node_chunks, - merge_chunks_by_first_page, - ) + from app.services.page_memory.node_assembler import build_node_rows from app.services.page_memory.skeleton_extractor import collapse_single_child_chains from shared.models.schemas.page_memory_config import PageMemoryConfig @@ -300,9 +286,9 @@ def main() -> int: anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) page_count = anatomy.page_count page_features = anatomy.page_features if anatomy else [] - from toc_page_policy import TocPagePolicy - - toc_policy = TocPagePolicy.from_anatomy(anatomy) + page_labels = anatomy.page_labels if anatomy else [] + toc_result = getattr(anatomy, "toc_result", None) + toc_pages = list(getattr(toc_result, "toc_pages", None) or []) scope_ids = resolve_debug_scope_ids( scopes_dir=scopes_dir, @@ -340,7 +326,7 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: page_count=page_count, page_texts=page_texts, page_features=page_features, - toc_policy=toc_policy, + toc_pages=toc_pages, args=args, token_cost_tracker=token_cost_tracker, ) @@ -440,22 +426,17 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: ) active_skeletons = all_skeletons - tags = _merge_static_toc_tags(all_tags, toc_policy) - nav_skeletons = _append_toc_nav_skeletons( - body_skeletons=active_skeletons, - anatomy=anatomy, - filename=filename, - ) + tags = all_tags # Write top-level artifacts write_top_level_artifacts( out_dir=out_dir, - hierarchy=nav_skeletons, + hierarchy=active_skeletons, tags=tags, assets_by_page=all_assets if all_assets else None, ) - # ── C7: Node assembly ── + # ── C7: Node assembly (same entry as production memory_service) ── logger.info("=" * 70) logger.info("🧱 C7: assemble canonical chunks") logger.info("=" * 70) @@ -473,25 +454,47 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: image_path_by_page[page] = rend.image_path page_memory_config = PageMemoryConfig.default() - body_chunks = build_node_chunks( + label_map: dict[int, str] = {} + if page_labels: + for lbl in page_labels: + label_map[int(lbl.page)] = str(lbl.kind) + rows = build_node_rows( skeletons=active_skeletons, raw_text_by_page=raw_text_by_page, image_path_by_page=image_path_by_page, + kind_by_page=label_map, tag_by_page=tag_map, filename=filename, + verdict="page", vlm_model=args.vlm_model or os.environ.get("IMAGE_MODEL"), page_assets_by_page=all_assets if all_assets else None, + node_summary_max_pages=page_memory_config.node_summary_max_pages, node_assembly_concurrency=page_memory_config.node_assembly_concurrency, - body_start_by_page=toc_policy.body_start_by_page(), - ) - toc_chunks = build_toc_node_chunks(anatomy=anatomy, filename=filename) - canonical_chunks = merge_chunks_by_first_page(toc_chunks, body_chunks) - from shared.services.chunks.canonical_chunk_builder import chunks_as_json - - chunks = cast( - list[dict[str, Any]], - chunks_as_json(canonical_chunks), ) + chunks = [ + { + "chunk_id": str(row.get("know_id") or ""), + "type": str(row.get("type") or "page"), + "content": str(row.get("content") or ""), + "path": str(row.get("path") or ""), + "metadata": { + "length": int(row.get("length") or 0), + "summary": str(row.get("summary") or ""), + "page_nums": [ + int(part) + for part in str(row.get("page_nums") or "").split(",") + if str(part).strip().isdigit() + ], + "keywords": [ + part.strip() + for part in str(row.get("keywords") or "").split(";") + if part.strip() + ], + **(row.get("extra_metadata") or {}), + }, + } + for row in rows + ] logger.info(f" C7: {len(chunks)} canonical chunks") record_stage( trace_stages, diff --git a/apps/worker/scripts/page_memory/toc_page_policy.py b/apps/worker/scripts/page_memory/toc_page_policy.py deleted file mode 100644 index 4a15f5f50..000000000 --- a/apps/worker/scripts/page_memory/toc_page_policy.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Debug-only TOC page policy for staged page_memory scripts. - -Production page_memory no longer ships this helper; debug stages still need a -single place to read ``toc_result.toc_pages`` and filter processing ranges. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - - -@dataclass(frozen=True) -class TocPagePolicy: - pure_toc_pages: frozenset[int] = field(default_factory=frozenset) - - @classmethod - def from_anatomy(cls, anatomy: Any | None) -> TocPagePolicy: - toc_result = getattr(anatomy, "toc_result", None) if anatomy is not None else None - pages = getattr(toc_result, "toc_pages", None) or [] - pure: set[int] = set() - for page in pages: - try: - pure.add(int(page)) - except (TypeError, ValueError): - continue - return cls(pure_toc_pages=frozenset(pure)) - - def filter_processing_pages(self, pages: list[int]) -> list[int]: - return [page for page in pages if page not in self.pure_toc_pages] - - def body_start_by_page(self) -> dict[int, float]: - return {} diff --git a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py index c78ba9dc5..d2eaf126a 100644 --- a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py +++ b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py @@ -53,6 +53,33 @@ def test_compute_fat_leaf_pages_uses_exclusive_boundaries() -> None: } +def test_compute_fat_leaf_pages_excludes_toc_pages_from_span() -> None: + """Closed range 1-6 with toc=[2,3,4] is only 3 body pages → not fat when min=4.""" + skeletons = [ + SectionSkeleton( + section_path="demo.pdf/Abbreviations", + level=1, + start_page=1, + end_page=6, + title="Abbreviations", + parent_path="demo.pdf", + ), + ] + assert ( + fine_hierarchy.compute_fat_leaf_pages( + skeletons, + min_pages=4, + toc_pages=[2, 3, 4], + ) + == set() + ) + assert fine_hierarchy.compute_fat_leaf_pages( + skeletons, + min_pages=2, + toc_pages=[2, 3, 4], + ) == {1, 5, 6} + + def test_refine_fat_leaf_skeletons_excludes_next_section_start_when_unordered( monkeypatch, ) -> None: diff --git a/apps/worker/tests/contract/test_toc_rehome_contract.py b/apps/worker/tests/contract/test_toc_rehome_contract.py new file mode 100644 index 000000000..92e8264fc --- /dev/null +++ b/apps/worker/tests/contract/test_toc_rehome_contract.py @@ -0,0 +1,355 @@ +"""Contract tests for same-forest TOC monotonic rehome. Synthetic trees only.""" + +from __future__ import annotations + +import os +from unittest.mock import patch + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from app.services.document_agent.manifest import TocResult, ToolContext +from app.services.document_agent.state import ProfileBlackboard +from app.services.document_agent.structure.anchoring_primitives import ( + SkeletonAnchor, + serialize_skeleton_anchor, + serialize_title_node, +) +from app.services.document_agent.structure.hierarchy_locator import ( + TitleMatch, + TitleNode, +) +from app.services.document_agent.structure.toc_anchoring import run_toc_anchoring +from app.services.document_agent.structure.toc_rehome import rehome_forest + + +def _match(title: str, page: int) -> TitleMatch: + return TitleMatch( + page=page, + source="anchored", + matched_line=title, + candidates=[page], + evidence={}, + ) + + +def _overrides(pages: dict[tuple[str, ...], int]) -> dict[tuple[str, ...], TitleMatch]: + return {path: _match(path[-1], page) for path, page in pages.items()} + + +def _titles(nodes: list[TitleNode]) -> list[str]: + return [node.title for node in nodes] + + +def test_same_parent_backjump_reorders_by_own_page() -> None: + nodes = [ + TitleNode(title="A", level=1, printed_page=10), + TitleNode(title="B", level=1, printed_page=20), + TitleNode(title="C", level=1, printed_page=5), + ] + result = rehome_forest( + nodes, + _overrides({("A",): 10, ("B",): 20, ("C",): 5}), + ) + assert _titles(result.nodes) == ["C", "A", "B"] + assert set(result.match_overrides) == {("A",), ("B",), ("C",)} + assert result.match_overrides[("C",)].page == 5 + + +def test_unpaged_shell_transparent_promotes_children_and_drops_shell() -> None: + nodes = [ + TitleNode(title="Preamble", level=1, printed_page=6), + TitleNode(title="References", level=1, printed_page=61), + TitleNode( + title="List of tables", + level=1, + children=[ + TitleNode(title="Table 1", level=2, printed_page=7), + TitleNode(title="Table 2", level=2, printed_page=9), + ], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Preamble",): 6, + ("References",): 61, + ("List of tables", "Table 1"): 7, + ("List of tables", "Table 2"): 9, + } + ), + ) + assert "List of tables" not in _titles(result.nodes) + assert _titles(result.nodes) == ["Preamble", "Table 1", "Table 2", "References"] + assert ("List of tables",) not in result.match_overrides + assert ("List of tables", "Table 1") not in result.match_overrides + assert result.match_overrides[("Table 1",)].page == 7 + assert result.match_overrides[("Table 2",)].page == 9 + assert result.nodes[1].level == 1 + + +def test_monotonic_forest_is_noop() -> None: + nodes = [ + TitleNode(title="A", level=1, printed_page=8), + TitleNode(title="B", level=1, printed_page=14), + TitleNode(title="C", level=1, printed_page=20), + ] + overrides = _overrides({("A",): 8, ("B",): 14, ("C",): 20}) + result = rehome_forest(nodes, overrides) + assert _titles(result.nodes) == ["A", "B", "C"] + assert result.match_overrides == overrides + assert result.events == [] + + +def test_child_level_backjump_does_not_use_parent_cursor() -> None: + nodes = [ + TitleNode( + title="Chapter", + level=1, + printed_page=10, + children=[ + TitleNode(title="Late", level=2, printed_page=30), + TitleNode(title="Early", level=2, printed_page=12), + ], + ), + TitleNode(title="Next", level=1, printed_page=40), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Chapter",): 10, + ("Chapter", "Late"): 30, + ("Chapter", "Early"): 12, + ("Next",): 40, + } + ), + ) + assert _titles(result.nodes) == ["Chapter", "Next"] + assert _titles(result.nodes[0].children) == ["Early", "Late"] + assert ("Chapter", "Early") in result.match_overrides + assert ("Chapter", "Late") in result.match_overrides + + +def test_transparent_shell_still_rehomes_children_of_paged_node() -> None: + """Unpaged shell expands into the parent layer; that paged node's kids still reorder.""" + nodes = [ + TitleNode(title="Chapter", level=1, printed_page=10), + TitleNode( + title="List", + level=1, + children=[ + TitleNode( + title="FrontMatter", + level=2, + printed_page=50, + children=[ + TitleNode(title="Late", level=3, printed_page=80), + TitleNode(title="Early", level=3, printed_page=55), + ], + ), + ], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Chapter",): 10, + ("List", "FrontMatter"): 50, + ("List", "FrontMatter", "Late"): 80, + ("List", "FrontMatter", "Early"): 55, + } + ), + ) + assert _titles(result.nodes) == ["Chapter", "List"] + front = result.nodes[1].children[0] + assert front.title == "FrontMatter" + assert _titles(front.children) == ["Early", "Late"] + assert result.match_overrides[("List", "FrontMatter", "Early")].page == 55 + assert result.match_overrides[("List", "FrontMatter", "Late")].page == 80 + + +def test_promote_paged_node_then_rehome_its_children() -> None: + """Backjump promotion of a paged node still fixes inversion under that node.""" + nodes = [ + TitleNode(title="Chapter", level=1, printed_page=10), + TitleNode( + title="List", + level=1, + children=[ + TitleNode( + title="FrontMatter", + level=2, + printed_page=5, + children=[ + TitleNode(title="Late", level=3, printed_page=8), + TitleNode(title="Early", level=3, printed_page=6), + ], + ), + ], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Chapter",): 10, + ("List", "FrontMatter"): 5, + ("List", "FrontMatter", "Late"): 8, + ("List", "FrontMatter", "Early"): 6, + } + ), + ) + assert "List" not in _titles(result.nodes) + assert _titles(result.nodes) == ["FrontMatter", "Chapter"] + assert _titles(result.nodes[0].children) == ["Early", "Late"] + assert result.nodes[0].level == 1 + assert result.match_overrides[("FrontMatter", "Early")].page == 6 + assert result.match_overrides[("FrontMatter", "Late")].page == 8 + assert ("List",) not in result.match_overrides + assert ("List", "FrontMatter") not in result.match_overrides + + +def _ctx(*, hierarchies: list[dict], page_count: int = 80) -> ToolContext: + ctx = ToolContext( + pdf_path="/tmp/rehome.pdf", + job_id="rehome-test", + blackboard=ProfileBlackboard(page_count=page_count), + trace=None, + settings={}, + ) + ctx.blackboard.toc_hierarchies = hierarchies + ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1]) + ctx.blackboard.page_full_text_cache = { + page: "body" for page in range(1, page_count + 1) + } + return ctx + + +def test_run_toc_anchoring_single_toc_still_runs_rehome() -> None: + hierarchies = [ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Alpha", "level": 1, "page_number": 10}, + {"heading": "Bravo", "level": 1, "page_number": 20}, + {"heading": "Charlie", "level": 1, "page_number": 5}, + ], + } + ] + ctx = _ctx(hierarchies=hierarchies, page_count=30) + + def _fake_anchor(**kwargs): + nodes = [ + TitleNode(title="Alpha", level=1, printed_page=10), + TitleNode(title="Bravo", level=1, printed_page=20), + TitleNode(title="Charlie", level=1, printed_page=5), + ] + anchor = SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides=_overrides( + {("Alpha",): 10, ("Bravo",): 20, ("Charlie",): 5} + ), + null_page_report=[], + bulk_count=3, + ) + return nodes, anchor + + with patch( + "app.services.document_agent.calibration.orchestrator.anchor_hierarchy", + side_effect=_fake_anchor, + ): + run_toc_anchoring(ctx) + + assert ctx.blackboard.skeleton_anchor is not None + assert ctx.blackboard.skeleton_nodes is not None + assert ctx.blackboard.pending_skeleton_anchors == [] + assert [node["title"] for node in ctx.blackboard.skeleton_nodes] == [ + "Charlie", + "Alpha", + "Bravo", + ] + + +def test_run_toc_anchoring_rehome_before_classify_keeps_contained_graft() -> None: + hierarchies = [ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Host", "level": 1, "page_number": 10}, + {"heading": "Tail", "level": 1, "page_number": 40}, + ], + }, + { + "toc_range": [5, 5], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Inner", "level": 1, "page_number": 12}, + ], + }, + ] + ctx = _ctx(hierarchies=hierarchies, page_count=50) + + def _fake_anchor(**kwargs): + nodes = [ + TitleNode(title="Host", level=1, printed_page=10), + TitleNode(title="Tail", level=1, printed_page=40), + ] + anchor = SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides=_overrides({("Host",): 10, ("Tail",): 40}), + null_page_report=[], + bulk_count=2, + ) + return nodes, anchor + + def _fake_pending(**kwargs): + return [ + { + "toc": hierarchies[1], + "nodes": [ + serialize_title_node( + TitleNode(title="Inner", level=1, printed_page=12) + ) + ], + "skeleton_anchor": serialize_skeleton_anchor( + SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides=_overrides({("Inner",): 12}), + null_page_report=[], + bulk_count=1, + ) + ), + } + ] + + with ( + patch( + "app.services.document_agent.calibration.orchestrator.anchor_hierarchy", + side_effect=_fake_anchor, + ), + patch( + "app.services.document_agent.structure.toc_anchoring._calibrate_pending_tocs", + side_effect=_fake_pending, + ), + ): + run_toc_anchoring(ctx) + + record = ctx.blackboard.pending_skeleton_anchors[0] + assert record["relationship"] == "contained" + assert record.get("grafted") is True + host = next( + node for node in ctx.blackboard.skeleton_nodes if node["title"] == "Host" + ) + assert any(child["title"] == "Inner" for child in host["children"]) From ec32a3c3f20138dd2f71ee132854a306bc2dd76d Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 25 Aug 2026 23:33:53 +0800 Subject: [PATCH 2/7] refactor: improve TOC rehome logic and enhance hierarchy resolution - Updated the hierarchy locator to better handle rehome matches and leaf nodes, ensuring accurate page range resolution. - Introduced new functions for managing rehome ranges and structural leaf identification, improving the clarity and efficiency of the code. - Enhanced the TOC anchoring process to support global paged-leaf rehome, allowing for more robust handling of TOC structures. - Refactored tests to validate the new rehome logic and ensure correct behavior in various scenarios. --- .../structure/hierarchy_locator.py | 136 +++++- .../document_agent/structure/toc_anchoring.py | 25 +- .../document_agent/structure/toc_rehome.py | 429 ++++++++++++++---- .../contract/test_toc_rehome_contract.py | 380 ++++++++++++++-- 4 files changed, 834 insertions(+), 136 deletions(-) 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 512689944..99e6e67aa 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -136,11 +136,11 @@ def first_leaf_start_under( parent_titles: tuple[str, ...], match_overrides: dict[tuple[str, ...], TitleMatch], ) -> int | None: - """Min start page among located leaves under *node*; None if none located.""" + """Min start page among structural leaves; rehome attachments do not bound scope.""" min_page: int | None = None for leaf_path, _leaf in iter_leaf_title_nodes([node], parent_titles=parent_titles): match = match_overrides.get(leaf_path) - if match is None: + if match is None or _is_rehome_match(match): continue if min_page is None or match.page < min_page: min_page = match.page @@ -179,7 +179,12 @@ def resolve_hierarchy_page_ranges( match_overrides=match_overrides or {}, resolved=resolved, ) - return resolved + rehome_ranges = _attach_rehome_ranges( + nodes, + match_overrides=match_overrides or {}, + structural_ranges=resolved, + ) + return _ranges_in_tree_order(nodes, [*resolved, *rehome_ranges]) def coverage_by_path( @@ -240,10 +245,19 @@ def _resolve_siblings( match_overrides: dict[tuple[str, ...], TitleMatch], resolved: list[ResolvedHierarchyRange], ) -> None: + boundary_nodes = [ + node + for node in nodes + if not _is_rehome_leaf( + node, + (*parent_titles, node.title), + match_overrides, + ) + ] located: list[tuple[TitleNode, int, TitleMatch | None]] = [] lower_bound = parent_scope.start - for index, node in enumerate(nodes): + for index, node in enumerate(boundary_nodes): path_titles = (*parent_titles, node.title) pages = _allowed_pages_between(lower_bound, parent_scope.end, allowed_pages) match = _locate_match_for_node( @@ -259,9 +273,9 @@ def _resolve_siblings( located.append((node, start_page, match)) if match is not None: lower_bound = start_page - elif index + 1 < len(nodes): + elif index + 1 < len(boundary_nodes): next_match = _find_next_located_sibling( - nodes=nodes, + nodes=boundary_nodes, start_index=index + 1, lower_bound=lower_bound, parent_end=parent_scope.end, @@ -334,6 +348,114 @@ def _resolve_siblings( ) +def _is_rehome_match(match: TitleMatch | None) -> bool: + if match is None: + return False + return bool((match.evidence or {}).get("toc_rehome")) + + +def _is_rehome_leaf( + node: TitleNode, + path_titles: tuple[str, ...], + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> bool: + return not node.children and _is_rehome_match(match_overrides.get(path_titles)) + + +def _rehome_host_range( + *, + structural_ranges: list[ResolvedHierarchyRange], + physical_page: int, +) -> ResolvedHierarchyRange: + candidates = [ + item + for item in structural_ranges + if item.start_page <= physical_page <= item.end_page + ] + if not candidates: + raise ValueError( + "rehome leaf has no resolved host scope " + f"page={physical_page}" + ) + latest_start = max(item.start_page for item in candidates) + return next( + item + for item in reversed(candidates) + if item.start_page == latest_start + ) + + +def _attach_rehome_ranges( + nodes: list[TitleNode], + *, + match_overrides: dict[tuple[str, ...], TitleMatch], + structural_ranges: list[ResolvedHierarchyRange], + parent_titles: tuple[str, ...] = (), +) -> list[ResolvedHierarchyRange]: + attached: list[ResolvedHierarchyRange] = [] + for node in nodes: + path_titles = (*parent_titles, node.title) + if _is_rehome_leaf(node, path_titles, match_overrides): + match = match_overrides.get(path_titles) + if match is None: + raise ValueError(f"rehome leaf override missing path={path_titles!r}") + host_range = _rehome_host_range( + structural_ranges=structural_ranges, + physical_page=match.page, + ) + attached.append( + ResolvedHierarchyRange( + title=node.title, + level=node.level, + start_page=host_range.start_page, + end_page=host_range.end_page, + path_titles=path_titles, + match=match, + evidence={ + **_range_evidence(match), + "status": "rehome_attached", + "skeleton_kind": "rehome_attachment", + "scope_host_path": list(host_range.path_titles), + }, + ) + ) + continue + if node.children: + attached.extend( + _attach_rehome_ranges( + node.children, + match_overrides=match_overrides, + structural_ranges=structural_ranges, + parent_titles=path_titles, + ) + ) + return attached + + +def _ranges_in_tree_order( + nodes: list[TitleNode], + ranges: list[ResolvedHierarchyRange], +) -> list[ResolvedHierarchyRange]: + order = { + path: index + for index, (path, _node) in enumerate(_iter_title_nodes(nodes)) + } + return sorted(ranges, key=lambda item: order[item.path_titles]) + + +def _iter_title_nodes( + nodes: list[TitleNode], + *, + parent_titles: tuple[str, ...] = (), +) -> list[tuple[tuple[str, ...], TitleNode]]: + items: list[tuple[tuple[str, ...], TitleNode]] = [] + for node in nodes: + path_titles = (*parent_titles, node.title) + items.append((path_titles, node)) + items.extend(_iter_title_nodes(node.children, parent_titles=path_titles)) + return items + + def _locate_match_for_node( node: TitleNode, *, @@ -391,7 +513,7 @@ def _infer_start_from_descendant_overrides( min_match: TitleMatch | None = None for leaf_path, _leaf_node in leaves: m = match_overrides.get(leaf_path) - if m is None: + if m is None or _is_rehome_match(m): continue if m.page not in scope_pages: continue diff --git a/apps/worker/app/services/document_agent/structure/toc_anchoring.py b/apps/worker/app/services/document_agent/structure/toc_anchoring.py index dd183d5ef..5e0f7c34a 100644 --- a/apps/worker/app/services/document_agent/structure/toc_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/toc_anchoring.py @@ -56,7 +56,7 @@ def run_toc_anchoring(ctx: ToolContext) -> None: physical overrides (no calibrate VLM). Otherwise keep the extracted TOC tree and run the existing VLM calibration path. - After calibration (primary + pending), each forest runs same-tree monotonic + After calibration (primary + pending), each forest runs global paged-leaf rehome, then pending forests are classified/grafted. Blackboard writes happen only after that sequence. """ @@ -122,9 +122,11 @@ def run_toc_anchoring(ctx: ToolContext) -> None: # scope-local TOC pre-pass: future hook (no-op this period) resolve_nodes, skeleton_anchor = _rehome_calibrated_forest( - resolve_nodes, skeleton_anchor + resolve_nodes, + skeleton_anchor, + toc_pages=toc_pages, ) - _rehome_pending_records(pending_records) + _rehome_pending_records(pending_records, toc_pages=toc_pages) if pending_records: _assign_toc_relationships( @@ -255,7 +257,9 @@ def _write_outline_skeleton( skeleton_anchor = replace(skeleton_anchor, source="pdf_outline") # scope-local TOC pre-pass: future hook (no-op this period) resolve_nodes, skeleton_anchor = _rehome_calibrated_forest( - resolve_nodes, skeleton_anchor + resolve_nodes, + skeleton_anchor, + toc_pages=toc_pages, ) ctx.blackboard.skeleton_anchor = serialize_skeleton_anchor(skeleton_anchor) ctx.blackboard.skeleton_nodes = [ @@ -275,16 +279,24 @@ def _write_outline_skeleton( def _rehome_calibrated_forest( nodes: list[TitleNode], skeleton_anchor: SkeletonAnchor, + *, + toc_pages: list[int] | None = None, ) -> tuple[list[TitleNode], SkeletonAnchor]: from app.services.document_agent.structure.toc_rehome import rehome_skeleton_forest rehomed_nodes, rehomed_anchor, _events = rehome_skeleton_forest( - nodes, skeleton_anchor + nodes, + skeleton_anchor, + toc_pages=toc_pages, ) return rehomed_nodes, rehomed_anchor -def _rehome_pending_records(pending_records: list[dict[str, Any]]) -> None: +def _rehome_pending_records( + pending_records: list[dict[str, Any]], + *, + toc_pages: list[int] | None = None, +) -> None: from app.services.document_agent.structure.toc_rehome import rehome_skeleton_forest for record in pending_records: @@ -302,6 +314,7 @@ def _rehome_pending_records(pending_records: list[dict[str, Any]]) -> None: rehomed_nodes, rehomed_anchor, _events = rehome_skeleton_forest( nodes, deserialize_skeleton_anchor(anchor_raw), + toc_pages=toc_pages, ) record["nodes"] = [serialize_title_node(node) for node in rehomed_nodes] record["skeleton_anchor"] = serialize_skeleton_anchor(rehomed_anchor) diff --git a/apps/worker/app/services/document_agent/structure/toc_rehome.py b/apps/worker/app/services/document_agent/structure/toc_rehome.py index f459aa76c..090294fcc 100644 --- a/apps/worker/app/services/document_agent/structure/toc_rehome.py +++ b/apps/worker/app/services/document_agent/structure/toc_rehome.py @@ -1,4 +1,4 @@ -"""Same-forest TOC rehome: layer-wise own-page monotonic repair before classify.""" +"""Same-forest TOC rehome: global paged-leaf monotonic repair before classify.""" from __future__ import annotations @@ -22,8 +22,11 @@ @dataclass(frozen=True) class RehomePlan: source_path: tuple[str, ...] - dest_parent_path: tuple[str, ...] physical_page: int + segment_index: int + source_toc_index: int + anchor_source_path: tuple[str, ...] + action: str = "moved" @dataclass(frozen=True) @@ -33,6 +36,13 @@ class RehomeResult: events: list[dict[str, Any]] +@dataclass(frozen=True) +class _LeafRef: + path: tuple[str, ...] + page: int + toc_index: int + + def own_page( path: tuple[str, ...], match_overrides: dict[tuple[str, ...], TitleMatch], @@ -47,33 +57,52 @@ def own_page( def rehome_forest( nodes: list[TitleNode], match_overrides: dict[tuple[str, ...], TitleMatch], + *, + toc_pages: list[int] | None = None, ) -> RehomeResult: - """Repair non-monotonic own-page order within one calibrated TOC forest.""" + """Repair TOC-order backjumps among paged leaves within one calibrated forest.""" overrides = dict(match_overrides) - plans = _collect_plans(siblings=list(nodes), parent_path=(), overrides=overrides) + toc_page_set = {int(page) for page in (toc_pages or [])} + plans = _collect_backjump_plans(nodes, overrides, toc_pages=toc_page_set) if not plans: return RehomeResult(nodes=list(nodes), match_overrides=overrides, events=[]) - events: list[dict[str, Any]] = [ - { - "action": "plan", - "source_path": list(plan.source_path), - "dest_parent_path": list(plan.dest_parent_path), - "physical_page": plan.physical_page, - } - for plan in plans - ] - logger.info("{} plans={}", _LOG_PREFIX, len(plans)) - new_nodes, overrides = _apply_plans(list(nodes), overrides, plans) - return RehomeResult(nodes=new_nodes, match_overrides=overrides, events=events) + prune_plans = [plan for plan in plans if plan.action == "pruned"] + move_plans = [plan for plan in plans if plan.action == "moved"] + logger.info( + "{} plans={} prunes={} moves={}", + _LOG_PREFIX, + len(plans), + len(prune_plans), + len(move_plans), + ) + working = list(nodes) + events: list[dict[str, Any]] = [] + if prune_plans: + working, overrides, prune_events = _apply_prunes( + working, overrides, prune_plans + ) + events.extend(prune_events) + if move_plans: + working, overrides, move_events = _apply_plans( + working, overrides, move_plans + ) + events.extend(move_events) + return RehomeResult(nodes=working, match_overrides=overrides, events=events) def rehome_skeleton_forest( nodes: list[TitleNode], anchor: SkeletonAnchor, + *, + toc_pages: list[int] | None = None, ) -> tuple[list[TitleNode], SkeletonAnchor, list[dict[str, Any]]]: """Run ``rehome_forest`` and return an updated ``SkeletonAnchor``.""" - result = rehome_forest(nodes, anchor.match_overrides) + result = rehome_forest( + nodes, + anchor.match_overrides, + toc_pages=toc_pages, + ) return ( result.nodes, replace(anchor, match_overrides=result.match_overrides), @@ -81,111 +110,340 @@ def rehome_skeleton_forest( ) -def _collect_plans( +def _collect_paged_leaves( + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], TitleMatch], + parent_path: tuple[str, ...] = (), *, - siblings: list[TitleNode], - parent_path: tuple[str, ...], + toc_index_start: int = 0, +) -> list[_LeafRef]: + """TOC-order list of paged leaves (nodes with no children and own page).""" + leaves: list[_LeafRef] = [] + toc_index = toc_index_start + for node in nodes: + path = (*parent_path, node.title) + if node.children: + child_leaves = _collect_paged_leaves( + list(node.children), + overrides, + path, + toc_index_start=toc_index, + ) + leaves.extend(child_leaves) + toc_index += len(child_leaves) + continue + page = own_page(path, overrides) + if page is None: + continue + leaves.append(_LeafRef(path=path, page=page, toc_index=toc_index)) + toc_index += 1 + return leaves + + +def _split_monotonic_segments(leaves: list[_LeafRef]) -> list[list[_LeafRef]]: + """Split TOC-order leaves at every physical-page backjump.""" + segments: list[list[_LeafRef]] = [] + current: list[_LeafRef] = [] + for leaf in leaves: + if current and leaf.page < current[-1].page: + segments.append(current) + current = [] + current.append(leaf) + if current: + segments.append(current) + return segments + + +def _nearest_leaf_in_segment( + *, + physical_page: int, + segment: list[_LeafRef], +) -> _LeafRef | None: + """Nearest leaf with ``page <= physical_page`` inside one fixed segment.""" + candidates = [ + leaf + for leaf in segment + if leaf.page <= physical_page + ] + if not candidates: + return None + return max(candidates, key=lambda leaf: (leaf.page, leaf.toc_index)) + + +def _first_segment_duplicate( + leaf: _LeafRef, + first_segment: list[_LeafRef], +) -> _LeafRef | None: + """First-segment leaf with the same hierarchy path and physical page, if any.""" + for candidate in first_segment: + if candidate.path == leaf.path and candidate.page == leaf.page: + return candidate + return None + + +def _collect_backjump_plans( + nodes: list[TitleNode], overrides: dict[tuple[str, ...], TitleMatch], + *, + toc_pages: set[int], ) -> list[RehomePlan]: + """Plan every post-break segment against the first monotonic segment only.""" + segments = _split_monotonic_segments(_collect_paged_leaves(nodes, overrides)) + if len(segments) < 2: + return [] + first_segment = segments[0] plans: list[RehomePlan] = [] - effective = _expand_effective( - siblings=siblings, - parent_path=parent_path, - overrides=overrides, - ) - cursor: int | None = None - for _node, source_path, page in effective: - if cursor is not None and page < cursor: - plans.append( - RehomePlan( - source_path=source_path, - dest_parent_path=parent_path, - physical_page=page, + for segment_index in range(1, len(segments)): + for leaf in segments[segment_index]: + duplicate = _first_segment_duplicate(leaf, first_segment) + if duplicate is not None: + logger.info( + "{} prune same_path_page path={} page={} segment={}", + _LOG_PREFIX, + leaf.path, + leaf.page, + segment_index, + ) + plans.append( + RehomePlan( + source_path=leaf.path, + physical_page=leaf.page, + segment_index=segment_index, + source_toc_index=leaf.toc_index, + anchor_source_path=duplicate.path, + action="pruned", + ) ) + continue + if leaf.page in toc_pages: + logger.info( + "{} skip toc_page path={} page={} segment={}", + _LOG_PREFIX, + leaf.path, + leaf.page, + segment_index, + ) + continue + anchor = _nearest_leaf_in_segment( + physical_page=leaf.page, + segment=first_segment, ) - continue - if cursor is None or page > cursor: - cursor = page - - # Drill into every own-paged effective node (includes transparent promotions). - # Unpaged shells are not parents here — their children already appear in effective. - for node, source_path, _page in effective: - if node.children: - plans.extend( - _collect_plans( - siblings=list(node.children), - parent_path=source_path, - overrides=overrides, + if anchor is None: + continue + plans.append( + RehomePlan( + source_path=leaf.path, + physical_page=leaf.page, + segment_index=segment_index, + source_toc_index=leaf.toc_index, + anchor_source_path=anchor.path, + action="moved", ) ) return plans -def _expand_effective( - *, - siblings: list[TitleNode], - parent_path: tuple[str, ...], +def _apply_prunes( + nodes: list[TitleNode], overrides: dict[tuple[str, ...], TitleMatch], -) -> list[tuple[TitleNode, tuple[str, ...], int]]: - """Build this level's ordered sequence of own-paged nodes (unpaged non-leaves transparent).""" - effective: list[tuple[TitleNode, tuple[str, ...], int]] = [] - for sibling in siblings: - path = (*parent_path, sibling.title) - page = own_page(path, overrides) - if page is not None: - effective.append((sibling, path, page)) - continue - if sibling.children: - effective.extend( - _expand_effective( - siblings=list(sibling.children), - parent_path=path, - overrides=overrides, - ) + plans: list[RehomePlan], +) -> tuple[list[TitleNode], dict[tuple[str, ...], TitleMatch], list[dict[str, Any]]]: + """Drop post-first-segment duplicates (same path + page as a first-segment leaf).""" + working = list(nodes) + events: list[dict[str, Any]] = [] + for plan in plans: + leaves = _collect_paged_leaves(working, overrides) + matches = [ + leaf + for leaf in leaves + if leaf.path == plan.source_path and leaf.page == plan.physical_page + ] + if len(matches) < 2: + raise ValueError( + f"{_LOG_PREFIX} prune miss path={plan.source_path!r} " + f"page={plan.physical_page} remaining={len(matches)}" ) - return effective + target = matches[-1] + working, detached, detached_path = _detach_by_toc_index( + working, + overrides, + target_toc_index=target.toc_index, + ) + if detached is None or detached_path is None: + raise ValueError( + f"{_LOG_PREFIX} prune detach miss toc_index={target.toc_index}" + ) + remaining = [ + leaf + for leaf in _collect_paged_leaves(working, overrides) + if leaf.path == detached_path + ] + if not remaining: + overrides.pop(detached_path, None) + working, overrides = _drop_empty_ancestors( + working, + overrides, + detached_path[:-1], + ) + events.append( + { + "action": "pruned", + "source_path": list(detached_path), + "physical_page": plan.physical_page, + "segment_index": plan.segment_index, + "anchor_path": list(plan.anchor_source_path), + } + ) + logger.info( + "{} pruned {} page={} segment={}", + _LOG_PREFIX, + detached_path, + plan.physical_page, + plan.segment_index, + ) + return working, overrides, events def _apply_plans( nodes: list[TitleNode], overrides: dict[tuple[str, ...], TitleMatch], plans: list[RehomePlan], -) -> tuple[list[TitleNode], dict[tuple[str, ...], TitleMatch]]: - ordered = sorted( - plans, - key=lambda plan: (-len(plan.source_path), plan.physical_page, plan.source_path), - ) +) -> tuple[list[TitleNode], dict[tuple[str, ...], TitleMatch], list[dict[str, Any]]]: working = list(nodes) - for plan in ordered: - working, detached = _detach_node(working, plan.source_path) + events: list[dict[str, Any]] = [] + current_paths = { + leaf.path: leaf.path for leaf in _collect_paged_leaves(working, overrides) + } + for plan in plans: + source_path = current_paths[plan.source_path] + anchor_path = current_paths[plan.anchor_source_path] + dest_parent = anchor_path[:-1] + if dest_parent and node_at_path(working, dest_parent) is None: + raise ValueError( + f"{_LOG_PREFIX} dest parent missing path={dest_parent!r}" + ) + + working, detached = _detach_node(working, source_path) if detached is None: raise ValueError( - f"{_LOG_PREFIX} detach miss source_path={plan.source_path!r}" + f"{_LOG_PREFIX} detach miss source_path={source_path!r}" ) - target_level = len(plan.dest_parent_path) + 1 + target_level = len(dest_parent) + 1 detached = rebase_levels(detached, target_level - detached.level) - new_path = (*plan.dest_parent_path, detached.title) + new_path = (*dest_parent, detached.title) working = _insert_under_parent( nodes=working, - parent_path=plan.dest_parent_path, + parent_path=dest_parent, new_node=detached, start=plan.physical_page, overrides=overrides, ) - _move_override_prefix(overrides, plan.source_path, new_path) - working, overrides = _drop_empty_ancestors( - working, + if new_path != source_path: + _move_override_prefix(overrides, source_path, new_path) + working, overrides = _drop_empty_ancestors( + working, + overrides, + source_path[:-1], + ) + current_paths[plan.source_path] = new_path + _mark_rehome_override( overrides, - plan.source_path[:-1], + new_path, + segment_index=plan.segment_index, + ) + events.append( + { + "action": "moved", + "source_path": list(source_path), + "dest_parent_path": list(dest_parent), + "anchor_path": list(anchor_path), + "new_path": list(new_path), + "physical_page": plan.physical_page, + "segment_index": plan.segment_index, + } ) logger.info( - "{} moved {} -> {} page={}", + "{} moved {} -> {} after {} page={} segment={}", _LOG_PREFIX, - plan.source_path, + source_path, new_path, + anchor_path, plan.physical_page, + plan.segment_index, ) - return working, overrides + return working, overrides, events + + +def _detach_by_toc_index( + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], TitleMatch], + *, + target_toc_index: int, + parent_path: tuple[str, ...] = (), + toc_index_start: int = 0, +) -> tuple[list[TitleNode], TitleNode | None, tuple[str, ...] | None]: + """Detach the paged leaf whose TOC-order index matches ``target_toc_index``.""" + toc_index = toc_index_start + result: list[TitleNode] = [] + detached: TitleNode | None = None + detached_path: tuple[str, ...] | None = None + for node in nodes: + if detached is not None: + result.append(node) + continue + path = (*parent_path, node.title) + if node.children: + new_children, child_detached, child_path = _detach_by_toc_index( + list(node.children), + overrides, + target_toc_index=target_toc_index, + parent_path=path, + toc_index_start=toc_index, + ) + child_leaves = _collect_paged_leaves( + list(node.children), + overrides, + path, + toc_index_start=toc_index, + ) + toc_index += len(child_leaves) + if child_detached is not None: + detached = child_detached + detached_path = child_path + result.append(replace(node, children=new_children)) + else: + result.append(node) + continue + page = own_page(path, overrides) + if page is None: + result.append(node) + continue + if toc_index == target_toc_index: + detached = node + detached_path = path + toc_index += 1 + continue + result.append(node) + toc_index += 1 + return result, detached, detached_path + + +def _mark_rehome_override( + overrides: dict[tuple[str, ...], TitleMatch], + path: tuple[str, ...], + *, + segment_index: int, +) -> None: + match = overrides.get(path) + if match is None: + raise ValueError(f"{_LOG_PREFIX} moved override missing path={path!r}") + overrides[path] = replace( + match, + evidence={ + **dict(match.evidence or {}), + "toc_rehome": {"segment_index": segment_index}, + }, + ) def _detach_node( @@ -225,7 +483,6 @@ def _drop_empty_ancestors( if not ancestor_path: return nodes, overrides working = nodes - # Drop from deepest ancestor upward. for depth in range(len(ancestor_path), 0, -1): path = ancestor_path[:depth] node = node_at_path(working, path) diff --git a/apps/worker/tests/contract/test_toc_rehome_contract.py b/apps/worker/tests/contract/test_toc_rehome_contract.py index 92e8264fc..77b6ec3df 100644 --- a/apps/worker/tests/contract/test_toc_rehome_contract.py +++ b/apps/worker/tests/contract/test_toc_rehome_contract.py @@ -1,4 +1,4 @@ -"""Contract tests for same-forest TOC monotonic rehome. Synthetic trees only.""" +"""Contract tests for same-forest TOC paged-leaf rehome. Synthetic trees only.""" from __future__ import annotations @@ -22,6 +22,7 @@ from app.services.document_agent.structure.hierarchy_locator import ( TitleMatch, TitleNode, + resolve_hierarchy_page_ranges, ) from app.services.document_agent.structure.toc_anchoring import run_toc_anchoring from app.services.document_agent.structure.toc_rehome import rehome_forest @@ -45,7 +46,24 @@ def _titles(nodes: list[TitleNode]) -> list[str]: return [node.title for node in nodes] -def test_same_parent_backjump_reorders_by_own_page() -> None: +def test_same_parent_backjump_inserts_after_nearest_prev_leaf() -> None: + nodes = [ + TitleNode(title="A", level=1, printed_page=10), + TitleNode(title="B", level=1, printed_page=20), + TitleNode(title="C", level=1, printed_page=15), + ] + result = rehome_forest( + nodes, + _overrides({("A",): 10, ("B",): 20, ("C",): 15}), + ) + # C@15 nearest prev leaf A@10 → after A under root. + assert _titles(result.nodes) == ["A", "C", "B"] + assert result.match_overrides[("C",)].page == 15 + assert not result.nodes[1].children + + +def test_no_prev_leaf_by_page_is_noop_for_that_backjump() -> None: + """Backjump with no TOC-earlier leaf page<=P is not moved (no fallback).""" nodes = [ TitleNode(title="A", level=1, printed_page=10), TitleNode(title="B", level=1, printed_page=20), @@ -55,12 +73,34 @@ def test_same_parent_backjump_reorders_by_own_page() -> None: nodes, _overrides({("A",): 10, ("B",): 20, ("C",): 5}), ) - assert _titles(result.nodes) == ["C", "A", "B"] - assert set(result.match_overrides) == {("A",), ("B",), ("C",)} - assert result.match_overrides[("C",)].page == 5 + assert _titles(result.nodes) == ["A", "B", "C"] + assert result.events == [] + + +def test_equal_page_backjump_inserts_after_front_same_page_leaves() -> None: + """Rehome @116 with one or more front @116 leaves → after that same-page group.""" + nodes = [ + TitleNode(title="A116", level=1, printed_page=116), + TitleNode(title="B116", level=1, printed_page=116), + TitleNode(title="Late", level=1, printed_page=200), + TitleNode(title="C116", level=1, printed_page=116), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("A116",): 116, + ("B116",): 116, + ("Late",): 200, + ("C116",): 116, + } + ), + ) + assert _titles(result.nodes) == ["A116", "B116", "C116", "Late"] + assert result.events[0]["anchor_path"] == ["B116"] -def test_unpaged_shell_transparent_promotes_children_and_drops_shell() -> None: +def test_unpaged_shell_leaf_inserts_after_nearest_prev_and_drops_shell() -> None: nodes = [ TitleNode(title="Preamble", level=1, printed_page=6), TitleNode(title="References", level=1, printed_page=61), @@ -85,12 +125,45 @@ def test_unpaged_shell_transparent_promotes_children_and_drops_shell() -> None: ), ) assert "List of tables" not in _titles(result.nodes) + # Table1 after Preamble@6; Table2 after Table1@7 — all stay leaves. assert _titles(result.nodes) == ["Preamble", "Table 1", "Table 2", "References"] - assert ("List of tables",) not in result.match_overrides - assert ("List of tables", "Table 1") not in result.match_overrides + assert not result.nodes[1].children + assert not result.nodes[2].children assert result.match_overrides[("Table 1",)].page == 7 assert result.match_overrides[("Table 2",)].page == 9 - assert result.nodes[1].level == 1 + + +def test_equal_page_points_at_same_first_segment_group() -> None: + """Both @117 leaves anchor only into the first monotonic segment.""" + nodes = [ + TitleNode(title="Prev", level=1, printed_page=116), + TitleNode(title="Late", level=1, printed_page=200), + TitleNode( + title="Shell", + level=1, + children=[ + TitleNode(title="T1", level=2, printed_page=117), + TitleNode(title="T2", level=2, printed_page=117), + ], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Prev",): 116, + ("Late",): 200, + ("Shell", "T1"): 117, + ("Shell", "T2"): 117, + } + ), + ) + assert "Shell" not in _titles(result.nodes) + assert _titles(result.nodes) == ["Prev", "T1", "T2", "Late"] + assert [event["anchor_path"] for event in result.events] == [ + ["Prev"], + ["Prev"], + ] def test_monotonic_forest_is_noop() -> None: @@ -106,13 +179,14 @@ def test_monotonic_forest_is_noop() -> None: assert result.events == [] -def test_child_level_backjump_does_not_use_parent_cursor() -> None: +def test_every_leaf_in_post_break_segment_attempts_first_segment() -> None: nodes = [ TitleNode( title="Chapter", level=1, printed_page=10, children=[ + TitleNode(title="Intro", level=2, printed_page=11), TitleNode(title="Late", level=2, printed_page=30), TitleNode(title="Early", level=2, printed_page=12), ], @@ -124,20 +198,259 @@ def test_child_level_backjump_does_not_use_parent_cursor() -> None: _overrides( { ("Chapter",): 10, + ("Chapter", "Intro"): 11, ("Chapter", "Late"): 30, ("Chapter", "Early"): 12, ("Next",): 40, } ), ) - assert _titles(result.nodes) == ["Chapter", "Next"] - assert _titles(result.nodes[0].children) == ["Early", "Late"] + assert _titles(result.nodes) == ["Chapter"] + assert _titles(result.nodes[0].children) == ["Intro", "Early", "Late", "Next"] assert ("Chapter", "Early") in result.match_overrides - assert ("Chapter", "Late") in result.match_overrides + assert ("Chapter", "Next") in result.match_overrides -def test_transparent_shell_still_rehomes_children_of_paged_node() -> None: - """Unpaged shell expands into the parent layer; that paged node's kids still reorder.""" +def test_each_later_segment_anchors_only_to_first_segment() -> None: + nodes = [ + TitleNode(title="A", level=1, printed_page=10), + TitleNode(title="B", level=1, printed_page=20), + TitleNode( + title="Shell 1", + level=1, + children=[ + TitleNode(title="C", level=2, printed_page=15), + TitleNode(title="D", level=2, printed_page=17), + ], + ), + TitleNode( + title="Shell 2", + level=1, + children=[ + TitleNode(title="E", level=2, printed_page=12), + TitleNode(title="F", level=2, printed_page=16), + ], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("A",): 10, + ("B",): 20, + ("Shell 1", "C"): 15, + ("Shell 1", "D"): 17, + ("Shell 2", "E"): 12, + ("Shell 2", "F"): 16, + } + ), + ) + + assert "Shell 1" not in _titles(result.nodes) + assert "Shell 2" not in _titles(result.nodes) + assert _titles(result.nodes) == ["A", "E", "C", "F", "D", "B"] + event_by_title = { + event["new_path"][-1]: event for event in result.events + } + assert event_by_title["C"]["anchor_path"] == ["A"] + assert event_by_title["D"]["anchor_path"] == ["A"] + assert event_by_title["E"]["anchor_path"] == ["A"] + assert event_by_title["F"]["anchor_path"] == ["A"] + + +def test_prune_duplicate_when_same_path_and_same_page_as_first_segment() -> None: + nodes = [ + TitleNode(title="Abbreviations", level=1, printed_page=4), + TitleNode(title="Body", level=1, printed_page=10), + TitleNode(title="Late", level=1, printed_page=50), + TitleNode(title="Abbreviations", level=1, printed_page=4), + ] + overrides = _overrides( + { + ("Abbreviations",): 4, + ("Body",): 10, + ("Late",): 50, + } + ) + result = rehome_forest(nodes, overrides) + assert _titles(result.nodes) == ["Abbreviations", "Body", "Late"] + assert len(result.events) == 1 + assert result.events[0]["action"] == "pruned" + assert result.events[0]["source_path"] == ["Abbreviations"] + assert result.match_overrides[("Abbreviations",)].page == 4 + assert "toc_rehome" not in (result.match_overrides[("Abbreviations",)].evidence or {}) + + +def test_prune_same_path_page_even_when_page_is_toc_page() -> None: + """same_path_page is judged before toc_page skip and still prunes.""" + nodes = [ + TitleNode(title="Abbreviations", level=1, printed_page=4), + TitleNode(title="Body", level=1, printed_page=10), + TitleNode(title="Late", level=1, printed_page=50), + TitleNode(title="Abbreviations", level=1, printed_page=4), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Abbreviations",): 4, + ("Body",): 10, + ("Late",): 50, + } + ), + toc_pages=[4], + ) + assert _titles(result.nodes) == ["Abbreviations", "Body", "Late"] + assert result.events[0]["action"] == "pruned" + assert "toc_rehome" not in (result.match_overrides[("Abbreviations",)].evidence or {}) + + +def test_skip_rehome_when_physical_page_is_toc_page() -> None: + nodes = [ + TitleNode(title="Lead", level=1, printed_page=5), + TitleNode(title="Body", level=1, printed_page=20), + TitleNode( + title="Shell", + level=1, + children=[TitleNode(title="Early", level=2, printed_page=3)], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Lead",): 5, + ("Body",): 20, + ("Shell", "Early"): 3, + } + ), + toc_pages=[2, 3, 4], + ) + assert result.events == [] + assert "Shell" in _titles(result.nodes) + assert _titles(result.nodes[2].children) == ["Early"] + assert "toc_rehome" not in ( + result.match_overrides[("Shell", "Early")].evidence or {} + ) + + +def test_rehome_leaves_attach_to_resolved_scope_without_cutting_boundaries() -> None: + nodes = [ + TitleNode(title="Lead", level=1, printed_page=21), + TitleNode( + title="Chapter 7", + level=1, + printed_page=22, + children=[ + TitleNode(title="7.1", level=2, printed_page=24), + TitleNode(title="7.2", level=2, printed_page=26), + ], + ), + TitleNode(title="Recommendation", level=1, printed_page=22), + TitleNode(title="Table 7", level=1, printed_page=23), + TitleNode(title="Chapter 8", level=1, printed_page=27), + ] + overrides = _overrides( + { + ("Lead",): 21, + ("Chapter 7",): 22, + ("Chapter 7", "7.1"): 24, + ("Chapter 7", "7.2"): 26, + ("Recommendation",): 22, + ("Table 7",): 23, + ("Chapter 8",): 27, + } + ) + for path in (("Recommendation",), ("Table 7",)): + overrides[path] = TitleMatch( + page=overrides[path].page, + source=overrides[path].source, + matched_line=overrides[path].matched_line, + candidates=overrides[path].candidates, + evidence={"toc_rehome": {"segment_index": 1}}, + ) + + ranges = resolve_hierarchy_page_ranges( + nodes, + page_count=30, + match_overrides=overrides, + ) + by_title = {item.title: item for item in ranges} + + assert by_title["7.1"].match is not None + assert by_title["7.1"].start_page == 24 + assert by_title["7.1"].evidence["source"] == "anchored" + assert by_title["7.2"].start_page == 26 + assert (by_title["Recommendation"].start_page, by_title["Recommendation"].end_page) == ( + 22, + 24, + ) + assert (by_title["Table 7"].start_page, by_title["Table 7"].end_page) == ( + 22, + 24, + ) + assert by_title["Recommendation"].evidence["status"] == "rehome_attached" + assert by_title["Table 7"].evidence["status"] == "rehome_attached" + + +def test_rehome_leaf_attaches_by_physical_page_after_structural_resolve() -> None: + nodes = [ + TitleNode( + title="Chapter 3", + level=1, + printed_page=7, + children=[ + TitleNode(title="3.1", level=2, printed_page=7), + TitleNode(title="3.2", level=2, printed_page=8), + TitleNode(title="Table 3", level=2, printed_page=9), + ], + ), + TitleNode(title="Chapter 4", level=1, printed_page=8), + TitleNode(title="Chapter 5", level=1, printed_page=12), + ] + overrides = _overrides( + { + ("Chapter 3",): 7, + ("Chapter 3", "3.1"): 7, + ("Chapter 3", "3.2"): 8, + ("Chapter 3", "Table 3"): 9, + ("Chapter 4",): 8, + ("Chapter 5",): 12, + } + ) + table_path = ("Chapter 3", "Table 3") + table_match = overrides[table_path] + overrides[table_path] = TitleMatch( + page=table_match.page, + source=table_match.source, + matched_line=table_match.matched_line, + candidates=table_match.candidates, + evidence={"toc_rehome": {"segment_index": 1}}, + ) + + ranges = resolve_hierarchy_page_ranges( + nodes, + page_count=15, + match_overrides=overrides, + ) + by_title = {item.title: item for item in ranges} + + assert (by_title["3.2"].start_page, by_title["3.2"].end_page) == (8, 8) + assert (by_title["Table 3"].start_page, by_title["Table 3"].end_page) == ( + 8, + 12, + ) + assert by_title["Table 3"].evidence["scope_host_path"] == ["Chapter 4"] + assert [item.title for item in ranges] == [ + "3.1", + "3.2", + "Table 3", + "Chapter 4", + "Chapter 5", + ] + + +def test_nested_leaf_inserts_after_nearest_root_leaf() -> None: nodes = [ TitleNode(title="Chapter", level=1, printed_page=10), TitleNode( @@ -167,16 +480,16 @@ def test_transparent_shell_still_rehomes_children_of_paged_node() -> None: } ), ) - assert _titles(result.nodes) == ["Chapter", "List"] + # Early@55 after Chapter@10 at root; List (parent shell with Late) stays. + assert _titles(result.nodes) == ["Chapter", "List", "Early"] + assert not result.nodes[2].children front = result.nodes[1].children[0] assert front.title == "FrontMatter" - assert _titles(front.children) == ["Early", "Late"] - assert result.match_overrides[("List", "FrontMatter", "Early")].page == 55 - assert result.match_overrides[("List", "FrontMatter", "Late")].page == 80 + assert _titles(front.children) == ["Late"] + assert result.match_overrides[("Early",)].page == 55 -def test_promote_paged_node_then_rehome_its_children() -> None: - """Backjump promotion of a paged node still fixes inversion under that node.""" +def test_non_leaf_not_rehomed_even_if_page_looks_early() -> None: nodes = [ TitleNode(title="Chapter", level=1, printed_page=10), TitleNode( @@ -188,8 +501,7 @@ def test_promote_paged_node_then_rehome_its_children() -> None: level=2, printed_page=5, children=[ - TitleNode(title="Late", level=3, printed_page=8), - TitleNode(title="Early", level=3, printed_page=6), + TitleNode(title="Child", level=3, printed_page=50), ], ), ], @@ -201,19 +513,13 @@ def test_promote_paged_node_then_rehome_its_children() -> None: { ("Chapter",): 10, ("List", "FrontMatter"): 5, - ("List", "FrontMatter", "Late"): 8, - ("List", "FrontMatter", "Early"): 6, + ("List", "FrontMatter", "Child"): 50, } ), ) - assert "List" not in _titles(result.nodes) - assert _titles(result.nodes) == ["FrontMatter", "Chapter"] - assert _titles(result.nodes[0].children) == ["Early", "Late"] - assert result.nodes[0].level == 1 - assert result.match_overrides[("FrontMatter", "Early")].page == 6 - assert result.match_overrides[("FrontMatter", "Late")].page == 8 - assert ("List",) not in result.match_overrides - assert ("List", "FrontMatter") not in result.match_overrides + assert _titles(result.nodes) == ["Chapter", "List"] + assert result.nodes[1].children[0].title == "FrontMatter" + assert result.events == [] def _ctx(*, hierarchies: list[dict], page_count: int = 80) -> ToolContext: @@ -240,7 +546,7 @@ def test_run_toc_anchoring_single_toc_still_runs_rehome() -> None: "toc_with_level": [ {"heading": "Alpha", "level": 1, "page_number": 10}, {"heading": "Bravo", "level": 1, "page_number": 20}, - {"heading": "Charlie", "level": 1, "page_number": 5}, + {"heading": "Charlie", "level": 1, "page_number": 15}, ], } ] @@ -250,13 +556,13 @@ def _fake_anchor(**kwargs): nodes = [ TitleNode(title="Alpha", level=1, printed_page=10), TitleNode(title="Bravo", level=1, printed_page=20), - TitleNode(title="Charlie", level=1, printed_page=5), + TitleNode(title="Charlie", level=1, printed_page=15), ] anchor = SkeletonAnchor( offset=0, offset_status="ok", match_overrides=_overrides( - {("Alpha",): 10, ("Bravo",): 20, ("Charlie",): 5} + {("Alpha",): 10, ("Bravo",): 20, ("Charlie",): 15} ), null_page_report=[], bulk_count=3, @@ -273,8 +579,8 @@ def _fake_anchor(**kwargs): assert ctx.blackboard.skeleton_nodes is not None assert ctx.blackboard.pending_skeleton_anchors == [] assert [node["title"] for node in ctx.blackboard.skeleton_nodes] == [ - "Charlie", "Alpha", + "Charlie", "Bravo", ] From a846893fea8b7f681f67322806e2962edb11c959 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 26 Aug 2026 10:50:44 +0800 Subject: [PATCH 3/7] refactor: enhance page rendering and inspection logic - Updated the rendering logic to support a new naming convention for PNG files, allowing for better organization based on the presence of a prefix. - Refactored the inspection process to render pages serially before conducting concurrent inspections, improving efficiency and clarity in the workflow. - Adjusted the handling of rendered pages in the inspection function to ensure proper coverage and error handling. - Enhanced tests to validate the new rendering and inspection behavior, ensuring robustness in various scenarios. --- .../document_agent/calibration/scan.py | 66 ++++++++++++++++--- .../document_agent/tools/inspect_pages.py | 46 +++++++++---- .../app/services/document_agent/visual.py | 8 ++- .../formats/pdf/pymupdf_subprocess.py | 25 ++++--- .../services/page_memory/memory_service.py | 3 + .../services/page_memory/node_assembler.py | 21 ++++-- .../app/services/page_memory/page_renderer.py | 23 +++++-- .../debug_pm_stage3_scope_fine_hierarchy.py | 2 + .../page_memory/debug_pm_stage4_assets.py | 1 + .../debug_pm_stage5_tagging_finalize.py | 2 + .../test_calibration_scan_contract.py | 60 +++++++++++++++++ ...est_page_memory_node_assembler_contract.py | 14 +++- ...emory_page_renderer_scope_path_contract.py | 64 ++++++++++++++++++ .../shared/services/ai/summary/engine.py | 6 +- 14 files changed, 292 insertions(+), 49 deletions(-) create mode 100644 apps/worker/tests/contract/test_page_memory_page_renderer_scope_path_contract.py diff --git a/apps/worker/app/services/document_agent/calibration/scan.py b/apps/worker/app/services/document_agent/calibration/scan.py index 7b773aa18..9050ab1d7 100644 --- a/apps/worker/app/services/document_agent/calibration/scan.py +++ b/apps/worker/app/services/document_agent/calibration/scan.py @@ -5,8 +5,10 @@ window, feeding each round's cursor into the next one, so a miss never re-opens pages that were already inspected. -Each round still covers ``window_schedule[i]`` pages, but pages are inspected -one-at-a-time concurrently (never batched into a single VLM call). +Each round still covers ``window_schedule[i]`` pages. PDF→PNG for the whole +window is rendered in one serial child-process call (same discipline as TOC +extract: never drive the gevent PyMuPDF pool from a ThreadPool). VLM inspect +then runs one page per call, concurrently. """ from __future__ import annotations @@ -25,6 +27,7 @@ ) from app.services.document_agent.manifest import ToolContext from app.services.document_agent.tools.inspect_pages import inspect_pages +from app.services.document_agent.visual import render_pages DEFAULT_WINDOW_SCHEDULE: tuple[int, ...] = (2, 4, 6, 10) @@ -111,6 +114,7 @@ def _inspect_one_page( ctx: ToolContext, title: str, page: int, + rendered_page: dict[str, Any], ) -> PageInspectResult: result = inspect_pages( ctx, @@ -122,6 +126,7 @@ def _inspect_one_page( "folder_name": "calibration_scan", "prefix": "scan", "usage_task": "calibration.scan_title_forward", + "rendered_pages": [rendered_page], }, ) if result.status != "ok": @@ -142,14 +147,58 @@ def _inspect_pages_concurrent( title: str, pages: list[int], ) -> list[PageInspectResult]: - """Inspect each page alone; keep page order in the returned list.""" + """Serial window render, then concurrent single-page VLM inspect.""" + rendered = render_pages( + ctx, + pages, + folder_name="calibration_scan", + prefix="scan", + timeout=120, + ) + rendered_by_page = { + int(item["page"]): { + "page": int(item["page"]), + "png_path": str(item["png_path"]), + } + for item in rendered + if item.get("page") is not None and item.get("png_path") + } + missing = [page for page in pages if page not in rendered_by_page] + if missing: + return [ + PageInspectResult( + page=page, + found=False, + error=( + f"render failed for pages={missing}" + if page in missing + else "render incomplete" + ), + ) + for page in pages + ] + if len(pages) == 1: - return [_inspect_one_page(ctx=ctx, title=title, page=pages[0])] + page = pages[0] + return [ + _inspect_one_page( + ctx=ctx, + title=title, + page=page, + rendered_page=rendered_by_page[page], + ) + ] by_page: dict[int, PageInspectResult] = {} with ThreadPoolExecutor(max_workers=len(pages)) as pool: futures = { - pool.submit(_inspect_one_page, ctx=ctx, title=title, page=page): page + pool.submit( + _inspect_one_page, + ctx=ctx, + title=title, + page=page, + rendered_page=rendered_by_page[page], + ): page for page in pages } for future in as_completed(futures): @@ -174,9 +223,10 @@ def scan_title_forward( """Scan forward from ``start_page`` until the title is found or rounds run out. Each round covers ``window_schedule[i]`` consecutive pages starting at the - cursor left by the previous round. Pages inside a round are inspected - concurrently, one page per VLM call; the earliest true page wins. A page - error without a hit is logged and the scan continues to the next window. + cursor left by the previous round. The window is rendered once serially, + then each page is inspected via VLM concurrently; the earliest true page + wins. A page error without a hit is logged and the scan continues to the + next window. """ scanned: list[int] = [] rounds: list[ScanRound] = [] diff --git a/apps/worker/app/services/document_agent/tools/inspect_pages.py b/apps/worker/app/services/document_agent/tools/inspect_pages.py index e04cd8152..856aaa5c3 100644 --- a/apps/worker/app/services/document_agent/tools/inspect_pages.py +++ b/apps/worker/app/services/document_agent/tools/inspect_pages.py @@ -90,23 +90,41 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) - from app.services.document_agent.visual import render_pages - folder_name = str(args.get("folder_name") or "inspect_pages") prefix = str(args.get("prefix") or "inspect") - rendered = render_pages( - ctx, - pages, - folder_name=folder_name, - prefix=prefix, - timeout=120, - ) - if not rendered: - return ToolResult( - status="error", - error="render failed", - latency_ms=int((time.monotonic() - start) * 1000), + pre_rendered = args.get("rendered_pages") + if isinstance(pre_rendered, list) and pre_rendered: + rendered = [ + {"page": int(item["page"]), "png_path": str(item["png_path"])} + for item in pre_rendered + if isinstance(item, dict) + and item.get("page") is not None + and item.get("png_path") + and int(item["page"]) in pages + ] + rendered = sorted(rendered, key=lambda item: int(item["page"])) + if len(rendered) != len(pages): + return ToolResult( + status="error", + error="rendered_pages missing coverage for requested pages", + latency_ms=int((time.monotonic() - start) * 1000), + ) + else: + from app.services.document_agent.visual import render_pages + + rendered = render_pages( + ctx, + pages, + folder_name=folder_name, + prefix=prefix, + timeout=120, ) + if not rendered: + return ToolResult( + status="error", + error="render failed", + latency_ms=int((time.monotonic() - start) * 1000), + ) model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") if not model: diff --git a/apps/worker/app/services/document_agent/visual.py b/apps/worker/app/services/document_agent/visual.py index e09f34aa7..a6feb0ee1 100644 --- a/apps/worker/app/services/document_agent/visual.py +++ b/apps/worker/app/services/document_agent/visual.py @@ -76,7 +76,13 @@ def _render_pages_worker( page = doc[idx] mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) pix = page.get_pixmap(matrix=mat) - png_name = f"{prefix}_page_{page_num}.png" + # Empty prefix → page_memory style ``page-{n}.png``. + # Non-empty keeps legacy ``{prefix}_page_{n}.png`` for PROFILE tools. + png_name = ( + f"page-{page_num}.png" + if not prefix + else f"{prefix}_page_{page_num}.png" + ) png_path = os.path.join(output_dir, png_name) pix.save(png_path) results.append({"page": page_num, "png_path": png_path}) diff --git a/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py index 0964825b0..5b3b0c3dd 100644 --- a/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py +++ b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py @@ -42,7 +42,6 @@ QUEUE_POLL_INTERVAL_SECONDS = 0.1 CHILD_EXIT_GRACE_SECONDS = 5 POST_RESULT_EXIT_GRACE_SECONDS = 5 -POST_KILL_JOIN_GRACE_SECONDS = 1 # Child can exit 0 before multiprocessing.Queue's feeder flushes; retry once. EMPTY_QUEUE_EXIT_RETRIES = 1 PROCESS_POOL_SIZE = read_pymupdf_max_concurrent() @@ -122,14 +121,16 @@ def _shutdown_process_pool() -> None: def _close_result_queue(result_queue: MultiprocessingQueue) -> None: """Release parent-side queue resources once the child result is no longer needed.""" try: - result_queue.close() + # Avoid joining the queue feeder thread after we already have the payload; + # otherwise parent/child teardown can stall each other. + result_queue.cancel_join_thread() except Exception as exc: - logger.debug(f"Failed to close PyMuPDF result queue: {exc}") + logger.debug(f"Failed to cancel PyMuPDF result queue join: {exc}") try: - result_queue.join_thread() + result_queue.close() except Exception as exc: - logger.debug(f"Failed to join PyMuPDF result queue thread: {exc}") + logger.debug(f"Failed to close PyMuPDF result queue: {exc}") def _is_empty_queue_exit(exc: Exception) -> bool: @@ -242,7 +243,8 @@ def _run_worker_in_spawned_process_once( ), ) - _close_result_queue(result_queue) + # Success path: wait for natural child exit before tearing down the queue. + # Do not kill — the payload is already in hand; forced kill only adds noise. proc.join(timeout=POST_RESULT_EXIT_GRACE_SECONDS) elapsed = time.monotonic() - t0 exit_lag = None @@ -250,17 +252,14 @@ def _run_worker_in_spawned_process_once( exit_lag = elapsed - (wait_result.received_at - t0) if proc.is_alive(): - proc.kill() - proc.join(timeout=POST_KILL_JOIN_GRACE_SECONDS) - elapsed = time.monotonic() - t0 - if wait_result.received_at is not None: - exit_lag = elapsed - (wait_result.received_at - t0) - logger.warning( + logger.debug( f"[pymupdf-subprocess] EXIT_DELAY pid={child_pid} fn={worker_fn.__name__} " f"elapsed={elapsed:.1f}s exit_lag={exit_lag:.1f}s " - f"— result returned before child exited; child killed after grace" + f"— result returned; leaving child to finish exit without kill" ) + _close_result_queue(result_queue) + result = wait_result.result or {} if not result.get("ok"): logger.error( diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index d08d4c501..dfc2ac196 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -441,6 +441,7 @@ def _build_page_dataframe( skeletons=skeletons, raw_text_by_page=raw_text_by_page, image_path_by_page=image_path_by_page, + output_dir=output_dir, kind_by_page=label_map, tag_by_page=tag_map, filename=filename, @@ -614,6 +615,7 @@ def _run_hierarchy_scope( pdf_path=pdf_path, page_count=page_count, output_dir=output_dir, + scope_id=scope.scope_id, pages=title_pages, page_features=page_features, page_texts=page_texts, @@ -697,6 +699,7 @@ def _run_hierarchy_scope( pdf_path=pdf_path, page_count=page_count, output_dir=output_dir, + scope_id=scope.scope_id, pages=final_pages, page_features=page_features, page_texts=page_texts, diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index 7273a0932..c4dc46257 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -382,6 +382,7 @@ def build_node_rows( skeletons: list[SectionSkeleton], raw_text_by_page: dict[int, str], image_path_by_page: dict[int, str], + output_dir: str, kind_by_page: dict[int, str], tag_by_page: dict[int, PageTagResult], filename: str, @@ -469,6 +470,7 @@ def _summarize_one( "extra_metadata": _build_page_extra_metadata( pages=view.pages, image_path_by_page=image_path_by_page, + output_dir=output_dir, ), } rows.append(row) @@ -498,10 +500,12 @@ def _build_page_extra_metadata( *, pages: list[int], image_path_by_page: dict[int, str], + output_dir: str, ) -> dict[str, Any]: page_assets = _build_page_citation_assets( pages=pages, image_path_by_page=image_path_by_page, + output_dir=output_dir, ) if not page_assets: return {} @@ -512,6 +516,7 @@ def _build_page_citation_assets( *, pages: list[int], image_path_by_page: dict[int, str], + output_dir: str, ) -> list[dict[str, Any]]: assets: list[dict[str, Any]] = [] seen_pages: set[int] = set() @@ -522,7 +527,11 @@ def _build_page_citation_assets( image_path = image_path_by_page.get(page) if not image_path or not os.path.exists(image_path): continue - artifact_ref = _promote_page_citation_asset(page=page, image_path=image_path) + artifact_ref = _promote_page_citation_asset( + page=page, + image_path=image_path, + output_dir=output_dir, + ) if not artifact_ref: continue width, height = _read_image_dimensions(image_path) @@ -540,10 +549,14 @@ def _build_page_citation_assets( return assets -def _promote_page_citation_asset(*, page: int, image_path: str) -> str: +def _promote_page_citation_asset( + *, + page: int, + image_path: str, + output_dir: str, +) -> str: source_path = Path(image_path) - output_dir = source_path.parent.parent - target_dir = output_dir / "page_citation_assets" + target_dir = Path(output_dir) / "page_citation_assets" target_path = target_dir / f"page-{page}.png" artifact_ref = f"page_citation_assets/page-{page}.png" try: diff --git a/apps/worker/app/services/page_memory/page_renderer.py b/apps/worker/app/services/page_memory/page_renderer.py index 944d7083a..925c15df8 100644 --- a/apps/worker/app/services/page_memory/page_renderer.py +++ b/apps/worker/app/services/page_memory/page_renderer.py @@ -23,7 +23,7 @@ class PageRenderResult: """1-based page number.""" image_path: str - """Absolute path to full-resolution PNG (``pages/page-N.png``).""" + """Absolute path to full-resolution PNG (``pages/{scope_id}/page-{n}.png``).""" raw_text: str """PROFILE scan text for this page.""" @@ -43,6 +43,7 @@ def render_document_pages( pdf_path: str, page_count: int, output_dir: str, + scope_id: str, pages: list[int] | None = None, page_features: list[PageFeature] | None = None, page_texts: dict[int, str] | None = None, @@ -61,7 +62,11 @@ def render_document_pages( pages: Optional 1-based page subset to render. Defaults to every page. output_dir: - Root output directory; pages are written to ``output_dir/pages/``. + Root output directory; pages are written to + ``output_dir/pages/{scope_id}/page-{n}.png``. + scope_id: + Scope isolation key (e.g. ``p58-60``). Concurrent scopes must not + share a write path for the same page index. page_features: If available, dimensions are read from here (avoiding a second PyMuPDF open). Otherwise falls back to 0/0/False. @@ -81,6 +86,10 @@ def render_document_pages( list[PageRenderResult] One entry per page, ordered by page_index. """ + resolved_scope = str(scope_id or "").strip() + if not resolved_scope: + raise ValueError("render_document_pages requires a non-empty scope_id") + requested_pages = ( sorted({page for page in pages if 1 <= page <= page_count}) if pages is not None @@ -90,14 +99,16 @@ def render_document_pages( return [] texts = page_texts or {} + folder_name = f"pages/{resolved_scope}" # ── full-resolution PNGs ────────────────────────────────────────── + # Empty prefix → ``page-{n}.png`` (see visual._render_pages_worker). if ctx is not None: pngs = render_pages( ctx, requested_pages, - folder_name="pages", - prefix="page", + folder_name=folder_name, + prefix="", dpi=dpi, timeout=timeout, ) @@ -117,8 +128,8 @@ def render_document_pages( pngs = render_pages( tmp_ctx, requested_pages, - folder_name="pages", - prefix="page", + folder_name=folder_name, + prefix="", dpi=dpi, timeout=timeout, ) diff --git a/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py b/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py index 700f03219..0b28900a3 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py @@ -171,6 +171,7 @@ def _run_fine_hierarchy_for_scope( pdf_path=pdf_path, page_count=page_count, output_dir=str(out_dir), + scope_id=scope_id, pages=title_pages, page_features=page_features, page_texts=page_texts, @@ -257,6 +258,7 @@ def _run_fine_hierarchy_for_scope( pdf_path=pdf_path, page_count=page_count, output_dir=str(out_dir), + scope_id=scope_id, pages=final_pages, page_features=page_features, page_texts=page_texts, diff --git a/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py b/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py index 9f657b55b..127121c2f 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py @@ -200,6 +200,7 @@ def main() -> int: pdf_path=pdf_path, page_count=page_count, output_dir=str(out_dir), + scope_id="c5_assets", pages=union_pages, page_features=page_features, page_texts=page_texts, diff --git a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py index 5321decf1..07da9530c 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py @@ -142,6 +142,7 @@ def _run_tagging_for_scope( pdf_path=pdf_path, page_count=page_count, output_dir=str(out_dir), + scope_id=scope_id, pages=final_pages, page_features=page_features, page_texts=page_texts, @@ -462,6 +463,7 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: skeletons=active_skeletons, raw_text_by_page=raw_text_by_page, image_path_by_page=image_path_by_page, + output_dir=str(out_dir), kind_by_page=label_map, tag_by_page=tag_map, filename=filename, diff --git a/apps/worker/tests/contract/test_calibration_scan_contract.py b/apps/worker/tests/contract/test_calibration_scan_contract.py index 34a5aa4d1..b28a114cb 100644 --- a/apps/worker/tests/contract/test_calibration_scan_contract.py +++ b/apps/worker/tests/contract/test_calibration_scan_contract.py @@ -60,10 +60,22 @@ def __call__(self, ctx: ToolContext, args: dict[str, Any]) -> ToolResult: ) +def _fake_render_pages( + ctx: ToolContext, + pages: list[int], + **kwargs: Any, +) -> list[dict[str, Any]]: + return [ + {"page": int(page), "png_path": f"/tmp/scan_page_{int(page)}.png"} + for page in pages + ] + + @pytest.fixture def patch_inspect(monkeypatch: pytest.MonkeyPatch): def _apply(fake: _FakeInspect) -> _FakeInspect: monkeypatch.setattr(scan_module, "inspect_pages", fake) + monkeypatch.setattr(scan_module, "render_pages", _fake_render_pages) return fake return _apply @@ -241,3 +253,51 @@ def __call__(self, ctx: ToolContext, args: dict[str, Any]) -> ToolResult: assert result.found_page == 13 assert sorted(fake.calls) == [[10], [11], [12], [13], [14], [15]] + + +def test_window_renders_once_serially_before_concurrent_vlm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """PDF→PNG is one serial batch per window; VLM stays one page per call.""" + render_calls: list[list[int]] = [] + inspect_args: list[dict[str, Any]] = [] + + def _recording_render( + ctx: ToolContext, + pages: list[int], + **kwargs: Any, + ) -> list[dict[str, Any]]: + render_calls.append(list(pages)) + return _fake_render_pages(ctx, pages, **kwargs) + + def _capture_inspect(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + inspect_args.append(dict(args)) + pages = list(args.get("pages") or []) + hit = 10 in pages + return ToolResult( + status="ok", + payload={ + "pages": pages, + "answer": "", + "fields": { + "found": hit, + "reason": "hit" if hit else "miss", + }, + }, + ) + + monkeypatch.setattr(scan_module, "inspect_pages", _capture_inspect) + monkeypatch.setattr(scan_module, "render_pages", _recording_render) + + result = scan_title_forward( + ctx=_ctx(), title="Chapter 1", start_page=10, page_count=60 + ) + + assert result.found_page == 10 + assert render_calls == [[10, 11]] + assert len(inspect_args) == 2 + assert sorted(args["pages"] for args in inspect_args) == [[10], [11]] + for args in inspect_args: + assert args["page_cap"] == 1 + assert len(args["rendered_pages"]) == 1 + assert args["rendered_pages"][0]["page"] == args["pages"][0] diff --git a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py index 58699fdeb..773f14a2e 100644 --- a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -113,6 +113,7 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={}, + output_dir="/tmp/knowhere-test-node-assembler", kind_by_page={}, tag_by_page={ 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), @@ -161,6 +162,7 @@ def _fake_compute_node_summary(**kwargs): skeletons=_ordered_page_skeletons(), raw_text_by_page={1: "text-1", 2: "text-2", 3: "text-3"}, image_path_by_page={}, + output_dir="/tmp/knowhere-test-node-assembler", kind_by_page={}, tag_by_page={}, filename="demo.pdf", @@ -202,6 +204,7 @@ def _fake_compute_node_summary(**kwargs): skeletons=_ordered_page_skeletons()[:1], raw_text_by_page={1: "text-1"}, image_path_by_page={}, + output_dir="/tmp/knowhere-test-node-assembler", kind_by_page={}, tag_by_page={}, filename="demo.pdf", @@ -212,14 +215,16 @@ def _fake_compute_node_summary(**kwargs): def test_build_node_rows_attaches_page_citation_assets_for_rendered_pages(tmp_path) -> None: - page_image = tmp_path / "pages" / "page-231.png" - page_image.parent.mkdir() + # Nested scope path matches production pages/{scope_id}/page-{n}.png. + page_image = tmp_path / "pages" / "p231-231" / "page-231.png" + page_image.parent.mkdir(parents=True) Image.new("RGB", (2, 3), color=(255, 255, 255)).save(page_image) rows = node_assembler.build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={231: str(page_image)}, + output_dir=str(tmp_path), kind_by_page={}, tag_by_page={ 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), @@ -244,6 +249,8 @@ def test_build_node_rows_attaches_page_citation_assets_for_rendered_pages(tmp_pa } ] assert (tmp_path / "page_citation_assets" / "page-231.png").is_file() + # Must not land under the temporary pages/ tree. + assert not (tmp_path / "pages" / "page_citation_assets").exists() def test_build_node_rows_keeps_internal_section_body_pages() -> None: @@ -268,6 +275,7 @@ def test_build_node_rows_keeps_internal_section_body_pages() -> None: skeletons=[parent, child], raw_text_by_page={233: "parent body", 234: "child body"}, image_path_by_page={}, + output_dir="/tmp/knowhere-test-node-assembler", kind_by_page={}, tag_by_page={ 233: PageTagResult(page_index=233, summary="s233", keywords=["parent"]), @@ -350,6 +358,7 @@ def chat_completion_with_usage(self, **kwargs): skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={231: str(img), 232: str(img)}, + output_dir=str(tmp_path), kind_by_page={}, tag_by_page={ 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), @@ -390,6 +399,7 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={}, + output_dir="/tmp/knowhere-test-node-assembler", kind_by_page={}, tag_by_page={ 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), diff --git a/apps/worker/tests/contract/test_page_memory_page_renderer_scope_path_contract.py b/apps/worker/tests/contract/test_page_memory_page_renderer_scope_path_contract.py new file mode 100644 index 000000000..6effce400 --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_page_renderer_scope_path_contract.py @@ -0,0 +1,64 @@ +"""Contract: page_memory renders into per-scope paths with page-{n}.png names.""" + +from __future__ import annotations + +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from pathlib import Path + +import pytest + +import app.services.page_memory.page_renderer as page_renderer + + +def test_render_document_pages_requires_scope_id(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="scope_id"): + page_renderer.render_document_pages( + pdf_path=str(tmp_path / "missing.pdf"), + page_count=1, + output_dir=str(tmp_path), + scope_id=" ", + pages=[1], + ) + + +def test_render_document_pages_writes_under_scope_folder( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + captured: dict[str, object] = {} + + def _fake_render_pages(ctx, pages, **kwargs): + captured["folder_name"] = kwargs.get("folder_name") + captured["prefix"] = kwargs.get("prefix") + captured["output_dir"] = ctx.output_dir + folder = Path(ctx.output_dir) / str(kwargs["folder_name"]) + folder.mkdir(parents=True, exist_ok=True) + results = [] + for page in pages: + png_path = folder / f"page-{page}.png" + png_path.write_bytes(b"png") + results.append({"page": page, "png_path": str(png_path)}) + return results + + monkeypatch.setattr(page_renderer, "render_pages", _fake_render_pages) + + results = page_renderer.render_document_pages( + pdf_path=str(tmp_path / "doc.pdf"), + page_count=78, + output_dir=str(tmp_path), + scope_id="p58-60", + pages=[60], + ) + + assert captured["folder_name"] == "pages/p58-60" + assert captured["prefix"] == "" + assert len(results) == 1 + assert results[0].image_path.endswith("/pages/p58-60/page-60.png") + assert Path(results[0].image_path).is_file() diff --git a/packages/shared-python/shared/services/ai/summary/engine.py b/packages/shared-python/shared/services/ai/summary/engine.py index 2066ab795..76e4129d0 100644 --- a/packages/shared-python/shared/services/ai/summary/engine.py +++ b/packages/shared-python/shared/services/ai/summary/engine.py @@ -49,7 +49,11 @@ def _read_image_b64(image_path: str) -> str | None: try: with open(image_path, "rb") as handle: - return base64.b64encode(handle.read()).decode() + raw = handle.read() + if not raw: + logger.warning("[summary] empty image file {}", image_path) + return None + return base64.b64encode(raw).decode() except Exception as exc: logger.warning("[summary] failed to read image {}: {}", image_path, exc) return None From 1d12b67c2ede9db4fadc2b171ff786060a33aeb3 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 26 Aug 2026 10:53:31 +0800 Subject: [PATCH 4/7] fix: pass toc_pages into outline skeleton and import extract_toc_nodes Unblocks make check before PR sync. Co-authored-by: Cursor --- .../app/services/document_agent/structure/toc_anchoring.py | 2 ++ .../scripts/page_memory/debug_pm_stage5_tagging_finalize.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/apps/worker/app/services/document_agent/structure/toc_anchoring.py b/apps/worker/app/services/document_agent/structure/toc_anchoring.py index 5e0f7c34a..3e8b57f8c 100644 --- a/apps/worker/app/services/document_agent/structure/toc_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/toc_anchoring.py @@ -215,6 +215,7 @@ def _try_outline_anchoring_route( page_texts=page_texts, body_pages=body_pages, page_count=page_count, + toc_pages=toc_pages, ): return False @@ -235,6 +236,7 @@ def _write_outline_skeleton( page_texts: dict[int, str], body_pages: list[int], page_count: int, + toc_pages: list[int] | None = None, ) -> bool: """Anchor outline rows via physical overrides (no calibrate VLM). diff --git a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py index 07da9530c..afa9d7183 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py @@ -574,6 +574,10 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: ) # ── Final trace + cross-stage cost rollup ── + from app.services.document_agent.structure.hierarchy_locator import ( + extract_toc_nodes, + ) + toc_nodes = ( extract_toc_nodes(anatomy.toc_hierarchies) if anatomy.toc_hierarchies else [] ) From ed0c3a90f6154d274c7f6f5ed18a6789a5018c74 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 27 Aug 2026 21:23:36 +0800 Subject: [PATCH 5/7] feat: enhance token tracking and usage reporting in document processing - Integrated token tracking into the concurrent page inspection process, allowing for detailed tracking of token usage during parsing. - Updated the summarization functions for DOCX and Excel formats to include usage task metadata, improving the granularity of token usage reporting. - Refactored the debug scripts to better manage token ledger operations, including loading and recording token usage across different stages. - Enhanced unit tests to validate the new token tracking features and ensure accurate reporting of token usage in various scenarios. --- .../document_agent/calibration/scan.py | 21 +- .../document_parser/formats/docx/parser.py | 7 +- .../formats/excel/table_parser.py | 1 + apps/worker/scripts/_debug_token_ledger.py | 128 ++++++ apps/worker/scripts/debug_text_track.py | 380 +++++++++++------- .../scripts/page_memory/_debug_pm_shared.py | 148 ++----- .../debug_pm_stage3_scope_fine_hierarchy.py | 2 - .../debug_pm_stage5_tagging_finalize.py | 4 - .../test_calibration_scan_contract.py | 37 ++ .../test_parse_usage_task_contract.py | 67 +++ .../test_debug_text_track_token_ledger.py | 67 +++ .../unit/test_processing_metadata_persist.py | 49 +++ .../shared/services/ai/summary/engine.py | 81 +--- 13 files changed, 654 insertions(+), 338 deletions(-) create mode 100644 apps/worker/scripts/_debug_token_ledger.py create mode 100644 apps/worker/tests/contract/test_parse_usage_task_contract.py create mode 100644 apps/worker/tests/unit/test_debug_text_track_token_ledger.py diff --git a/apps/worker/app/services/document_agent/calibration/scan.py b/apps/worker/app/services/document_agent/calibration/scan.py index 9050ab1d7..82a220f83 100644 --- a/apps/worker/app/services/document_agent/calibration/scan.py +++ b/apps/worker/app/services/document_agent/calibration/scan.py @@ -28,6 +28,10 @@ from app.services.document_agent.manifest import ToolContext from app.services.document_agent.tools.inspect_pages import inspect_pages from app.services.document_agent.visual import render_pages +from shared.services.ai.token_tracking import ( + bind_token_tracker, + get_current_token_tracker_root_id, +) DEFAULT_WINDOW_SCHEDULE: tuple[int, ...] = (2, 4, 6, 10) @@ -189,16 +193,21 @@ def _inspect_pages_concurrent( ) ] - by_page: dict[int, PageInspectResult] = {} - with ThreadPoolExecutor(max_workers=len(pages)) as pool: - futures = { - pool.submit( - _inspect_one_page, + token_tracker_root_id = get_current_token_tracker_root_id() + + def _inspect_one_page_with_tracking(page: int) -> PageInspectResult: + with bind_token_tracker(token_tracker_root_id): + return _inspect_one_page( ctx=ctx, title=title, page=page, rendered_page=rendered_by_page[page], - ): page + ) + + by_page: dict[int, PageInspectResult] = {} + with ThreadPoolExecutor(max_workers=len(pages)) as pool: + futures = { + pool.submit(_inspect_one_page_with_tracking, page): page for page in pages } for future in as_completed(futures): diff --git a/apps/worker/app/services/document_parser/formats/docx/parser.py b/apps/worker/app/services/document_parser/formats/docx/parser.py index f978ab77d..00ef0772c 100755 --- a/apps/worker/app/services/document_parser/formats/docx/parser.py +++ b/apps/worker/app/services/document_parser/formats/docx/parser.py @@ -432,7 +432,12 @@ def handle_table( from shared.services.ai.summary.engine import summarize # Tables are Contract B assets: title + summary + entities from HTML. - result = summarize(mode="asset", text=tb_html_str, max_keywords=3) + result = summarize( + mode="asset", + text=tb_html_str, + max_keywords=3, + usage_task="parser.docx.table", + ) llm_title = result.title or None tb_keywords = result.keywords_str() llm_summary = result.summary or None diff --git a/apps/worker/app/services/document_parser/formats/excel/table_parser.py b/apps/worker/app/services/document_parser/formats/excel/table_parser.py index 4aa4c5e78..367a56dd5 100644 --- a/apps/worker/app/services/document_parser/formats/excel/table_parser.py +++ b/apps/worker/app/services/document_parser/formats/excel/table_parser.py @@ -325,6 +325,7 @@ def _summarize_excel_table( mode="asset", text=table_html, max_keywords=3, + usage_task="parser.excel.table", ) return ( result.title or None, diff --git a/apps/worker/scripts/_debug_token_ledger.py b/apps/worker/scripts/_debug_token_ledger.py new file mode 100644 index 000000000..c0d81308e --- /dev/null +++ b/apps/worker/scripts/_debug_token_ledger.py @@ -0,0 +1,128 @@ +"""Shared token-ledger helpers for debug parse scripts.""" + +from __future__ import annotations + +import json +from copy import deepcopy +from pathlib import Path +from typing import Any + +_NUMERIC_USAGE_FIELDS = ("prompt_tokens", "completion_tokens", "total_tokens", "calls") + + +def empty_token_usage() -> dict[str, Any]: + return { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "calls": 0, + "by_model": {}, + "by_task": {}, + } + + +def token_usage_delta(prev: dict[str, Any], cur: dict[str, Any]) -> dict[str, Any]: + numeric = _NUMERIC_USAGE_FIELDS + + def _sub(left: dict[str, Any], right: dict[str, Any]) -> dict[str, int]: + return { + field: int(right.get(field, 0)) - int(left.get(field, 0)) for field in numeric + } + + def _bucket(prev_bucket: dict[str, Any], cur_bucket: dict[str, Any]) -> dict[str, Any]: + merged: dict[str, Any] = {} + for key in set(prev_bucket) | set(cur_bucket): + prev_item = prev_bucket.get(key, {}) + cur_item = cur_bucket.get(key, {}) + if not isinstance(prev_item, dict) or not isinstance(cur_item, dict): + continue + entry = { + field: value for field, value in _sub(prev_item, cur_item).items() if value + } + prev_models = prev_item.get("models", {}) + cur_models = cur_item.get("models", {}) + if prev_models or cur_models: + models = _bucket(prev_models, cur_models) + if models: + entry["models"] = models + if entry: + merged[key] = entry + return merged + + delta = _sub(prev, cur) + for bucket_key in ("by_model", "by_task"): + bucket = _bucket(prev.get(bucket_key, {}), cur.get(bucket_key, {})) + if bucket: + delta[bucket_key] = bucket + return delta + + +def merge_token_usage(destination: dict[str, Any], source: dict[str, Any]) -> None: + for key, value in source.items(): + if isinstance(value, dict): + child = destination.setdefault(str(key), {}) + if isinstance(child, dict): + merge_token_usage(child, value) + elif isinstance(value, int | float) and not isinstance(value, bool): + destination[str(key)] = destination.get(str(key), 0) + value + + +def load_stage_ledger(path: Path, *, version: str = "1.0") -> dict[str, Any]: + if not path.exists(): + return {"version": version, "stages": {}} + data = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(data, dict): + return {"version": version, "stages": {}} + data.setdefault("version", version) + data.setdefault("stages", {}) + return data + + +def write_stage_ledger(path: Path, ledger: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temp_path = path.with_suffix(f"{path.suffix}.tmp") + temp_path.write_text(json.dumps(ledger, ensure_ascii=False, indent=2), encoding="utf-8") + temp_path.replace(path) + + +def record_stage_delta( + ledger: dict[str, Any], + *, + stage: str, + stage_keys: tuple[str, ...], + prev: dict[str, Any], + current: dict[str, Any], + out_path: Path, +) -> dict[str, Any]: + """Store per-stage delta and drop stale downstream stage keys.""" + stages = ledger.setdefault("stages", {}) + if not isinstance(stages, dict): + stages = {} + ledger["stages"] = stages + stage_index = stage_keys.index(stage) + for stale in stage_keys[stage_index + 1 :]: + stages.pop(stale, None) + stages[stage] = {"token_usage": token_usage_delta(prev, current)} + write_stage_ledger(out_path, ledger) + return deepcopy(current) + + +def aggregate_stage_deltas( + ledger: dict[str, Any], + stage_keys: tuple[str, ...], + *, + remainder: dict[str, Any] | None = None, +) -> dict[str, Any]: + usage = empty_token_usage() + stages = ledger.get("stages") + if isinstance(stages, dict): + for stage in stage_keys: + row = stages.get(stage) + if not isinstance(row, dict): + continue + raw_usage = row.get("token_usage") + if isinstance(raw_usage, dict): + merge_token_usage(usage, raw_usage) + if remainder: + merge_token_usage(usage, remainder) + return usage diff --git a/apps/worker/scripts/debug_text_track.py b/apps/worker/scripts/debug_text_track.py index d80c15bd9..e14176a1f 100644 --- a/apps/worker/scripts/debug_text_track.py +++ b/apps/worker/scripts/debug_text_track.py @@ -29,6 +29,7 @@ import shutil import sys import time +from copy import deepcopy from pathlib import Path from typing import Any @@ -36,6 +37,7 @@ ROOT = Path(__file__).resolve().parents[3] WORKER_ROOT = ROOT / "apps" / "worker" sys.path.insert(0, str(WORKER_ROOT)) +sys.path.insert(0, str(WORKER_ROOT / "scripts")) sys.path.insert(0, str(ROOT / "packages" / "shared-python")) from dotenv import load_dotenv @@ -47,8 +49,17 @@ from loguru import logger from shared.services.ai.token_tracking import ( - init_token_tracker, + cleanup_token_tracker, get_current_token_tracker, + init_token_tracker, +) + +from _debug_token_ledger import ( + aggregate_stage_deltas, + empty_token_usage, + load_stage_ledger, + record_stage_delta, + token_usage_delta, ) # ── Constants ─────────────────────────────────────────────────────────────── @@ -58,6 +69,8 @@ ) DEFAULT_SPACEX_PDF = Path("/Users/wuchengke/Desktop/temp/test_docs/spacex-s1.pdf") OUTPUT_ROOT = Path("~/.knowhere/_debug_parse").expanduser() +TOKEN_LEDGER_NAME = "token_ledger.json" +TOKEN_LEDGER_STAGES = ("profile", "mineru", "hierarchy", "full") def _write_json(path: Path, data: Any) -> None: @@ -66,6 +79,61 @@ def _write_json(path: Path, data: Any) -> None: logger.info(" → {}", path) +def _load_token_ledger(out_dir: Path) -> dict[str, Any]: + return load_stage_ledger(out_dir / TOKEN_LEDGER_NAME) + + +def _record_token_stage( + ledger: dict[str, Any], + stage: str, + *, + prev: dict[str, Any], + out_dir: Path, +) -> dict[str, Any]: + current = deepcopy(get_current_token_tracker() or {}) + record_stage_delta( + ledger, + stage=stage, + stage_keys=TOKEN_LEDGER_STAGES, + prev=prev, + current=current, + out_path=out_dir / TOKEN_LEDGER_NAME, + ) + logger.info(" → {}", out_dir / TOKEN_LEDGER_NAME) + return current + + +def _aggregate_token_ledger( + ledger: dict[str, Any], + *, + remainder: dict[str, Any] | None = None, +) -> dict[str, Any]: + return aggregate_stage_deltas( + ledger, + TOKEN_LEDGER_STAGES, + remainder=remainder, + ) + + +def _apply_token_usage_to_outputs( + out_dir: Path, + trace: dict[str, Any], + usage: dict[str, Any], +) -> None: + trace["token_usage"] = usage + _write_json(out_dir / "trace.json", trace) + manifest_path = out_dir / "manifest.json" + if not manifest_path.exists(): + return + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + if not isinstance(manifest, dict): + return + processing = manifest.setdefault("processing", {}) + if isinstance(processing, dict): + processing["token_usage"] = usage + _write_json(manifest_path, manifest) + + # ── Stage 1: Profile + Shard Plan (PDF only) ─────────────────────────────── def _stage_profile(pdf_path: str, filename: str, out_dir: Path, model: str | None): @@ -716,6 +784,8 @@ def main() -> int: out_dir.mkdir(parents=True, exist_ok=True) init_token_tracker() + ledger = _load_token_ledger(out_dir) + prev_usage = deepcopy(get_current_token_tracker() or empty_token_usage()) logger.info("█" * 70) logger.info(" TEXT-TRACK DEBUG: {}", filename) @@ -731,158 +801,188 @@ def main() -> int: } t_start = time.time() - # ── Format dispatch ────────────────────────────────────────────────────── - if ext == ".pdf": - # Stage 1: Profile - if args.reuse_profile: - anatomy = _load_anatomy_cache(out_dir, file_path, filename) - else: - anatomy, profile_elapsed, profile_meta = _stage_profile( - file_path, filename, out_dir, args.model + try: + # ── Format dispatch ────────────────────────────────────────────────── + if ext == ".pdf": + # Stage 1: Profile + if args.reuse_profile: + anatomy = _load_anatomy_cache(out_dir, file_path, filename) + else: + anatomy, profile_elapsed, profile_meta = _stage_profile( + file_path, filename, out_dir, args.model + ) + trace["stages"]["profile"] = { + "elapsed_s": round(profile_elapsed, 1), + **profile_meta, + } + prev_usage = _record_token_stage( + ledger, "profile", prev=prev_usage, out_dir=out_dir + ) + + if args.stop_at == "profile": + trace["stages"].setdefault("profile", {})["shard_count"] = len( + anatomy.shard_plan.shards + ) + logger.info("⏸️ Stopped at profile → {}", out_dir) + return 0 + + # Stage 2: MinerU extraction + if not args.reuse_mineru: + _shard_dirs, mineru_elapsed = _stage_mineru_pdf( + file_path, filename, out_dir, anatomy, + ) + trace["stages"]["mineru"] = { + "elapsed_s": round(mineru_elapsed, 1), + "shard_count": len(_shard_dirs), + } + prev_usage = _record_token_stage( + ledger, "mineru", prev=prev_usage, out_dir=out_dir + ) + else: + logger.info("⏩ Reusing cached MinerU shard dirs") + + if args.stop_at == "mineru": + logger.info("⏸️ Stopped at mineru → {}", out_dir) + return 0 + + # Stage 3: Heading prediction + hierarchy tree + if args.reuse_hierarchy: + merged_path = out_dir / "_shards" / "merged_lines.json" + if not merged_path.exists(): + raise FileNotFoundError(f"No cached hierarchy: {merged_path}") + logger.info("⏩ Reusing cached merged_lines: {}", merged_path) + merged_lines = json.loads(merged_path.read_text(encoding="utf-8")) + else: + merged_lines, hier_elapsed = _stage_hierarchy_pdf( + out_dir, anatomy, args.model, + ) + trace["stages"]["hierarchy"] = { + "elapsed_s": round(hier_elapsed, 1), + "merged_lines_count": len(merged_lines), + "heading_count": sum(1 for ln in merged_lines if ln.startswith("#")), + } + prev_usage = _record_token_stage( + ledger, "hierarchy", prev=prev_usage, out_dir=out_dir + ) + + if args.stop_at == "hierarchy": + logger.info("⏸️ Stopped at hierarchy → {}", out_dir) + return 0 + + # Stage 4: Full extraction + chunks, full_elapsed = _stage_full_pdf( + out_dir, filename, merged_lines, args.model ) - trace["stages"]["profile"] = { - "elapsed_s": round(profile_elapsed, 1), - **profile_meta, + trace["stages"]["full"] = { + "elapsed_s": round(full_elapsed, 1), + "chunk_count": len(chunks), } + prev_usage = _record_token_stage( + ledger, "full", prev=prev_usage, out_dir=out_dir + ) + + elif ext in (".docx", ".doc"): + if args.stop_at in ("profile", "mineru"): + logger.info("ℹ️ No profiling/MinerU for DOCX format. Nothing to do.") + return 0 - if args.stop_at == "profile": - trace["stages"].setdefault("profile", {})["shard_count"] = len( - anatomy.shard_plan.shards + parsed_df, hier_elapsed = _stage_hierarchy_docx( + file_path, filename, out_dir, args.model + ) + trace["stages"]["hierarchy"] = { + "elapsed_s": round(hier_elapsed, 1), + "row_count": len(parsed_df) if parsed_df is not None else 0, + } + prev_usage = _record_token_stage( + ledger, "hierarchy", prev=prev_usage, out_dir=out_dir ) - trace["token_usage"] = get_current_token_tracker() - _write_json(out_dir / "trace.json", trace) - logger.info("⏸️ Stopped at profile → {}", out_dir) - return 0 - - # Stage 2: MinerU extraction - if not args.reuse_mineru: - _shard_dirs, mineru_elapsed = _stage_mineru_pdf( - file_path, filename, out_dir, anatomy, + + if args.stop_at == "hierarchy": + logger.info("⏸️ Stopped at hierarchy → {}", out_dir) + return 0 + + from app.services.document_parser.orchestration.postprocess import ( + apply_parse_postprocess, ) - trace["stages"]["mineru"] = { - "elapsed_s": round(mineru_elapsed, 1), - "shard_count": len(_shard_dirs), + parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) + chunks, full_elapsed = _finalize_df(out_dir, filename, parsed_df) + trace["stages"]["full"] = { + "elapsed_s": round(full_elapsed, 1), + "chunk_count": len(chunks), } - else: - logger.info("⏩ Reusing cached MinerU shard dirs") - - if args.stop_at == "mineru": - trace["token_usage"] = get_current_token_tracker() - _write_json(out_dir / "trace.json", trace) - logger.info("⏸️ Stopped at mineru → {}", out_dir) - return 0 - - # Stage 3: Heading prediction + hierarchy tree - if args.reuse_hierarchy: - merged_path = out_dir / "_shards" / "merged_lines.json" - if not merged_path.exists(): - raise FileNotFoundError(f"No cached hierarchy: {merged_path}") - logger.info("⏩ Reusing cached merged_lines: {}", merged_path) - merged_lines = json.loads(merged_path.read_text(encoding="utf-8")) - else: - merged_lines, hier_elapsed = _stage_hierarchy_pdf( - out_dir, anatomy, args.model, + prev_usage = _record_token_stage( + ledger, "full", prev=prev_usage, out_dir=out_dir + ) + + elif ext in (".md", ".markdown"): + if args.stop_at in ("profile", "mineru"): + logger.info("ℹ️ No profiling/MinerU for Markdown format. Nothing to do.") + return 0 + + parsed_df, hier_elapsed = _stage_hierarchy_md( + file_path, filename, out_dir, args.model ) trace["stages"]["hierarchy"] = { "elapsed_s": round(hier_elapsed, 1), - "merged_lines_count": len(merged_lines), - "heading_count": sum(1 for ln in merged_lines if ln.startswith("#")), + "row_count": len(parsed_df) if parsed_df is not None else 0, } + prev_usage = _record_token_stage( + ledger, "hierarchy", prev=prev_usage, out_dir=out_dir + ) - if args.stop_at == "hierarchy": - trace["token_usage"] = get_current_token_tracker() - _write_json(out_dir / "trace.json", trace) - logger.info("⏸️ Stopped at hierarchy → {}", out_dir) - return 0 - - # Stage 4: Full extraction - chunks, full_elapsed = _stage_full_pdf(out_dir, filename, merged_lines, args.model) - trace["stages"]["full"] = { - "elapsed_s": round(full_elapsed, 1), - "chunk_count": len(chunks), - } - - elif ext in (".docx", ".doc"): - if args.stop_at in ("profile", "mineru"): - logger.info("ℹ️ No profiling/MinerU for DOCX format. Nothing to do.") - return 0 - - # Stage 2: parse_docx - parsed_df, hier_elapsed = _stage_hierarchy_docx(file_path, filename, out_dir, args.model) - trace["stages"]["hierarchy"] = { - "elapsed_s": round(hier_elapsed, 1), - "row_count": len(parsed_df) if parsed_df is not None else 0, - } - - if args.stop_at == "hierarchy": - _write_json(out_dir / "trace.json", trace) - logger.info("⏸️ Stopped at hierarchy → {}", out_dir) - return 0 - - # Stage 4: DataFrame → chunks - from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess - parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) - chunks, full_elapsed = _finalize_df(out_dir, filename, parsed_df) - trace["stages"]["full"] = { - "elapsed_s": round(full_elapsed, 1), - "chunk_count": len(chunks), - } - - elif ext in (".md", ".markdown"): - if args.stop_at in ("profile", "mineru"): - logger.info("ℹ️ No profiling/MinerU for Markdown format. Nothing to do.") - return 0 - - # Stage 2: parse_md - parsed_df, hier_elapsed = _stage_hierarchy_md(file_path, filename, out_dir, args.model) - trace["stages"]["hierarchy"] = { - "elapsed_s": round(hier_elapsed, 1), - "row_count": len(parsed_df) if parsed_df is not None else 0, - } - - if args.stop_at == "hierarchy": - _write_json(out_dir / "trace.json", trace) - logger.info("⏸️ Stopped at hierarchy → {}", out_dir) - return 0 - - # Stage 4: DataFrame → chunks - from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess - parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) - chunks, full_elapsed = _finalize_df(out_dir, filename, parsed_df) - trace["stages"]["full"] = { - "elapsed_s": round(full_elapsed, 1), - "chunk_count": len(chunks), - } + if args.stop_at == "hierarchy": + logger.info("⏸️ Stopped at hierarchy → {}", out_dir) + return 0 - else: - logger.error("Unsupported format: {}", ext) - return 1 - - # ── Optional DB publication ────────────────────────────────────────────── - if args.run_db: - from scripts._debug_publish import publish_debug_result_dir - publish_result = publish_debug_result_dir( - result_dir=out_dir, - source_file_name=filename, - chunks=chunks, - parse_track="text_track", - upload_assets=True, - ) - trace["stages"]["db_publish"] = { - "job_id": publish_result.job_id, - "document_id": publish_result.document_id, - } - - # ── Final trace ────────────────────────────────────────────────────────── - trace["total_elapsed_s"] = round(time.time() - t_start, 1) - trace["token_usage"] = get_current_token_tracker() - _write_json(out_dir / "trace.json", trace) + from app.services.document_parser.orchestration.postprocess import ( + apply_parse_postprocess, + ) + parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) + chunks, full_elapsed = _finalize_df(out_dir, filename, parsed_df) + trace["stages"]["full"] = { + "elapsed_s": round(full_elapsed, 1), + "chunk_count": len(chunks), + } + prev_usage = _record_token_stage( + ledger, "full", prev=prev_usage, out_dir=out_dir + ) - logger.info("") - logger.info("═" * 70) - logger.info(" ✅ DONE in {:.1f}s → {}", time.time() - t_start, out_dir) - logger.info("═" * 70) - return 0 + else: + logger.error("Unsupported format: {}", ext) + return 1 + + # ── Optional DB publication ────────────────────────────────────────── + if args.run_db: + from scripts._debug_publish import publish_debug_result_dir + publish_result = publish_debug_result_dir( + result_dir=out_dir, + source_file_name=filename, + chunks=chunks, + parse_track="text_track", + upload_assets=True, + ) + trace["stages"]["db_publish"] = { + "job_id": publish_result.job_id, + "document_id": publish_result.document_id, + } + + trace["total_elapsed_s"] = round(time.time() - t_start, 1) + logger.info("") + logger.info("═" * 70) + logger.info(" ✅ DONE in {:.1f}s → {}", time.time() - t_start, out_dir) + logger.info("═" * 70) + return 0 + finally: + remainder = token_usage_delta( + prev_usage, + deepcopy(get_current_token_tracker() or {}), + ) + usage = _aggregate_token_ledger( + ledger, + remainder=remainder if remainder.get("calls") or remainder.get("total_tokens") else None, + ) + _apply_token_usage_to_outputs(out_dir, trace, usage) + cleanup_token_tracker() if __name__ == "__main__": diff --git a/apps/worker/scripts/page_memory/_debug_pm_shared.py b/apps/worker/scripts/page_memory/_debug_pm_shared.py index 8699ee7c1..43c913d45 100644 --- a/apps/worker/scripts/page_memory/_debug_pm_shared.py +++ b/apps/worker/scripts/page_memory/_debug_pm_shared.py @@ -26,6 +26,7 @@ WORKER_ROOT = ROOT / "apps" / "worker" sys.path.insert(0, str(WORKER_ROOT)) sys.path.insert(0, str(ROOT / "packages" / "shared-python")) +sys.path.insert(0, str(WORKER_ROOT / "scripts")) from dotenv import load_dotenv @@ -50,10 +51,12 @@ serialize_scope_skeletons as _serialize_scope_skeletons, ) from shared.services.ai.token_tracking import ( + cleanup_token_tracker, init_token_tracker, get_current_token_tracker, ) from shared.services.ai.token_costing import build_token_cost_estimate +from _debug_token_ledger import empty_token_usage, merge_token_usage, token_usage_delta, write_stage_ledger # Re-exported for staged debug scripts. Listing them here marks the imports as # intentional so CodeQL does not treat them as unused. @@ -76,62 +79,17 @@ # ── Token cost tracker ──────────────────────────────────────────────────────── -def _usage_delta(prev: dict[str, Any], cur: dict[str, Any]) -> dict[str, Any]: - _NUM = ("prompt_tokens", "completion_tokens", "total_tokens", "calls") - - def _sub(a: dict, b: dict) -> dict: - return {f: int(b.get(f, 0)) - int(a.get(f, 0)) for f in _NUM} - - def _bucket(pa: dict, pb: dict) -> dict: - r: dict[str, Any] = {} - for k in set(pa) | set(pb): - pk, ck = pa.get(k, {}), pb.get(k, {}) - if not isinstance(pk, dict) or not isinstance(ck, dict): - continue - e = {f: v for f, v in _sub(pk, ck).items() if v} - pm, cm = pk.get("models", {}), ck.get("models", {}) - if pm or cm: - md = _bucket(pm, cm) - if md: - e["models"] = md - if e: - r[k] = e - return r - - d = _sub(prev, cur) - for bk in ("by_model", "by_task"): - bd = _bucket(prev.get(bk, {}), cur.get(bk, {})) - if bd: - d[bk] = bd - return d - - class TokenCostTracker: """Incremental token usage & cost tracker for debug pipeline stages.""" def __init__(self) -> None: self._dict = init_token_tracker() - self._root_gid = self._gid() self._prev: dict[str, Any] = deepcopy(self._dict) self._stages: list[dict[str, Any]] = [] - @staticmethod - def _gid() -> int: - from shared.services.ai.token_tracking import _current_greenlet_id - - return _current_greenlet_id() - - def register_child_thread(self) -> None: - from shared.services.ai.token_tracking import _root_ids, _lock - - gid = self._gid() - if gid != self._root_gid: - with _lock: - _root_ids[gid] = self._root_gid - def snapshot_stage(self, stage: str) -> None: cur = deepcopy(get_current_token_tracker() or {}) - delta = _usage_delta(self._prev, cur) + delta = token_usage_delta(self._prev, cur) self._stages.append({ "stage": stage, "prompt_tokens": delta.get("prompt_tokens", 0), @@ -953,20 +911,6 @@ def _stage_usage_snapshot(tracker: TokenCostTracker | None) -> dict[str, Any]: } -def _merge_token_usage( - destination: dict[str, Any], - source: dict[str, Any], -) -> None: - """Merge raw production token-tracker snapshots recursively.""" - for key, value in source.items(): - if isinstance(value, dict): - child = destination.setdefault(str(key), {}) - if isinstance(child, dict): - _merge_token_usage(child, value) - elif isinstance(value, int | float) and not isinstance(value, bool): - destination[str(key)] = destination.get(str(key), 0) + value - - def record_stage_cost( out_dir: Path, *, @@ -1008,13 +952,7 @@ def record_stage_cost( ledger["updated_at"] = updated_at path = stage_costs_path(out_dir) - path.parent.mkdir(parents=True, exist_ok=True) - temp_path = path.with_suffix(f"{path.suffix}.tmp") - temp_path.write_text( - json.dumps(ledger, ensure_ascii=False, indent=2), - encoding="utf-8", - ) - temp_path.replace(path) + write_stage_ledger(path, ledger) logger.info( " stage cost → {} ({} {:.1f}s / ${:.6f})", path, @@ -1034,14 +972,7 @@ def aggregate_stage_costs(ledger: dict[str, Any]) -> dict[str, Any]: by_pipeline_stage: dict[str, Any] = {} by_substage: list[dict[str, Any]] = [] merged_trace_stages: list[dict[str, Any]] = [] - usage = { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0, - "calls": 0, - "by_model": {}, - "by_task": {}, - } + usage = empty_token_usage() elapsed_s = 0.0 for stage_key in _COST_STAGE_KEYS: @@ -1060,7 +991,7 @@ def aggregate_stage_costs(ledger: dict[str, Any]) -> dict[str, Any]: raw_usage = row.get("token_usage") if isinstance(raw_usage, dict): # Numeric top-level fields were already added above. - _merge_token_usage( + merge_token_usage( usage, { key: value @@ -1200,38 +1131,41 @@ def stop_with_trace( extra_summary: dict[str, Any] | None = None, ) -> int: """Write TRACE.JSON, optionally recording this pipeline stage's cost ledger.""" - if pipeline_stage is not None and elapsed_s is not None: - record_stage_cost( - out_dir, - pipeline_stage=pipeline_stage, - elapsed_s=elapsed_s, - token_cost_tracker=token_cost_tracker, - trace_stages=stages, - stop_at=stop_at, - ) + try: + if pipeline_stage is not None and elapsed_s is not None: + record_stage_cost( + out_dir, + pipeline_stage=pipeline_stage, + elapsed_s=elapsed_s, + token_cost_tracker=token_cost_tracker, + trace_stages=stages, + stop_at=stop_at, + ) - aggregated = aggregate_stage_costs(load_stage_costs(out_dir)) - merged_stages = aggregated.get("stages") or stages - summary: dict[str, Any] = { - "page_count": page_count, - "scope_id": scope_id, - "rows_count": None, - "elapsed_s": aggregated.get("elapsed_s"), - "completed_pipeline_stages": aggregated.get("completed_pipeline_stages"), - "token_cost": aggregated.get("token_cost"), - } - if extra_summary: - summary.update(jsonable(extra_summary)) - - write_trace( - out_dir=out_dir, - stages=merged_stages, - final_status=final_status or f"stopped_at_{stop_at}", - summary=summary, - ) - remove_nested_doc_agent_trace(out_dir) - maybe_purge_debug_visuals(out_dir) - return 0 + aggregated = aggregate_stage_costs(load_stage_costs(out_dir)) + merged_stages = aggregated.get("stages") or stages + summary: dict[str, Any] = { + "page_count": page_count, + "scope_id": scope_id, + "rows_count": None, + "elapsed_s": aggregated.get("elapsed_s"), + "completed_pipeline_stages": aggregated.get("completed_pipeline_stages"), + "token_cost": aggregated.get("token_cost"), + } + if extra_summary: + summary.update(jsonable(extra_summary)) + + write_trace( + out_dir=out_dir, + stages=merged_stages, + final_status=final_status or f"stopped_at_{stop_at}", + summary=summary, + ) + remove_nested_doc_agent_trace(out_dir) + maybe_purge_debug_visuals(out_dir) + return 0 + finally: + cleanup_token_tracker() def remove_nested_doc_agent_trace(out_dir: Path) -> None: diff --git a/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py b/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py index 0b28900a3..bc885c352 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py @@ -127,8 +127,6 @@ def _run_fine_hierarchy_for_scope( ) scope_stages: list[dict[str, Any]] = [] - if token_cost_tracker is not None: - token_cost_tracker.register_child_thread() skel_path = scope_dir / "skeletons.json" require_file(skel_path, hint=f"Stage 3 should have created {skel_path}") diff --git a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py index afa9d7183..079a3349e 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py @@ -73,15 +73,12 @@ def _run_tagging_for_scope( page_features: list[Any], toc_pages: list[int], args: Any, - token_cost_tracker: TokenCostTracker | None = None, ) -> ScopeResult: """Load Stage-3 combined tags and rehydrate renders for final assembly.""" from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.page_renderer import render_document_pages scope_stages: list[dict[str, Any]] = [] - if token_cost_tracker is not None: - token_cost_tracker.register_child_thread() fine_hierarchy_path = scope_dir / "fine_hierarchy.json" require_file(fine_hierarchy_path, hint=f"Run Stage 3 to produce {fine_hierarchy_path}") @@ -329,7 +326,6 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: page_features=page_features, toc_pages=toc_pages, args=args, - token_cost_tracker=token_cost_tracker, ) if args.max_workers > 1 and len(scope_ids) > 1: diff --git a/apps/worker/tests/contract/test_calibration_scan_contract.py b/apps/worker/tests/contract/test_calibration_scan_contract.py index b28a114cb..72e053aa6 100644 --- a/apps/worker/tests/contract/test_calibration_scan_contract.py +++ b/apps/worker/tests/contract/test_calibration_scan_contract.py @@ -301,3 +301,40 @@ def _capture_inspect(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: assert args["page_cap"] == 1 assert len(args["rendered_pages"]) == 1 assert args["rendered_pages"][0]["page"] == args["pages"][0] + + +def test_concurrent_inspect_records_tokens_on_parse_tracker( + patch_inspect, +) -> None: + from shared.services.ai.token_tracking import ( + cleanup_token_tracker, + init_token_tracker, + record_tokens, + ) + + class _RecordingInspect(_FakeInspect): + def __call__(self, ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + record_tokens( + { + "prompt_tokens": 3, + "completion_tokens": 1, + "total_tokens": 4, + }, + model="test-vlm", + task="calibration.scan_title_forward", + ) + return super().__call__(ctx, args) + + tracker = init_token_tracker() + try: + patch_inspect(_RecordingInspect(hit_page=10)) + result = scan_title_forward( + ctx=_ctx(), title="Chapter 1", start_page=10, page_count=60 + ) + assert result.found_page == 10 + assert tracker["calls"] == 2 + assert tracker["total_tokens"] == 8 + assert tracker["by_task"]["calibration.scan_title_forward"]["total_tokens"] == 8 + assert tracker["by_model"]["test-vlm"]["total_tokens"] == 8 + finally: + cleanup_token_tracker() diff --git a/apps/worker/tests/contract/test_parse_usage_task_contract.py b/apps/worker/tests/contract/test_parse_usage_task_contract.py new file mode 100644 index 000000000..c108e21d5 --- /dev/null +++ b/apps/worker/tests/contract/test_parse_usage_task_contract.py @@ -0,0 +1,67 @@ +"""Parse-pipeline LLM calls must carry an explicit usage_task.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +_WORKER_SERVICES = Path(__file__).resolve().parents[2] / "app" / "services" +_WORKER_SCRIPTS = Path(__file__).resolve().parents[2] / "scripts" +_PARSE_ROOTS = ( + _WORKER_SERVICES / "document_parser", + _WORKER_SERVICES / "document_agent", + _WORKER_SERVICES / "page_memory", + _WORKER_SERVICES / "connect_builder", + _WORKER_SCRIPTS, +) +_LLM_CALL_NAMES = { + "chat_completion", + "chat_completion_with_usage", + "chat_completion_raw_with_usage", + "summarize", + "transcribe", +} + + +def _call_name(node: ast.Call) -> str | None: + func = node.func + if isinstance(func, ast.Name): + return func.id + if isinstance(func, ast.Attribute): + return func.attr + return None + + +def _iter_parse_python_files() -> list[Path]: + files: list[Path] = [] + for root in _PARSE_ROOTS: + if not root.exists(): + continue + files.extend(path for path in root.rglob("*.py") if path.is_file()) + return files + + +def _relative_path(path: Path) -> str: + for anchor in (_WORKER_SERVICES, _WORKER_SCRIPTS): + try: + return str(path.relative_to(anchor)) + except ValueError: + continue + return str(path) + + +def test_parse_llm_calls_pass_explicit_usage_task() -> None: + missing: list[str] = [] + for path in _iter_parse_python_files(): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + name = _call_name(node) + if name not in _LLM_CALL_NAMES: + continue + if any(keyword.arg == "usage_task" for keyword in node.keywords): + continue + rel = _relative_path(path) + missing.append(f"{rel}:{node.lineno}:{name}") + assert missing == [] diff --git a/apps/worker/tests/unit/test_debug_text_track_token_ledger.py b/apps/worker/tests/unit/test_debug_text_track_token_ledger.py new file mode 100644 index 000000000..4d7078322 --- /dev/null +++ b/apps/worker/tests/unit/test_debug_text_track_token_ledger.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +_SCRIPTS_ROOT = Path(__file__).resolve().parents[2] / "scripts" +sys.path.insert(0, str(_SCRIPTS_ROOT)) + +from _debug_token_ledger import ( + aggregate_stage_deltas, + empty_token_usage, + token_usage_delta, +) + +_TEXT_TRACK_STAGES = ("profile", "mineru", "hierarchy", "full") + + +def test_token_usage_delta_and_reuse_merge_match_stage_costs_semantics() -> None: + prev = empty_token_usage() + after_profile = { + "prompt_tokens": 10, + "completion_tokens": 4, + "total_tokens": 14, + "calls": 1, + "by_model": {"vlm": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14, "calls": 1}}, + "by_task": { + "document_agent.coarse_profile": { + "prompt_tokens": 10, + "completion_tokens": 4, + "total_tokens": 14, + "calls": 1, + } + }, + } + profile_delta = token_usage_delta(prev, after_profile) + assert profile_delta["total_tokens"] == 14 + assert profile_delta["by_task"]["document_agent.coarse_profile"]["calls"] == 1 + + # reuse-profile starts a fresh tracker; only the new stage exists in this run. + reuse_run = token_usage_delta(empty_token_usage(), { + "prompt_tokens": 8, + "completion_tokens": 2, + "total_tokens": 10, + "calls": 1, + "by_model": { + "text": {"prompt_tokens": 8, "completion_tokens": 2, "total_tokens": 10, "calls": 1} + }, + "by_task": { + "parser.heading_hierarchy": { + "prompt_tokens": 8, + "completion_tokens": 2, + "total_tokens": 10, + "calls": 1, + } + }, + }) + ledger = { + "stages": { + "profile": {"token_usage": profile_delta}, + "hierarchy": {"token_usage": reuse_run}, + } + } + merged = aggregate_stage_deltas(ledger, _TEXT_TRACK_STAGES) + assert merged["total_tokens"] == 24 + assert merged["calls"] == 2 + assert merged["by_task"]["document_agent.coarse_profile"]["total_tokens"] == 14 + assert merged["by_task"]["parser.heading_hierarchy"]["total_tokens"] == 10 diff --git a/apps/worker/tests/unit/test_processing_metadata_persist.py b/apps/worker/tests/unit/test_processing_metadata_persist.py index 4cd3d8200..340426e49 100644 --- a/apps/worker/tests/unit/test_processing_metadata_persist.py +++ b/apps/worker/tests/unit/test_processing_metadata_persist.py @@ -104,3 +104,52 @@ def test_record_processing_completion_persists_token_usage_to_job_row( assert updates["stages"]["token_usage"]["total_tokens"] == 10 assert "processing_completed_at" in updates assert "processing_duration_ms" in updates + + +def test_run_parse_job_failure_persists_token_usage(monkeypatch: object) -> None: + import app.services.document_ingestion.processing_run as processing_run + + persisted: dict[str, object] = {} + + def fake_persist(*, job_id: str, job_context: object, metadata_updates: dict[str, object]) -> None: + persisted.update(metadata_updates) + + monkeypatch.setattr(processing_run, "persist_job_metadata_updates", fake_persist) + monkeypatch.setattr( + processing_run, + "init_token_tracker", + lambda: {"prompt_tokens": 6, "completion_tokens": 1, "total_tokens": 7, "calls": 1}, + ) + monkeypatch.setattr(processing_run, "init_stage_tracker", lambda: {"worker.parse.document": 12}) + monkeypatch.setattr(processing_run, "init_llm_overrides", lambda *_args: None) + monkeypatch.setattr(processing_run, "cleanup_llm_overrides", lambda: None) + monkeypatch.setattr(processing_run, "cleanup_token_tracker", lambda: None) + monkeypatch.setattr(processing_run, "cleanup_stage_tracker", lambda: None) + + def _boom(**_kwargs: object) -> object: + raise RuntimeError("parse failed after LLM work") + + monkeypatch.setattr(processing_run, "prepare_source_file", _boom) + + lifecycle = Mock() + job_context = _job_context() + try: + processing_run._run_parse_job( + job_id="job-fail", + job_context=job_context, + lifecycle_service=lifecycle, + task_workspace=SimpleNamespace( + input_dir="/tmp", + output_dir="/tmp", + root_dir="/tmp", + ), + ) + except RuntimeError as exc: + assert "parse failed after LLM work" in str(exc) + else: + raise AssertionError("expected parse failure") + + stages = persisted["stages"] + assert isinstance(stages, dict) + assert stages["token_usage"]["total_tokens"] == 7 + assert stages["timing_ms"]["worker.parse.document"] == 12 diff --git a/packages/shared-python/shared/services/ai/summary/engine.py b/packages/shared-python/shared/services/ai/summary/engine.py index 76e4129d0..970a08657 100644 --- a/packages/shared-python/shared/services/ai/summary/engine.py +++ b/packages/shared-python/shared/services/ai/summary/engine.py @@ -7,7 +7,6 @@ - prompt construction (via the shared ``build_prompt`` registry), - the text-or-vision LLM call (with optional page/asset images), - JSON parsing with one retry, -- budget reserve/commit/refund for visual calls, - deterministic language locking for text input. Two public functions: @@ -38,12 +37,10 @@ BodySummary, Entity, ) -from shared.utils.token_estimate import estimate_tokens SummaryMode = Literal["text", "page", "asset"] _MAX_JSON_RETRIES = 1 -_IMAGE_TOKEN_EST = 800 def _read_image_b64(image_path: str) -> str | None: @@ -109,17 +106,12 @@ def _call_llm( image_paths: list[str], usage_task: str, expect_json: bool, - budget: Any | None, - budget_pool: str, - budget_stage: str | None, channel: Literal["text", "vision"] = "text", ) -> Any | None: - """One text-or-vision call with budget accounting and a single JSON retry. + """One text-or-vision call with a single JSON retry. Returns the parsed object (``expect_json``) or the raw string, or ``None`` on - failure / exhausted budget. Budget is reserved before the call, committed on - success, refunded on failure — matching the prior per-caller bookkeeping but - in one place. + failure. Token usage is recorded by the shared client into the parse tracker. ``channel`` selects BYOK text vs vision credentials when overrides are active. """ @@ -140,13 +132,6 @@ def _call_llm( # Every requested image failed to load; nothing to send. return None - est = estimate_tokens(prompt) + _IMAGE_TOKEN_EST * max(1, len(image_paths)) - if budget is not None and not budget.try_reserve( - budget_pool, est, stage=budget_stage - ): - logger.debug("[summary] budget exhausted for task {}", usage_task) - return None - api_kwargs: dict[str, Any] = {} if expect_json: api_kwargs["response_format"] = {"type": "json_object"} @@ -154,8 +139,6 @@ def _call_llm( resolve = resolve_vision if channel == "vision" else resolve_text effective_model, api_key, api_url = resolve(model) if not effective_model: - if budget is not None: - budget.refund(budget_pool, est=est, stage=budget_stage) return None client = _client_mod.get_openai_client( @@ -165,7 +148,7 @@ def _call_llm( ) for attempt in range(_MAX_JSON_RETRIES + 1): try: - raw, usage = client.chat_completion_with_usage( + raw, _ = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=effective_model, temperature=temperature, @@ -174,14 +157,6 @@ def _call_llm( usage_task=usage_task, **api_kwargs, ) - if budget is not None: - budget.commit( - budget_pool, - actual=usage.get("total_tokens", est), - est=est, - stage=budget_stage, - ) - budget = None # commit once even across the retry loop if not expect_json: if isinstance(raw, str) and raw.strip().lower() in ("null", "none"): return None @@ -199,15 +174,11 @@ def _call_llm( continue return None except UnavailableException: - if budget is not None: - budget.refund(budget_pool, est=est, stage=budget_stage) if usage_task.startswith("page_memory."): raise return None except Exception as exc: logger.warning("[summary] LLM call failed for {}: {}", usage_task, exc) - if budget is not None: - budget.refund(budget_pool, est=est, stage=budget_stage) return None return None @@ -222,9 +193,6 @@ def summarize( max_keywords: int = ..., model: str | None = ..., usage_task: str | None = ..., - budget: Any | None = ..., - budget_pool: str = ..., - budget_stage: str | None = ..., asset_title_hint: str = ..., prompt_task: str | None = ..., prompt_paras: dict[str, Any] | None = ..., @@ -242,9 +210,6 @@ def summarize( max_keywords: int = ..., model: str | None = ..., usage_task: str | None = ..., - budget: Any | None = ..., - budget_pool: str = ..., - budget_stage: str | None = ..., asset_title_hint: str = ..., prompt_task: str | None = ..., prompt_paras: dict[str, Any] | None = ..., @@ -261,13 +226,6 @@ def summarize( max_keywords: int = 5, model: str | None = None, usage_task: str | None = None, - # TODO(parse-budget-cleanup): no live caller passes a non-None budget after - # PROFILE BudgetTracker removal. Drop budget/budget_pool/budget_stage once - # remaining formats stop needing this duck-typed hook, or redirect any - # future limit to token_tracking instead. - budget: Any | None = None, - budget_pool: str = "visual", - budget_stage: str | None = None, asset_title_hint: str = "", prompt_task: str | None = None, prompt_paras: dict[str, Any] | None = None, @@ -283,9 +241,6 @@ def summarize( image_paths: Page or asset image(s). Required for ``page``/``asset`` modes that render from an image; ignored for plain ``text``. - budget: - Optional external reservation ledger. Visual calls reserve from - ``budget_stage``. prompt_task / prompt_paras: Override the prompt used for the image-based page path. Lets a bounded node summary (``page-memory-node-summary`` with ``node_title`` / @@ -304,9 +259,6 @@ def summarize( summary_len=summary_len, model=model, usage_task=usage_task or "summary.asset", - budget=budget, - budget_pool=budget_pool, - budget_stage=budget_stage, asset_title_hint=asset_title_hint, ) return _summarize_body( @@ -317,9 +269,6 @@ def summarize( max_keywords=max_keywords, model=model, usage_task=usage_task or f"summary.{mode}", - budget=budget, - budget_pool=budget_pool, - budget_stage=budget_stage, prompt_task=prompt_task, prompt_paras=prompt_paras, ) @@ -334,9 +283,6 @@ def _summarize_body( max_keywords: int, model: str | None, usage_task: str, - budget: Any | None, - budget_pool: str, - budget_stage: str | None, prompt_task: str | None = None, prompt_paras: dict[str, Any] | None = None, ) -> BodySummary: @@ -359,9 +305,6 @@ def _summarize_body( image_paths=image_paths, usage_task=usage_task, expect_json=True, - budget=budget, - budget_pool=budget_pool, - budget_stage=budget_stage, channel="vision", ) else: @@ -389,9 +332,6 @@ def _summarize_body( image_paths=[], usage_task=usage_task, expect_json=True, - budget=budget, - budget_pool="plan", - budget_stage=None, channel="text", ) @@ -412,9 +352,6 @@ def _summarize_asset( summary_len: int, model: str | None, usage_task: str, - budget: Any | None, - budget_pool: str, - budget_stage: str | None, asset_title_hint: str, ) -> AssetSummary: if not image_paths and not text.strip(): @@ -440,9 +377,6 @@ def _summarize_asset( image_paths=[], usage_task=usage_task, expect_json=False, - budget=budget, - budget_pool="plan", - budget_stage=None, channel="text", ) if isinstance(raw, str) and raw.strip(): @@ -469,9 +403,6 @@ def _summarize_asset( image_paths=image_paths, usage_task=usage_task, expect_json=True, - budget=budget, - budget_pool=budget_pool, - budget_stage=budget_stage, channel="vision", ) if isinstance(parsed, dict): @@ -490,9 +421,6 @@ def transcribe( model: str | None = None, max_tokens: int = 1500, usage_task: str = "summary.transcribe", - budget: Any | None = None, - budget_pool: str = "visual", - budget_stage: str | None = None, ) -> str: """Single OCR primitive (§4.2): transcribe page/image text verbatim. @@ -518,9 +446,6 @@ def transcribe( image_paths=image_paths, usage_task=usage_task, expect_json=True, - budget=budget, - budget_pool=budget_pool, - budget_stage=budget_stage, channel="vision", ) if isinstance(parsed, dict): From 1390d4bf9fedd52a1f918e25fd492e1fade24e9c Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 27 Aug 2026 21:49:05 +0800 Subject: [PATCH 6/7] fix: stub pending TOC calibration internals in rehome contract Patching _calibrate_pending_tocs by string path was unreliable in the full contract suite, so the real path tried to render a missing pdf and dropped the pending graft record. Align with toc_graft stubs instead. Co-authored-by: Cursor --- .../contract/test_toc_rehome_contract.py | 47 ++++++++++--------- 1 file changed, 25 insertions(+), 22 deletions(-) diff --git a/apps/worker/tests/contract/test_toc_rehome_contract.py b/apps/worker/tests/contract/test_toc_rehome_contract.py index 77b6ec3df..93284f88e 100644 --- a/apps/worker/tests/contract/test_toc_rehome_contract.py +++ b/apps/worker/tests/contract/test_toc_rehome_contract.py @@ -586,6 +586,13 @@ def _fake_anchor(**kwargs): def test_run_toc_anchoring_rehome_before_classify_keeps_contained_graft() -> None: + """Pending calibrate must succeed without opening pdf_path; graft after rehome. + + Stub the calibration *internals* (same pattern as toc_graft contracts), not + ``_calibrate_pending_tocs`` itself. Patching that helper by string path is + brittle across the full contract suite and lets the real pending path try to + render ``pdf_path`` (``/tmp/rehome.pdf``), then drop the pending record. + """ hierarchies = [ { "toc_range": [1, 1], @@ -619,26 +626,14 @@ def _fake_anchor(**kwargs): ) return nodes, anchor - def _fake_pending(**kwargs): - return [ - { - "toc": hierarchies[1], - "nodes": [ - serialize_title_node( - TitleNode(title="Inner", level=1, printed_page=12) - ) - ], - "skeleton_anchor": serialize_skeleton_anchor( - SkeletonAnchor( - offset=0, - offset_status="ok", - match_overrides=_overrides({("Inner",): 12}), - null_page_report=[], - bulk_count=1, - ) - ), - } - ] + inner = TitleNode(title="Inner", level=1, printed_page=12) + inner_anchor = SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides=_overrides({("Inner",): 12}), + null_page_report=[], + bulk_count=1, + ) with ( patch( @@ -646,8 +641,16 @@ def _fake_pending(**kwargs): side_effect=_fake_anchor, ), patch( - "app.services.document_agent.structure.toc_anchoring._calibrate_pending_tocs", - side_effect=_fake_pending, + "app.services.document_agent.calibration.service.calibrate_offset", + return_value=object(), + ), + patch( + "app.services.document_agent.calibration.procedure.pick_primary_offset", + return_value=0, + ), + patch( + "app.services.document_agent.calibration.procedure.finalize_calibration_result", + return_value=([inner], inner_anchor, True), ), ): run_toc_anchoring(ctx) From f7d9a02c846256344c6fb9b6350c69abd2a46df1 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 27 Aug 2026 21:59:01 +0800 Subject: [PATCH 7/7] fix: drop unused serialize imports from toc rehome contract Co-authored-by: Cursor --- apps/worker/tests/contract/test_toc_rehome_contract.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/apps/worker/tests/contract/test_toc_rehome_contract.py b/apps/worker/tests/contract/test_toc_rehome_contract.py index 93284f88e..9e4f4b87f 100644 --- a/apps/worker/tests/contract/test_toc_rehome_contract.py +++ b/apps/worker/tests/contract/test_toc_rehome_contract.py @@ -16,8 +16,6 @@ from app.services.document_agent.state import ProfileBlackboard from app.services.document_agent.structure.anchoring_primitives import ( SkeletonAnchor, - serialize_skeleton_anchor, - serialize_title_node, ) from app.services.document_agent.structure.hierarchy_locator import ( TitleMatch,