From 81427b096b4ee3bb3b9bdc8418ca0367fd173766 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 13 Aug 2026 15:04:32 +0800 Subject: [PATCH 1/9] feat: enhance TOC anchoring and state management - Introduced `_clear_toc_anchor_state` method to reset TOC anchor-related attributes in `ProfileCoordinator`. - Updated `PageAnatomyMap` and `AgentBlackboard` to include new fields for `skeleton_anchor`, `skeleton_nodes`, and `pending_skeleton_anchors`. - Integrated `run_toc_anchoring` in the extraction pipeline to improve TOC processing. - Added serialization and deserialization functions for `skeleton_anchor` and `title_node` to support new data structures. - Refactored `extract_section_skeletons` to utilize persisted `skeleton_anchor` for section boundary resolution. --- .../services/document_agent/coordinator.py | 12 + .../app/services/document_agent/manifest.py | 6 + .../app/services/document_agent/state.py | 3 + .../structure/anchoring_primitives.py | 38 ++ .../document_agent/structure/toc_anchoring.py | 332 ++++++++++++++ .../tools/persist_anatomy_map.py | 3 + .../tools/propose_shard_plan.py | 45 +- .../services/page_memory/memory_service.py | 42 -- .../page_memory/skeleton_extractor.py | 397 ++++------------ ...profile_skeleton_anchor_wiring_contract.py | 429 ++++++++++++++++++ 10 files changed, 920 insertions(+), 387 deletions(-) create mode 100644 apps/worker/app/services/document_agent/structure/toc_anchoring.py create mode 100644 apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 00caba70a..ce8715f70 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -26,6 +26,7 @@ from app.services.document_agent.planner import ProfilePlanner from app.services.document_agent.registry import REGISTRY from app.services.document_agent.state import AgentBlackboard, DocumentAgentState +from app.services.document_agent.structure.toc_anchoring import run_toc_anchoring from app.services.document_agent import tools as _registered_tools # noqa: F401 from app.services.document_agent.trace import ParseRunRecorder from app.services.document_agent.validators import single_shard_plan @@ -125,6 +126,7 @@ def run_toc(self) -> TocResult: failure_kind="degraded", ) self.blackboard.toc_hierarchies = None + self._clear_toc_anchor_state() return self.blackboard.toc_result def run_lightweight_anatomy( @@ -333,6 +335,7 @@ def _ensure_disabled_toc_placeholder(self) -> None: notes="TOC profiling disabled by PDF_PROFILE_TOC_ENABLED", ) self.blackboard.toc_hierarchies = None + self._clear_toc_anchor_state() self.blackboard.global_signals["toc_profile_attempted"] = False def _ensure_toc_profile(self, *, strict: bool) -> None: @@ -340,6 +343,7 @@ def _ensure_toc_profile(self, *, strict: bool) -> None: if strict and self._toc_result_requires_strict_retry(): self.blackboard.toc_result = None self.blackboard.toc_hierarchies = None + self._clear_toc_anchor_state() should_run = True if not should_run: @@ -361,6 +365,7 @@ def _ensure_toc_profile(self, *, strict: bool) -> None: failure_kind="degraded", ) self.blackboard.toc_hierarchies = None + self._clear_toc_anchor_state() return if self.blackboard.toc_result is None: @@ -403,10 +408,17 @@ def _dispatch_profile_tool(self, *, tool_name: str, actor: str) -> ToolResult: self.round_index += 1 return result + def _clear_toc_anchor_state(self) -> None: + self.blackboard.toc_page_offset = None + self.blackboard.skeleton_anchor = None + self.blackboard.skeleton_nodes = None + self.blackboard.pending_skeleton_anchors = [] + def _run_toc_extraction_pipeline(self) -> None: for tool_name in ("find.toc_anchor_pages", "extract.toc_with_boundaries"): self._dispatch_profile_tool( tool_name=tool_name, actor=f"toc:{tool_name}", ) + run_toc_anchoring(self.ctx) diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index ad9367150..e545c63c0 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -231,6 +231,9 @@ class PageAnatomyMap: document_profile: DocumentProfile | None = None toc_hierarchies: list[dict[str, Any]] | None = None toc_page_offset: int | None = None + skeleton_anchor: dict[str, Any] | None = None + skeleton_nodes: list[dict[str, Any]] | None = None + pending_skeleton_anchors: list[dict[str, Any]] = field(default_factory=list) global_signals: dict[str, Any] = field(default_factory=dict) trace_summary: dict[str, Any] = field(default_factory=dict) created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) @@ -252,6 +255,9 @@ def to_dict(self) -> dict[str, Any]: if self.document_profile else None, "toc_page_offset": self.toc_page_offset, + "skeleton_anchor": self.skeleton_anchor, + "skeleton_nodes": self.skeleton_nodes, + "pending_skeleton_anchors": list(self.pending_skeleton_anchors), "global_signals": dict(self.global_signals), "trace_summary": dict(self.trace_summary), "created_at": self.created_at.isoformat(), diff --git a/apps/worker/app/services/document_agent/state.py b/apps/worker/app/services/document_agent/state.py index c717cbf07..93ddfd83d 100644 --- a/apps/worker/app/services/document_agent/state.py +++ b/apps/worker/app/services/document_agent/state.py @@ -38,6 +38,9 @@ class AgentBlackboard: toc_hierarchies: list[dict[str, Any]] | None = None h1_result: H1BoundaryResult | None = None toc_page_offset: int | None = None + skeleton_anchor: dict[str, Any] | None = None + skeleton_nodes: list[dict[str, Any]] | None = None + pending_skeleton_anchors: list[dict[str, Any]] = field(default_factory=list) shard_plan: ShardPlan | None = None validation_report: dict[str, Any] | None = None verdict: AgentVerdict | None = None diff --git a/apps/worker/app/services/document_agent/structure/anchoring_primitives.py b/apps/worker/app/services/document_agent/structure/anchoring_primitives.py index 390e8ba33..d0865bde2 100644 --- a/apps/worker/app/services/document_agent/structure/anchoring_primitives.py +++ b/apps/worker/app/services/document_agent/structure/anchoring_primitives.py @@ -699,6 +699,44 @@ def deserialize_title_match(data: dict[str, Any]) -> TitleMatch: ) +def serialize_title_node(node: TitleNode) -> dict[str, Any]: + return { + "title": node.title, + "level": node.level, + "printed_page": node.printed_page, + "printed_label": node.printed_label, + "page_kind": node.page_kind, + "physical_page_hint": node.physical_page_hint, + "children": [serialize_title_node(child) for child in node.children], + } + + +def deserialize_title_node(data: dict[str, Any]) -> TitleNode: + children_raw = data.get("children") or [] + children = [ + deserialize_title_node(child) + for child in children_raw + if isinstance(child, dict) + ] + printed_page = data.get("printed_page") + physical_page_hint = data.get("physical_page_hint") + return TitleNode( + title=str(data.get("title") or ""), + level=int(data.get("level") or 1), + printed_page=None if printed_page is None else int(printed_page), + printed_label=data.get("printed_label") + if isinstance(data.get("printed_label"), str) + else None, + page_kind=data.get("page_kind") + if isinstance(data.get("page_kind"), str) + else None, + physical_page_hint=( + None if physical_page_hint is None else int(physical_page_hint) + ), + children=children, + ) + + def deserialize_skeleton_anchor(data: dict[str, Any]) -> SkeletonAnchor: raw_overrides = data.get("match_overrides") or {} overrides: dict[tuple[str, ...], TitleMatch] = {} diff --git a/apps/worker/app/services/document_agent/structure/toc_anchoring.py b/apps/worker/app/services/document_agent/structure/toc_anchoring.py new file mode 100644 index 000000000..1e774a91c --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/toc_anchoring.py @@ -0,0 +1,332 @@ +"""Profile-time TOC anchoring: run existing calibration after TOC extract.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from loguru import logger + +from app.services.document_agent.agents.calibration import calibrate_offset +from app.services.document_agent.agents.calibration.orchestrator import anchor_hierarchy +from app.services.document_agent.agents.calibration.procedure import ( + finalize_calibration_result, + pick_primary_offset, +) +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.pdf_text import read_page_texts +from app.services.document_agent.structure.anchoring_primitives import ( + serialize_skeleton_anchor, + serialize_title_node, + toc_range_end, + toc_range_start, +) +from app.services.document_agent.structure.hierarchy_locator import ( + ResolvedHierarchyRange, + TitleNode, + collapse_intermediate_single_child_chains, + extract_toc_nodes, + iter_leaf_title_nodes, + resolve_hierarchy_page_ranges, +) + +_FRONT_TOC_REGION_GAP_PAGES = 5 +_LOG_PREFIX = "[profile.toc_anchoring]" + + +def run_toc_anchoring(ctx: ToolContext) -> None: + """Anchor extracted TOC hierarchies onto the profile blackboard.""" + page_count = int(ctx.blackboard.page_count or 0) + hierarchies = list(ctx.blackboard.toc_hierarchies or []) + if page_count <= 0 or not hierarchies: + return + + if not ctx.blackboard.page_full_text_cache: + ctx.blackboard.page_full_text_cache = read_page_texts( + ctx.pdf_path, + list(range(1, page_count + 1)), + timeout=300, + ) + page_texts = dict(ctx.blackboard.page_full_text_cache) + filename = Path(ctx.pdf_path).name + primary, pending, _summary = select_global_toc_hierarchies( + hierarchies=hierarchies, + filename=filename, + ) + nodes = extract_toc_nodes(primary) + if not nodes: + return + + nodes = collapse_intermediate_single_child_chains(nodes) + toc_result = ctx.blackboard.toc_result + body_pages = body_pages_excluding_toc( + getattr(toc_result, "toc_pages", None), + page_count, + ) + pending_starts: list[int] = [] + for toc in pending: + start = toc_range_start(toc) + if start is not None: + pending_starts.append(start) + primary_page_count = page_count + primary_body_pages = body_pages + if pending_starts: + primary_page_count = min(pending_starts) - 1 + primary_body_pages = [ + page for page in body_pages if page <= primary_page_count + ] + + resolve_nodes, skeleton_anchor = anchor_hierarchy( + nodes=nodes, + toc_hierarchies=primary, + page_texts=page_texts, + body_pages=primary_body_pages, + page_count=page_count, + ctx=ctx, + ) + ctx.blackboard.skeleton_anchor = serialize_skeleton_anchor(skeleton_anchor) + ctx.blackboard.skeleton_nodes = [serialize_title_node(node) for node in resolve_nodes] + ctx.blackboard.toc_page_offset = skeleton_anchor.offset + pending_records: list[dict[str, Any]] = [] + if pending: + primary_ranges = resolve_hierarchy_page_ranges( + resolve_nodes, + page_count=primary_page_count, + page_texts=page_texts, + body_pages=primary_body_pages, + match_overrides=skeleton_anchor.match_overrides, + ) + pending_records = _anchor_pending_tocs( + pending_tocs=pending, + ctx=ctx, + page_texts=page_texts, + page_count=page_count, + body_pages=body_pages, + primary_ranges=primary_ranges, + ) + ctx.blackboard.pending_skeleton_anchors = pending_records + + +def select_global_toc_hierarchies( + *, + hierarchies: list[dict[str, Any]], + filename: str, +) -> tuple[list[dict[str, Any]] | None, list[dict[str, Any]], dict[str, Any]]: + """Split TOC hierarchies into primary (front cluster) and pending.""" + if len(hierarchies) <= 1: + return (hierarchies or None), [], {} + + page_based = [ + hierarchy + for hierarchy in hierarchies + if hierarchy.get("toc_range_unit") == "page" + and toc_range_start(hierarchy) is not None + ] + if not page_based or len(page_based) != len(hierarchies): + return hierarchies, [], {} + + sorted_items = sorted( + enumerate(hierarchies), + key=lambda item: toc_range_start(item[1]) or 0, + ) + selected_indices: set[int] = set() + pending_indices: list[int] = [] + cluster_end: int | None = None + + for original_index, hierarchy in sorted_items: + start = toc_range_start(hierarchy) + end = toc_range_end(hierarchy) + if start is None or end is None: + selected_indices.add(original_index) + continue + if cluster_end is None: + selected_indices.add(original_index) + cluster_end = end + continue + if start <= cluster_end + _FRONT_TOC_REGION_GAP_PAGES: + selected_indices.add(original_index) + cluster_end = max(cluster_end, end) + continue + pending_indices.append(original_index) + + selected = [ + hierarchy + for index, hierarchy in enumerate(hierarchies) + if index in selected_indices + ] + pending = [hierarchies[i] for i in pending_indices] + + if pending: + logger.info( + "{} toc split: primary={} pending={} filename={}", + _LOG_PREFIX, + len(selected), + len(pending), + filename, + ) + summary = { + "strategy": "front_cluster_with_pending", + "input_count": len(hierarchies), + "primary_count": len(selected), + "pending_count": len(pending), + } + return (selected or None), pending, summary + + +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] + + +def pending_toc_body_scope( + *, + pending_tocs: list[dict[str, Any]], + index: int, + page_count: int, + body_pages: list[int], +) -> tuple[int, list[int]]: + pending_toc = pending_tocs[index] + toc_end = toc_range_end(pending_toc) + toc_scope_start = (toc_end + 1) if toc_end is not None else None + next_starts: list[int] = [] + for j in range(index + 1, len(pending_tocs)): + start = toc_range_start(pending_tocs[j]) + if start is not None: + next_starts.append(start) + toc_scope_end = (min(next_starts) - 1) if next_starts else page_count + toc_body_pages = [ + page + for page in body_pages + if page <= toc_scope_end + and (toc_scope_start is None or page >= toc_scope_start) + ] + return toc_scope_end, toc_body_pages + + +def classify_toc_relationship( + *, + offset: int, + nodes: list[TitleNode], + primary_ranges: list[ResolvedHierarchyRange], + page_count: int, +) -> str: + """Classify a pending TOC as parallel or contained vs primary ranges. + + parallel: the pending TOC covers pages beyond the primary tree's *anchored* + content (i.e. the last explicitly-located section start page). + contained: the pending TOC's content falls strictly within a primary + section's explicitly-anchored range. + """ + leaves = [ + node for _, node in iter_leaf_title_nodes(nodes) if node.printed_page is not None + ] + if not leaves: + return "unresolvable" + + first_printed = leaves[0].printed_page + last_printed = leaves[-1].printed_page + if first_printed is None or last_printed is None: + return "unresolvable" + first_physical = first_printed + offset + last_physical = last_printed + offset + + if first_physical < 1 or first_physical > page_count: + return "unresolvable" + + if not primary_ranges: + return "parallel" + + # Use the last *start_page* among primary ranges as the boundary of + # explicitly-anchored content. The end_page of the last section is often + # extended to page_count by default and doesn't reflect real content coverage. + last_anchored_start = max( + (r.start_page for r in primary_ranges if r.start_page is not None), default=0 + ) + + if first_physical > last_anchored_start: + return "parallel" + + min_level = min(r.level for r in primary_ranges) + top_level_ranges = [r for r in primary_ranges if r.level == min_level] + for r in top_level_ranges: + if r.start_page and r.end_page: + if r.start_page <= first_physical and last_physical <= r.end_page: + return "contained" + + return "parallel" + + +def _anchor_pending_tocs( + *, + pending_tocs: list[dict[str, Any]], + ctx: ToolContext, + page_texts: dict[int, str], + page_count: int, + body_pages: list[int], + primary_ranges: list[ResolvedHierarchyRange], +) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + for i, pending_toc in enumerate(pending_tocs): + nodes = extract_toc_nodes([pending_toc]) + if not nodes: + continue + nodes = collapse_intermediate_single_child_chains(nodes) + toc_scope_end, toc_body_pages = pending_toc_body_scope( + pending_tocs=pending_tocs, + index=i, + page_count=page_count, + body_pages=body_pages, + ) + phase1 = calibrate_offset( + nodes=nodes, + toc_hierarchies=[pending_toc], + ctx=ctx, + page_texts=page_texts, + page_count=toc_scope_end, + ) + offset = pick_primary_offset(phase1) + if offset is None: + logger.info( + "{} pending TOC toc_range={}: calibration failed, skipping", + _LOG_PREFIX, + pending_toc.get("toc_range"), + ) + continue + relationship = classify_toc_relationship( + offset=offset, + nodes=nodes, + primary_ranges=primary_ranges, + page_count=page_count, + ) + if relationship == "unresolvable": + logger.info( + "{} pending TOC toc_range={}: unresolvable, skipping", + _LOG_PREFIX, + pending_toc.get("toc_range"), + ) + records.append( + { + "toc": pending_toc, + "relationship": relationship, + } + ) + continue + resolve_nodes, skeleton_anchor, _finalized = finalize_calibration_result( + result=phase1, + entries=list(pending_toc.get("toc_with_level") or []), + toc_hierarchies=[pending_toc], + ctx=ctx, + page_count=toc_scope_end, + page_texts=page_texts, + body_pages=toc_body_pages, + nodes=nodes, + ) + records.append( + { + "toc": pending_toc, + "relationship": relationship, + "nodes": [serialize_title_node(node) for node in resolve_nodes], + "skeleton_anchor": serialize_skeleton_anchor(skeleton_anchor), + } + ) + return records diff --git a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py index f5865d856..b7683f76b 100644 --- a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py @@ -40,6 +40,9 @@ def build_anatomy_map(ctx: ToolContext) -> PageAnatomyMap: document_profile=ctx.blackboard.document_profile, toc_hierarchies=ctx.blackboard.toc_hierarchies, toc_page_offset=ctx.blackboard.toc_page_offset, + skeleton_anchor=ctx.blackboard.skeleton_anchor, + skeleton_nodes=ctx.blackboard.skeleton_nodes, + pending_skeleton_anchors=list(ctx.blackboard.pending_skeleton_anchors), global_signals=ctx.blackboard.global_signals, trace_summary={ "budget": ctx.budget.snapshot(), diff --git a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index 4f5b91c5b..fd0e82e08 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -473,32 +473,6 @@ def _parse_llm_plan( return enabled, cuts, reason, rationale -def _deterministic_guardrail_plan( - *, - page_count: int, - min_pages: int, - max_pages: int, - leaf_pages: list[int], -) -> tuple[list[tuple[int, str, str, float]], str]: - cuts: list[tuple[int, str, str, float]] = [] - previous = 0 - while page_count - previous > max_pages: - target = previous + max_pages - eligible = [ - page for page in leaf_pages if previous + min_pages < page <= target - ] - if eligible: - chosen = max(eligible) - cut_page = chosen - 1 - cuts.append((cut_page, "h1_boundary", f"guardrail leaf node at page {chosen}", 0.35)) - previous = cut_page - else: - cut_page = previous + max_pages - cuts.append((cut_page, "forced_max_size", "no leaf node in range", 0.2)) - previous = cut_page - return cuts, "too_large" - - def _deterministic_chapter_plan( *, chapters: list[dict[str, Any]], @@ -616,24 +590,7 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) - offset_hint: int | None = None - if ctx.blackboard.toc_hierarchies: - from app.services.document_agent.agents.calibration import calibrate_offset - from app.services.document_agent.agents.calibration.procedure import ( - pick_primary_offset, - ) - from app.services.document_agent.structure.hierarchy_locator import extract_toc_nodes - - nodes = extract_toc_nodes(ctx.blackboard.toc_hierarchies) - phase1 = calibrate_offset( - nodes=nodes, - toc_hierarchies=ctx.blackboard.toc_hierarchies, - ctx=ctx, - page_texts={}, - page_count=page_count, - ) - offset_hint = pick_primary_offset(phase1) - ctx.blackboard.toc_page_offset = offset_hint + offset_hint = ctx.blackboard.toc_page_offset # Try TOC chapter-based planning first chapters = derive_chapter_boundaries( diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 0c09005ba..da5f8602b 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -13,8 +13,6 @@ purge_debug_visual_dirs, visual_debug_enabled, ) -from app.services.document_agent.manifest import ToolContext -from app.services.document_agent.state import AgentBlackboard from app.services.document_parser.profiling.doc_profiler import profile_document from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time from app.services.document_parser.support.parser_rows import PARSER_ROW_COLUMNS @@ -238,15 +236,6 @@ def _build_page_dataframe( asset_extraction_enabled = page_memory_config.asset_extraction_enabled - # ── build ToolContext for sub-agent VLM calls ───────────────────── - ctx = _build_page_ctx( - pdf_path=pdf_path, - job_id=filename, - output_dir=output_dir, - page_count=page_count, - trace_recorder=trace_recorder, - ) - # ── C4: skeleton (from profile anatomy) ─────────────────────────── with stage_timer("page_memory.read_page_texts", page_count=page_count): page_texts = read_page_texts( @@ -258,7 +247,6 @@ def _build_page_dataframe( anatomy=anatomy, filename=filename, page_texts=page_texts, - ctx=ctx, ) else: skeletons = [] @@ -469,36 +457,6 @@ def _build_page_dataframe( return pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) -def _build_page_ctx( - *, - pdf_path: str, - job_id: str, - output_dir: str, - page_count: int, - trace_recorder: Any | None = None, -) -> ToolContext: - """Construct a ToolContext for C4 sub-agent and C3 tagger VLM calls.""" - blackboard = AgentBlackboard() - blackboard.page_count = page_count - vlm_model = os.environ.get("IMAGE_MODEL") - reason_model = ( - os.environ.get("PAGE_LOCATE_REASON_MODEL") - or os.environ.get("NORMOL_MODEL") - ) - return ToolContext( - pdf_path=pdf_path, - job_id=job_id, - blackboard=blackboard, - budget=None, - trace=trace_recorder, - output_dir=output_dir, - settings={ - "vlm_model": vlm_model, - "model": reason_model, - }, - ) - - def _build_hierarchy_scopes( *, skeletons: list[Any], diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 3dfd19c80..976c9ecfe 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -1,8 +1,7 @@ """Build page-memory section skeletons from profile-time anatomy. -Step 1 of the page-memory native hierarchy plan: -- Full TOC-depth grep anchoring + on-demand VLM confirmation -- Section boundaries come purely from TOC anchoring +Section boundaries come from TOC anchoring persisted on anatomy +(``skeleton_anchor``). This module only resolves page ranges. """ from __future__ import annotations @@ -12,29 +11,28 @@ from app.services.document_agent.manifest import ( PageAnatomyMap, - ToolContext, ) from app.services.document_agent.structure.hierarchy_locator import ( ResolvedHierarchyRange, - TitleNode, extract_toc_nodes, - iter_leaf_title_nodes, resolve_hierarchy_page_ranges, ) -from app.services.document_agent.agents.calibration import calibrate_offset -from app.services.document_agent.agents.calibration.orchestrator import anchor_hierarchy from app.services.document_agent.structure.anchoring_primitives import ( - toc_range_end, + deserialize_skeleton_anchor, + deserialize_title_node, toc_range_start, ) +from app.services.document_agent.structure.toc_anchoring import ( + body_pages_excluding_toc, + pending_toc_body_scope, + select_global_toc_hierarchies, +) from loguru import logger from shared.services.chunks.path_segments import ( append_document_path, join_document_path, ) -_FRONT_TOC_REGION_GAP_PAGES = 5 - @@ -57,74 +55,46 @@ def extract_section_skeletons( anatomy: PageAnatomyMap | Any | None, filename: str, page_texts: dict[int, str], - ctx: ToolContext | None = None, - hierarchy_nodes: list[TitleNode] | None = None, ) -> list[SectionSkeleton]: """Convert PageAnatomyMap hierarchy evidence into section skeletons. Section page ranges are anchored purely from the TOC hierarchy (every - level, every document). + level, every document). Calibration runs at PROFILE; this stage only + resolves persisted ``skeleton_anchor``. """ page_count = _page_count(anatomy) root_path = f"{filename}/Root" if page_count <= 0: return [_root_skeleton(root_path=root_path, filename=filename, page_count=0)] - toc_selection: dict[str, Any] = {} - pending_tocs: list[dict[str, Any]] = [] - toc_hierarchies: list[dict[str, Any]] | None = None - if hierarchy_nodes: - nodes = hierarchy_nodes - else: - toc_hierarchies, pending_tocs, toc_selection = _select_global_toc_hierarchies( - anatomy=anatomy, - filename=filename, - ) - toc_nodes = extract_toc_nodes(toc_hierarchies) - if not toc_nodes: - # TODO: explore lightweight hierarchy inference for no-TOC documents - # (e.g. heading font-size clustering, visual layout analysis). - # For now, no TOC → flat page tagging + asset extraction only. - return [ - _root_skeleton( - root_path=root_path, - filename=filename, - page_count=page_count, - reason="no_toc", - ) - ] - nodes = toc_nodes - - # Collapse degenerate single-child intermediate chains before locate. - # Rule: only merge a parent with its only child when that child is NOT a - # leaf (i.e. the child still has children of its own). This preserves the - # original leaf title so offset-guided anchoring can find it in the PDF. - nodes = _collapse_intermediate_single_child_chains(nodes) - - body_pages = _body_pages(anatomy=anatomy, page_count=page_count) - - # When pending TOCs exist, limit primary scope so the last sibling's - # end_page doesn't extend into the pending TOC region. - primary_page_count = page_count - primary_body_pages = body_pages - if pending_tocs: - pending_starts: list[int] = [] - for t in pending_tocs: - start = toc_range_start(t) - if start is not None: - pending_starts.append(start) - if pending_starts: - primary_page_count = min(pending_starts) - 1 - primary_body_pages = [p for p in body_pages if p <= primary_page_count] - - resolve_nodes, skeleton_anchor = anchor_hierarchy( - nodes=nodes, - toc_hierarchies=toc_hierarchies if not hierarchy_nodes else None, - page_texts=page_texts, - body_pages=primary_body_pages, - page_count=page_count, - ctx=ctx, + toc_hierarchies, pending_tocs, toc_selection = select_global_toc_hierarchies( + hierarchies=list(getattr(anatomy, "toc_hierarchies", None) or []), + filename=filename, ) + toc_nodes = extract_toc_nodes(toc_hierarchies) + if not toc_nodes: + return [ + _root_skeleton( + root_path=root_path, + filename=filename, + page_count=page_count, + reason="no_toc", + ) + ] + + skeleton_anchor_raw = getattr(anatomy, "skeleton_anchor", None) + skeleton_nodes_raw = getattr(anatomy, "skeleton_nodes", None) + if not isinstance(skeleton_anchor_raw, dict) or not isinstance( + skeleton_nodes_raw, list + ): + raise ValueError("anatomy.skeleton_anchor/skeleton_nodes missing after TOC extract") + + skeleton_anchor = deserialize_skeleton_anchor(skeleton_anchor_raw) + resolve_nodes = [ + deserialize_title_node(node) + for node in skeleton_nodes_raw + if isinstance(node, dict) + ] if skeleton_anchor.pruned_count and not resolve_nodes: return [ _root_skeleton( @@ -135,6 +105,23 @@ def extract_section_skeletons( ) ] + toc_result = getattr(anatomy, "toc_result", None) + body_pages = body_pages_excluding_toc( + getattr(toc_result, "toc_pages", None), + page_count, + ) + primary_page_count = page_count + primary_body_pages = body_pages + if pending_tocs: + pending_starts: list[int] = [] + for pending_toc in pending_tocs: + start = toc_range_start(pending_toc) + if start is not None: + pending_starts.append(start) + if pending_starts: + primary_page_count = min(pending_starts) - 1 + primary_body_pages = [page for page in body_pages if page <= primary_page_count] + match_overrides = skeleton_anchor.match_overrides null_page_report = skeleton_anchor.null_page_report if skeleton_anchor.locate_agent == "offset_guided_bulk": @@ -191,12 +178,11 @@ def extract_section_skeletons( for item in ranges ] - # Phase B: graft pending TOCs (appendix / parallel sections) - if pending_tocs: + pending_records = list(getattr(anatomy, "pending_skeleton_anchors", None) or []) + if pending_tocs and pending_records: secondary_skeletons = _resolve_pending_tocs( pending_tocs=pending_tocs, - primary_ranges=ranges, - ctx=ctx, + pending_records=pending_records, page_texts=page_texts, page_count=page_count, filename=filename, @@ -245,31 +231,6 @@ def _range_to_skeleton( ) -# ── Single-child intermediate chain collapse ───────────────────────────────── -# -# Motivation: TOC hierarchies often contain "structural" intermediate nodes -# (category codes, volume identifiers) that add depth but carry no locatable -# text. Compressing them before locate keeps emit_depth small and lets the -# offset-guided anchoring focus on meaningful leaf titles. -# -# Critical invariant: a node whose only child is a LEAF (no grandchildren) is -# NOT merged, so the leaf's original title survives unchanged into -# offset-guided anchoring. Only pure-intermediate chains are compressed. - - -def _collapse_intermediate_single_child_chains( - nodes: list[TitleNode], -) -> list[TitleNode]: - """Collapse single-child chains of intermediate (non-leaf) nodes. - - Leaf nodes (children=[]) are never absorbed into their parent title. - """ - from app.services.document_agent.structure.hierarchy_locator import ( - collapse_intermediate_single_child_chains, - ) - - return collapse_intermediate_single_child_chains(nodes) - def _root_skeleton( *, root_path: str, @@ -293,179 +254,66 @@ def _page_count(anatomy: Any | None) -> int: return max(int(getattr(anatomy, "page_count", 0) or 0), 0) -def _toc_hierarchies(anatomy: Any | None) -> list[dict[str, Any]] | None: - return getattr(anatomy, "toc_hierarchies", None) if anatomy is not None else None - - -def _select_global_toc_hierarchies( - *, - anatomy: Any | None, - filename: str, -) -> tuple[list[dict[str, Any]] | None, list[dict[str, Any]], dict[str, Any]]: - """Split TOC hierarchies into primary (front cluster) and pending (for probe). - - Profile-time TOC extraction can find multiple TOCs in a long document. - The front cluster is selected by physical page proximity. Remaining TOCs - are returned as *pending* for downstream independent calibration rather - than being unconditionally discarded. - - Returns (primary_hierarchies, pending_hierarchies, summary). - """ - hierarchies = list(_toc_hierarchies(anatomy) or []) - if len(hierarchies) <= 1: - return (hierarchies or None), [], {} - - page_based = [ - hierarchy - for hierarchy in hierarchies - if hierarchy.get("toc_range_unit") == "page" and toc_range_start(hierarchy) is not None - ] - if not page_based or len(page_based) != len(hierarchies): - return hierarchies, [], {} - - sorted_items = sorted(enumerate(hierarchies), key=lambda item: toc_range_start(item[1]) or 0) - selected_indices: set[int] = set() - pending_indices: list[int] = [] - cluster_end: int | None = None - - for original_index, hierarchy in sorted_items: - start = toc_range_start(hierarchy) - end = toc_range_end(hierarchy) - if start is None or end is None: - selected_indices.add(original_index) - continue - if cluster_end is None: - selected_indices.add(original_index) - cluster_end = end - continue - if start <= cluster_end + _FRONT_TOC_REGION_GAP_PAGES: - selected_indices.add(original_index) - cluster_end = max(cluster_end, end) - continue - pending_indices.append(original_index) - - selected = [ - hierarchy - for index, hierarchy in enumerate(hierarchies) - if index in selected_indices - ] - pending = [hierarchies[i] for i in pending_indices] - - if pending: - logger.info( - "[page_memory.skeleton] toc split: primary={} pending={} filename={}", - len(selected), - len(pending), - filename, - ) - summary = { - "strategy": "front_cluster_with_pending", - "input_count": len(hierarchies), - "primary_count": len(selected), - "pending_count": len(pending), - } - return (selected or None), pending, summary - - - - -def _body_pages(*, anatomy: Any | None, page_count: int) -> list[int]: - excluded: set[int] = set() - toc_result = getattr(anatomy, "toc_result", None) - excluded.update(int(page) for page in getattr(toc_result, "toc_pages", []) or []) - return [page for page in range(1, page_count + 1) if page not in excluded] - - - - -# ── Multi-TOC grafting (Track B) ───────────────────────────────────────────── - - def _resolve_pending_tocs( *, pending_tocs: list[dict[str, Any]], - primary_ranges: list[ResolvedHierarchyRange], - ctx: ToolContext | None, + pending_records: list[dict[str, Any]], page_texts: dict[int, str], page_count: int, filename: str, body_pages: list[int], ) -> list[SectionSkeleton]: - """Independently calibrate and anchor each pending TOC, then graft results. - - Each pending TOC gets its own offset via VLM calibration + tail verify, - then entries are bulk-anchored (or fallback to residual agent). - Classification is PARALLEL (append at root level) or CONTAINED (skip). - """ - if not pending_tocs or ctx is None: + """Graft pending TOCs from PROFILE-persisted skeleton anchors.""" + if not pending_tocs or not pending_records: return [] + records_by_range: dict[tuple[Any, ...], dict[str, Any]] = {} + for record in pending_records: + toc = record.get("toc") + if not isinstance(toc, dict): + continue + key = tuple(toc.get("toc_range") or []) + records_by_range[key] = record + all_secondary_skeletons: list[SectionSkeleton] = [] for i, pending_toc in enumerate(pending_tocs): toc_range = pending_toc.get("toc_range") - nodes = extract_toc_nodes([pending_toc]) - if not nodes: + record = records_by_range.get(tuple(toc_range or [])) + if record is None: continue - nodes = _collapse_intermediate_single_child_chains(nodes) - - # Each TOC's content scope: [toc_range_end + 1, next_toc_start - 1] - toc_end = toc_range_end(pending_toc) - toc_scope_start = (toc_end + 1) if toc_end is not None else None - next_starts: list[int] = [] - for j in range(i + 1, len(pending_tocs)): - start = toc_range_start(pending_tocs[j]) - if start is not None: - next_starts.append(start) - toc_scope_end = (min(next_starts) - 1) if next_starts else page_count - toc_body_pages = [ - p for p in body_pages - if p <= toc_scope_end and (toc_scope_start is None or p >= toc_scope_start) - ] - - from app.services.document_agent.agents.calibration.procedure import ( - finalize_calibration_result, - pick_primary_offset, - ) - - phase1 = calibrate_offset( - nodes=nodes, - toc_hierarchies=[pending_toc], - ctx=ctx, - page_texts=page_texts, - page_count=toc_scope_end, - ) - offset = pick_primary_offset(phase1) - - if offset is None: - logger.info( - "[page_memory.skeleton] pending TOC toc_range={}: calibration failed, skipping", - toc_range, - ) - continue - - relationship = _classify_toc_relationship( - offset=offset, - nodes=nodes, - primary_ranges=primary_ranges, - page_count=page_count, - ) + relationship = record.get("relationship") if relationship == "unresolvable": logger.info( "[page_memory.skeleton] pending TOC toc_range={}: unresolvable, skipping", toc_range, ) continue + if relationship not in {"parallel", "contained"}: + raise ValueError( + "pending TOC relationship missing after PROFILE classify" + ) + resolve_nodes_raw = record.get("nodes") or [] + resolve_nodes = [ + deserialize_title_node(node) + for node in resolve_nodes_raw + if isinstance(node, dict) + ] + skeleton_anchor_raw = record.get("skeleton_anchor") + if not isinstance(skeleton_anchor_raw, dict) or not resolve_nodes: + raise ValueError( + "pending TOC skeleton_anchor/nodes missing after PROFILE classify" + ) + skeleton_anchor = deserialize_skeleton_anchor(skeleton_anchor_raw) + offset = skeleton_anchor.offset + if offset is None: + raise ValueError("pending TOC offset missing after PROFILE classify") - resolve_nodes, skeleton_anchor, _finalized = finalize_calibration_result( - result=phase1, - entries=list(pending_toc.get("toc_with_level") or []), - toc_hierarchies=[pending_toc], - ctx=ctx, - page_count=toc_scope_end, - page_texts=page_texts, - body_pages=toc_body_pages, - nodes=nodes, + toc_scope_end, toc_body_pages = pending_toc_body_scope( + pending_tocs=pending_tocs, + index=i, + page_count=page_count, + body_pages=body_pages, ) match_overrides = skeleton_anchor.match_overrides null_page_report = skeleton_anchor.null_page_report @@ -523,59 +371,6 @@ def _resolve_pending_tocs( return all_secondary_skeletons -def _classify_toc_relationship( - *, - offset: int, - nodes: list[TitleNode], - primary_ranges: list[ResolvedHierarchyRange], - page_count: int, -) -> str: - """Classify a pending TOC as parallel or contained vs primary ranges. - - parallel: the pending TOC covers pages beyond the primary tree's *anchored* - content (i.e. the last explicitly-located section start page). - contained: the pending TOC's content falls strictly within a primary - section's explicitly-anchored range. - """ - leaves = [ - node for _, node in iter_leaf_title_nodes(nodes) if node.printed_page is not None - ] - if not leaves: - return "unresolvable" - - first_printed = leaves[0].printed_page - last_printed = leaves[-1].printed_page - if first_printed is None or last_printed is None: - return "unresolvable" - first_physical = first_printed + offset - last_physical = last_printed + offset - - if first_physical < 1 or first_physical > page_count: - return "unresolvable" - - if not primary_ranges: - return "parallel" - - # Use the last *start_page* among primary ranges as the boundary of - # explicitly-anchored content. The end_page of the last section is often - # extended to page_count by default and doesn't reflect real content coverage. - last_anchored_start = max( - (r.start_page for r in primary_ranges if r.start_page is not None), default=0 - ) - - if first_physical > last_anchored_start: - return "parallel" - - min_level = min(r.level for r in primary_ranges) - top_level_ranges = [r for r in primary_ranges if r.level == min_level] - for r in top_level_ranges: - if r.start_page and r.end_page: - if r.start_page <= first_physical and last_physical <= r.end_page: - return "contained" - - return "parallel" - - def _clamp_page(page: int, page_count: int) -> int: return min(max(page, 1), max(page_count, 1)) diff --git a/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py b/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py new file mode 100644 index 000000000..388ca1d55 --- /dev/null +++ b/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py @@ -0,0 +1,429 @@ +"""PROFILE writes skeleton_anchor; C4 and shard plan do not recalibrate.""" + +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.budget import BudgetTracker +from app.services.document_agent.manifest import ( + PageAnatomyMap, + PageFeature, + PageLabel, + Shard, + ShardPlan, + TocResult, + ToolContext, +) +from app.services.document_agent.state import AgentBlackboard +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 ( + ResolvedHierarchyRange, + TitleMatch, + TitleNode, +) +from app.services.document_agent.structure.toc_anchoring import ( + classify_toc_relationship, + run_toc_anchoring, +) +from app.services.document_agent.tools.propose_shard_plan import propose_shard_plan +from app.services.page_memory.skeleton_extractor import extract_section_skeletons + + +def _ctx(*, page_count: int = 10) -> ToolContext: + return ToolContext( + pdf_path="/tmp/doc.pdf", + job_id="job-wire", + blackboard=AgentBlackboard(page_count=page_count), + budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), + trace=None, + settings={}, + ) + + +def _toc() -> list[dict[str, object]]: + return [ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 2}, + ], + } + ] + + +def _node() -> TitleNode: + return TitleNode(title="Ch1", level=1, printed_page=2, children=[]) + + +def _anchor(*, title: str = "Ch1", page: int = 2) -> SkeletonAnchor: + return SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides={ + (title,): TitleMatch( + page=page, + confidence=1.0, + source="anchored", + matched_line=title, + score=1.0, + candidates=[page], + evidence={}, + ) + }, + null_page_report=[], + bulk_count=1, + pruned_count=0, + locate_agent="offset_guided_bulk", + ) + + +def _anatomy(*, with_anchor: bool) -> PageAnatomyMap: + kwargs: dict[str, object] = { + "job_id": "job-wire", + "file_path": "/tmp/doc.pdf", + "page_count": 10, + "page_features": [ + PageFeature( + page=page, + raw_text_length=20, + text_density=0.1, + image_coverage=0.0, + image_count=0, + table_count=0, + drawings_count=0, + orientation="portrait", + width=72.0, + height=72.0, + has_asset=False, + is_blank_like=False, + ) + for page in range(1, 11) + ], + "page_labels": [ + PageLabel(page=page, kind="normal", confidence=1.0) + for page in range(1, 11) + ], + "toc_result": TocResult(method="vlm_batch", toc_pages=[1]), + "shard_plan": ShardPlan( + enabled=False, + reason="not_needed", + shards=[ + Shard( + shard_index=0, + page_start=1, + page_end=10, + page_offset=0, + anchor_type="forced_max_size", + anchor_evidence="test", + confidence=1.0, + ) + ], + ), + "toc_hierarchies": _toc(), + "toc_page_offset": 0 if with_anchor else None, + } + if with_anchor: + kwargs["skeleton_anchor"] = serialize_skeleton_anchor(_anchor()) + kwargs["skeleton_nodes"] = [serialize_title_node(_node())] + return PageAnatomyMap(**kwargs) # type: ignore[arg-type] + + +def test_profile_toc_anchoring_writes_skeleton_anchor() -> None: + ctx = _ctx() + ctx.blackboard.toc_hierarchies = _toc() + ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1]) + + def fake_anchor_hierarchy(**_kwargs): + return [_node()], _anchor() + + with ( + patch( + "app.services.document_agent.structure.toc_anchoring.read_page_texts", + return_value={page: "Ch1" for page in range(1, 11)}, + ), + patch( + "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + side_effect=fake_anchor_hierarchy, + ), + ): + run_toc_anchoring(ctx) + + assert ctx.blackboard.toc_page_offset == 0 + assert isinstance(ctx.blackboard.skeleton_anchor, dict) + assert ctx.blackboard.skeleton_anchor["offset"] == 0 + assert ctx.blackboard.skeleton_nodes + assert ctx.blackboard.skeleton_nodes[0]["title"] == "Ch1" + + +def test_c4_resolve_does_not_call_calibration() -> None: + anatomy = _anatomy(with_anchor=True) + page_texts = {page: "Ch1 body" for page in range(1, 11)} + + def _boom(*_args, **_kwargs): + raise AssertionError("C4 must not recalibrate") + + with ( + patch( + "app.services.document_agent.agents.calibration.service.calibrate_offset", + side_effect=_boom, + ), + patch( + "app.services.document_agent.agents.calibration.orchestrator.anchor_hierarchy", + side_effect=_boom, + ), + patch( + "app.services.document_agent.agents.calibration.procedure.finalize_calibration_result", + side_effect=_boom, + ), + ): + skeletons = extract_section_skeletons( + anatomy=anatomy, + filename="doc.pdf", + page_texts=page_texts, + ) + + assert skeletons + assert skeletons[0].title == "Ch1" + assert skeletons[0].start_page == 2 + + +def test_shard_plan_reads_offset_and_does_not_calibrate() -> None: + ctx = _ctx(page_count=250) + ctx.blackboard.toc_hierarchies = [ + { + "toc_range": [1, 2], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 3}, + {"heading": "Ch2", "level": 1, "page_number": 120}, + ], + } + ] + ctx.blackboard.toc_page_offset = 0 + ctx.blackboard.toc_result = TocResult(method="vlm_batch") + ctx.blackboard.doc_stats = {"page_count": 250} + ctx.settings["shard_threshold"] = 200 + ctx.settings["max_pages_per_shard"] = 200 + ctx.settings["min_pages_per_shard"] = 20 + + def _boom(*_args, **_kwargs): + raise AssertionError("shard plan must not recalibrate") + + with patch( + "app.services.document_agent.agents.calibration.service.calibrate_offset", + side_effect=_boom, + ): + result = propose_shard_plan(ctx, {}) + + assert result.status == "ok" + assert ctx.blackboard.toc_page_offset == 0 + assert ctx.blackboard.shard_plan is not None + + +def _pending_tocs() -> list[dict[str, object]]: + return [ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 2}, + ], + }, + { + "toc_range": [20, 21], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "App", "level": 1, "page_number": 22}, + ], + }, + ] + + +def test_profile_classifies_pending_toc_before_finalize() -> None: + ctx = _ctx(page_count=30) + ctx.blackboard.toc_hierarchies = _pending_tocs() + ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1, 20, 21]) + pending_node = TitleNode(title="App", level=1, printed_page=22, children=[]) + finalize_calls: list[object] = [] + + def fake_finalize(**kwargs): + finalize_calls.append(kwargs["nodes"]) + return [pending_node], _anchor(title="App", page=22), True + + with ( + patch( + "app.services.document_agent.structure.toc_anchoring.read_page_texts", + return_value={page: "body" for page in range(1, 31)}, + ), + patch( + "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + return_value=([_node()], _anchor()), + ), + patch( + "app.services.document_agent.structure.toc_anchoring.resolve_hierarchy_page_ranges", + return_value=[ + ResolvedHierarchyRange( + title="Ch1", + level=1, + start_page=2, + end_page=19, + path_titles=("Ch1",), + match=None, + ) + ], + ), + patch( + "app.services.document_agent.structure.toc_anchoring.calibrate_offset", + return_value=object(), + ), + patch( + "app.services.document_agent.structure.toc_anchoring.pick_primary_offset", + return_value=0, + ), + patch( + "app.services.document_agent.structure.toc_anchoring.finalize_calibration_result", + side_effect=fake_finalize, + ), + ): + run_toc_anchoring(ctx) + + records = ctx.blackboard.pending_skeleton_anchors + assert len(records) == 1 + assert records[0]["relationship"] == "parallel" + assert finalize_calls + assert records[0]["nodes"][0]["title"] == "App" + + +def test_profile_skips_finalize_for_unresolvable_pending_toc() -> None: + ctx = _ctx(page_count=30) + hierarchies = _pending_tocs() + hierarchies[1]["toc_with_level"] = [{"heading": "App", "level": 1}] + ctx.blackboard.toc_hierarchies = hierarchies + ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1, 20, 21]) + + def _boom(*_args, **_kwargs): + raise AssertionError("unresolvable pending TOC must not finalize") + + with ( + patch( + "app.services.document_agent.structure.toc_anchoring.read_page_texts", + return_value={page: "body" for page in range(1, 31)}, + ), + patch( + "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + return_value=([_node()], _anchor()), + ), + patch( + "app.services.document_agent.structure.toc_anchoring.resolve_hierarchy_page_ranges", + return_value=[ + ResolvedHierarchyRange( + title="Ch1", + level=1, + start_page=2, + end_page=19, + path_titles=("Ch1",), + match=None, + ) + ], + ), + patch( + "app.services.document_agent.structure.toc_anchoring.calibrate_offset", + return_value=object(), + ), + patch( + "app.services.document_agent.structure.toc_anchoring.pick_primary_offset", + return_value=0, + ), + patch( + "app.services.document_agent.structure.toc_anchoring.finalize_calibration_result", + side_effect=_boom, + ), + ): + run_toc_anchoring(ctx) + + records = ctx.blackboard.pending_skeleton_anchors + assert len(records) == 1 + assert records[0]["relationship"] == "unresolvable" + assert "nodes" not in records[0] + assert "skeleton_anchor" not in records[0] + + +def test_c4_uses_persisted_pending_relationship_and_does_not_classify() -> None: + pending_toc = _pending_tocs()[1] + anatomy = _anatomy(with_anchor=True) + anatomy.page_count = 30 + anatomy.page_features = [ + PageFeature( + page=page, + raw_text_length=20, + text_density=0.1, + image_coverage=0.0, + image_count=0, + table_count=0, + drawings_count=0, + orientation="portrait", + width=72.0, + height=72.0, + has_asset=False, + is_blank_like=False, + ) + for page in range(1, 31) + ] + anatomy.page_labels = [ + PageLabel(page=page, kind="normal", confidence=1.0) + for page in range(1, 31) + ] + anatomy.toc_hierarchies = _pending_tocs() + anatomy.pending_skeleton_anchors = [ + { + "toc": pending_toc, + "relationship": "parallel", + "nodes": [ + serialize_title_node( + TitleNode(title="App", level=1, printed_page=22, children=[]) + ) + ], + "skeleton_anchor": serialize_skeleton_anchor(_anchor(title="App", page=22)), + } + ] + page_texts = {page: "Ch1 body" for page in range(1, 31)} + page_texts[22] = "App" + + def _boom(*_args, **_kwargs): + raise AssertionError("C4 must not classify pending TOC") + + with patch( + "app.services.document_agent.structure.toc_anchoring.classify_toc_relationship", + side_effect=_boom, + ): + skeletons = extract_section_skeletons( + anatomy=anatomy, + filename="doc.pdf", + page_texts=page_texts, + ) + + titles = [skeleton.title for skeleton in skeletons] + assert "Ch1" in titles + assert "App" in titles + app = next(skeleton for skeleton in skeletons if skeleton.title == "App") + assert app.evidence["page_locate_summary"]["toc_relationship"] == "parallel" + + +def test_classify_toc_relationship_is_not_on_c4_module() -> None: + import app.services.page_memory.skeleton_extractor as skeleton_extractor + + assert not hasattr(skeleton_extractor, "_classify_toc_relationship") + assert callable(classify_toc_relationship) From 712e7e765fb5f3db781a190df8deb0422ca6f18f Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 13 Aug 2026 16:30:51 +0800 Subject: [PATCH 2/9] feat: integrate OCR text scanning and enhance document processing - Added `rapidocr-onnxruntime` as a dependency for OCR capabilities. - Implemented `_run_text_scan` method in `ProfileCoordinator` to handle text extraction from PDF pages. - Updated `PageAnatomyMap` to include `page_full_text_cache` for storing scanned text. - Enhanced TOC anchoring logic to utilize cached page text for improved accuracy. - Refactored various tools to leverage the new text scanning functionality, ensuring seamless integration with existing workflows. --- .../services/document_agent/coordinator.py | 53 +- .../app/services/document_agent/manifest.py | 12 +- .../app/services/document_agent/pdf_text.py | 6 + .../app/services/document_agent/registry.py | 7 + .../document_agent/structure/toc_anchoring.py | 9 +- .../services/document_agent/tools/__init__.py | 1 + .../tools/find_toc_anchor_pages.py | 60 +-- .../document_agent/tools/grep_text.py | 21 +- .../document_agent/tools/ocr_pages.py | 102 ++++ .../tools/persist_anatomy_map.py | 1 + .../tools/propose_shard_plan.py | 459 +++--------------- .../app/services/document_agent/visual.py | 1 + .../document_parser/profiling/doc_profiler.py | 4 +- .../profiling/profile_model.py | 2 + .../services/page_memory/memory_service.py | 17 +- .../app/services/page_memory/page_renderer.py | 20 +- apps/worker/pyproject.toml | 1 + .../test_doc_profile_anatomy_contract.py | 94 +++- .../tests/contract/test_ocr_pages_contract.py | 78 +++ ...profile_skeleton_anchor_wiring_contract.py | 53 +- .../test_propose_shard_plan_contract.py | 136 ++++++ uv.lock | 136 ++++++ 22 files changed, 765 insertions(+), 508 deletions(-) create mode 100644 apps/worker/app/services/document_agent/tools/ocr_pages.py create mode 100644 apps/worker/tests/contract/test_ocr_pages_contract.py create mode 100644 apps/worker/tests/contract/test_propose_shard_plan_contract.py diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index ce8715f70..b1be64991 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -22,6 +22,7 @@ ToolContext, ToolResult, ) +from app.services.document_agent.pdf_text import read_page_texts from app.services.document_agent.persist import build_anatomy_map, persist_anatomy_map from app.services.document_agent.planner import ProfilePlanner from app.services.document_agent.registry import REGISTRY @@ -142,11 +143,14 @@ def _run_coarse(self) -> DocumentProfile: self.state = DocumentAgentState.RUNNING if not self.blackboard.page_features: self._run_bootstrap() - if self._should_run_toc_before_coarse(): - self._ensure_toc_profile(strict=False) profile, _initial_decision, _planner_result = self._propose_profile( actor="planner:coarse" ) + self._run_text_scan() + if self._toc_profile_enabled(): + self._ensure_toc_profile(strict=False) + else: + self._ensure_disabled_toc_placeholder() self._ensure_asset_probe() return profile @@ -321,14 +325,48 @@ def _toc_result_requires_strict_retry(self) -> bool: and toc_result.failure_kind in {"confirm_failed", "degraded"} ) - def _should_run_toc_before_coarse(self) -> bool: - return self._toc_profile_enabled() and bool( - self.ctx.settings.get("toc_before_coarse") - ) - def _toc_profile_enabled(self) -> bool: return bool(self.ctx.settings.get("toc_profile_enabled", True)) + def _run_text_scan(self) -> None: + profile = self.blackboard.document_profile + if profile is None: + raise RuntimeError("document_profile missing; run planner first") + page_count = int(self.blackboard.page_count or 0) + pages = list(range(1, page_count + 1)) + if not pages: + self.blackboard.page_full_text_cache = {} + return + if profile.is_scanned: + result = REGISTRY.dispatch("ocr.pages", self.ctx, {"pages": pages}) + self.trace.record_step( + round_index=self.round_index, + actor="scan:ocr.pages", + action_type="scan", + result=result, + tool_name="ocr.pages", + tool_args={"pages": pages}, + ) + if result.status != "ok": + raise RuntimeError(result.error or "ocr.pages failed") + self.round_index += 1 + return + texts = read_page_texts(self.ctx.pdf_path, pages, timeout=300) + self.blackboard.page_full_text_cache = texts + self.trace.record_step( + round_index=self.round_index, + actor="scan:read_page_texts", + action_type="scan", + result=ToolResult( + status="ok", + payload={"page_count": len(texts)}, + output_summary={"page_count": len(texts)}, + ), + tool_name="read_page_texts", + tool_args={"pages": pages}, + ) + self.round_index += 1 + def _ensure_disabled_toc_placeholder(self) -> None: self.blackboard.toc_result = TocResult( method="none", @@ -349,7 +387,6 @@ def _ensure_toc_profile(self, *, strict: bool) -> None: if not should_run: return - self._planner_cache = None self.blackboard.global_signals["toc_profile_attempted"] = True try: self._run_toc_extraction_pipeline() diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index e545c63c0..fda87f8b0 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -186,7 +186,13 @@ class Shard: page_start: int page_end: int page_offset: int - anchor_type: Literal["h1_boundary", "blank_separator", "forced_max_size", "toc_chapter_boundary"] + anchor_type: Literal[ + "h1_boundary", + "blank_separator", + "forced_max_size", + "toc_chapter_boundary", + "toc_leaf_boundary", + ] anchor_evidence: str confidence: float @@ -234,6 +240,7 @@ class PageAnatomyMap: skeleton_anchor: dict[str, Any] | None = None skeleton_nodes: list[dict[str, Any]] | None = None pending_skeleton_anchors: list[dict[str, Any]] = field(default_factory=list) + page_full_text_cache: dict[int, str] = field(default_factory=dict) global_signals: dict[str, Any] = field(default_factory=dict) trace_summary: dict[str, Any] = field(default_factory=dict) created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) @@ -258,6 +265,9 @@ def to_dict(self) -> dict[str, Any]: "skeleton_anchor": self.skeleton_anchor, "skeleton_nodes": self.skeleton_nodes, "pending_skeleton_anchors": list(self.pending_skeleton_anchors), + "page_full_text_cache": { + str(page): text for page, text in self.page_full_text_cache.items() + }, "global_signals": dict(self.global_signals), "trace_summary": dict(self.trace_summary), "created_at": self.created_at.isoformat(), diff --git a/apps/worker/app/services/document_agent/pdf_text.py b/apps/worker/app/services/document_agent/pdf_text.py index ca9684370..7a6da6522 100644 --- a/apps/worker/app/services/document_agent/pdf_text.py +++ b/apps/worker/app/services/document_agent/pdf_text.py @@ -35,6 +35,12 @@ def _read_page_texts_worker(queue, pdf_path: str, pages: list[int]) -> None: queue.put({"ok": True, "texts": texts}) +def coerce_page_text_cache(raw: Any) -> dict[int, str]: + if not isinstance(raw, dict): + return {} + return {int(page): str(text) for page, text in raw.items()} + + def read_page_texts( pdf_path: str, pages: list[int], diff --git a/apps/worker/app/services/document_agent/registry.py b/apps/worker/app/services/document_agent/registry.py index 9fd4966ca..330290d20 100644 --- a/apps/worker/app/services/document_agent/registry.py +++ b/apps/worker/app/services/document_agent/registry.py @@ -135,6 +135,13 @@ def has_document_profile(blackboard: AgentBlackboard) -> tuple[bool, str]: return blackboard.document_profile is not None, "document_profile missing; run planner first" +def has_page_full_text(blackboard: AgentBlackboard) -> tuple[bool, str]: + return ( + bool(blackboard.page_full_text_cache), + "page_full_text_cache missing; run text scan first", + ) + + def not_is_scanned(blackboard: AgentBlackboard) -> tuple[bool, str]: profile = blackboard.document_profile return ( 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 1e774a91c..2f7fe1b70 100644 --- a/apps/worker/app/services/document_agent/structure/toc_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/toc_anchoring.py @@ -14,7 +14,6 @@ pick_primary_offset, ) from app.services.document_agent.manifest import ToolContext -from app.services.document_agent.pdf_text import read_page_texts from app.services.document_agent.structure.anchoring_primitives import ( serialize_skeleton_anchor, serialize_title_node, @@ -41,13 +40,9 @@ def run_toc_anchoring(ctx: ToolContext) -> None: if page_count <= 0 or not hierarchies: return - if not ctx.blackboard.page_full_text_cache: - ctx.blackboard.page_full_text_cache = read_page_texts( - ctx.pdf_path, - list(range(1, page_count + 1)), - timeout=300, - ) page_texts = dict(ctx.blackboard.page_full_text_cache) + if not page_texts: + raise ValueError("page_full_text_cache missing; run text scan before TOC anchoring") filename = Path(ctx.pdf_path).name primary, pending, _summary = select_global_toc_hierarchies( hierarchies=hierarchies, diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index 43b91226a..6ccb1a8a7 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -5,6 +5,7 @@ from . import extract_toc_with_boundaries as extract_toc_with_boundaries # noqa: F401 from . import find_toc_anchor_pages as find_toc_anchor_pages # noqa: F401 from . import grep_text as grep_text # noqa: F401 +from . import ocr_pages as ocr_pages # noqa: F401 from . import propose_shard_plan as propose_shard_plan # noqa: F401 from . import validate_anatomy_map as validate_anatomy_map # noqa: F401 from . import verdict as verdict # noqa: F401 diff --git a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py index e2acb4e51..14a8e213f 100644 --- a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py +++ b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py @@ -9,7 +9,7 @@ from typing import Any from app.services.document_agent.manifest import TocAnchorPage, ToolContext, ToolResult -from app.services.document_agent.registry import has_page_labels, register_tool +from app.services.document_agent.registry import has_page_full_text, has_page_labels, register_tool from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( run_in_child_process, worker, @@ -98,35 +98,22 @@ def _find_toc_text_matches( return matches -@worker -def _scan_toc_text_worker( - queue, pdf_path: str, cross_line_window: int -) -> None: - import pymupdf # type: ignore[import] - +def _scan_toc_from_page_texts( + page_texts: dict[int, str], + *, + page_count: int, + cross_line_window: int, +) -> list[dict[str, Any]]: matches: list[dict[str, Any]] = [] - page_count = 0 - doc = None - try: - doc = pymupdf.open(pdf_path) - page_count = doc.page_count - for page_idx in range(doc.page_count): - page_num = page_idx + 1 - text = str(doc[page_idx].get_text("text") or "") - lines = _meaningful_text_lines(text) - for match in _find_toc_text_matches( - lines, - cross_line_window=cross_line_window, - ): - matches.append({"page": page_num, **match}) - finally: - if doc is not None: - try: - doc.close() - except Exception: - pass - gc.collect() - queue.put({"ok": True, "matches": matches, "page_count": page_count}) + for page_num in range(1, page_count + 1): + text = page_texts.get(page_num, "") + lines = _meaningful_text_lines(text) + for match in _find_toc_text_matches( + lines, + cross_line_window=cross_line_window, + ): + matches.append({"page": page_num, **match}) + return matches @worker @@ -222,26 +209,21 @@ def _filter_recurring_elements( "Scan full PDF page text for TOC keywords, filter recurring " "navigation elements, then render candidate PNGs for VLM confirmation." ), - preconditions=(has_page_labels,), + preconditions=(has_page_labels, has_page_full_text), ) def find_toc_anchor_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() total_pages = ctx.blackboard.page_count - scan_timeout = int(ctx.settings.get("toc_text_scan_timeout", "180")) cross_line_window = int( ctx.settings.get("toc_cross_line_window", TOC_CROSS_LINE_WINDOW) ) - scan_result = run_in_child_process( - _scan_toc_text_worker, - ctx.pdf_path, - cross_line_window, - timeout=scan_timeout, + keyword_matches = _scan_toc_from_page_texts( + ctx.blackboard.page_full_text_cache, + page_count=total_pages, + cross_line_window=cross_line_window, ) - keyword_matches = list(scan_result.get("matches") or []) raw_hit_pages = {int(match["page"]) for match in keyword_matches} - if scan_result.get("page_count"): - total_pages = int(scan_result["page_count"]) # Apply recurring element fingerprint filter if keyword_matches: diff --git a/apps/worker/app/services/document_agent/tools/grep_text.py b/apps/worker/app/services/document_agent/tools/grep_text.py index 6c59f5268..3663c8268 100644 --- a/apps/worker/app/services/document_agent/tools/grep_text.py +++ b/apps/worker/app/services/document_agent/tools/grep_text.py @@ -7,17 +7,12 @@ from typing import Any from app.services.document_agent.manifest import ToolContext, ToolResult -from app.services.document_agent.pdf_text import read_page_texts -from app.services.document_agent.registry import has_page_features, not_is_scanned, register_tool - - -def _load_page_texts(ctx: ToolContext) -> dict[int, str]: - if ctx.blackboard.page_full_text_cache: - return dict(ctx.blackboard.page_full_text_cache) - pages = list(range(1, ctx.blackboard.page_count + 1)) - texts = read_page_texts(ctx.pdf_path, pages, timeout=300) - ctx.blackboard.page_full_text_cache = texts - return texts +from app.services.document_agent.registry import ( + has_page_features, + has_page_full_text, + not_is_scanned, + register_tool, +) @register_tool( @@ -34,7 +29,7 @@ def _load_page_texts(ctx: ToolContext) -> dict[int, str]: }, "required": ["query"], }, - preconditions=(has_page_features, not_is_scanned), + preconditions=(has_page_features, has_page_full_text, not_is_scanned), ) def grep_text(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: start = time.monotonic() @@ -52,7 +47,7 @@ def grep_text(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: flags = 0 if case_sensitive else re.IGNORECASE pattern = re.compile(query if use_regex else re.escape(query), flags) results: list[dict[str, Any]] = [] - for page, text in sorted(_load_page_texts(ctx).items()): + for page, text in sorted(ctx.blackboard.page_full_text_cache.items()): for match in pattern.finditer(text): start_idx = max(match.start() - context_chars, 0) end_idx = min(match.end() + context_chars, len(text)) diff --git a/apps/worker/app/services/document_agent/tools/ocr_pages.py b/apps/worker/app/services/document_agent/tools/ocr_pages.py new file mode 100644 index 000000000..737a8c377 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/ocr_pages.py @@ -0,0 +1,102 @@ +"""OCR specified PDF pages with RapidOCR and persist page text on the blackboard.""" + +from __future__ import annotations + +import time +from typing import Any + +from app.services.document_agent.manifest import ToolContext, ToolResult +from app.services.document_agent.registry import has_page_features, register_tool +from app.services.document_agent.visual import render_pages + + +def _line_text(item: Any) -> str: + if isinstance(item, (list, tuple)) and len(item) >= 2: + return str(item[1] or "") + return "" + + +def _line_box(item: Any) -> Any: + if isinstance(item, (list, tuple)) and len(item) >= 1: + return item[0] + return None + + +def _line_score(item: Any) -> float: + if isinstance(item, (list, tuple)) and len(item) >= 3: + try: + return float(item[2]) + except (TypeError, ValueError): + return 0.0 + return 0.0 + + +@register_tool( + name="ocr.pages", + description="Run RapidOCR on specified pages and return positioned text lines.", + parameters={ + "type": "object", + "properties": { + "pages": { + "type": "array", + "items": {"type": "integer"}, + }, + }, + "required": ["pages"], + }, + preconditions=(has_page_features,), +) +def ocr_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + raw_pages = args.get("pages") + if not isinstance(raw_pages, list) or not raw_pages: + return ToolResult( + status="error", + error="ocr.pages requires pages", + latency_ms=int((time.monotonic() - start) * 1000), + ) + pages = [int(page) for page in raw_pages] + pngs = render_pages( + ctx, + pages, + folder_name="ocr_pages", + prefix="ocr", + timeout=300, + ) + png_by_page = { + int(item["page"]): str(item["png_path"]) + for item in pngs + if item.get("page") is not None and item.get("png_path") + } + + from rapidocr_onnxruntime import RapidOCR + + engine = RapidOCR() + page_texts: dict[int, str] = {} + page_lines: dict[int, list[dict[str, Any]]] = {} + for page in pages: + image_path = png_by_page.get(page) + lines: list[dict[str, Any]] = [] + if image_path: + result, _elapse = engine(image_path) + for item in result or []: + text = _line_text(item) + lines.append( + { + "box": _line_box(item), + "text": text, + "score": _line_score(item), + } + ) + page_lines[page] = lines + page_texts[page] = "\n".join(line["text"] for line in lines if line["text"]) + + cache = dict(ctx.blackboard.page_full_text_cache) + cache.update(page_texts) + ctx.blackboard.page_full_text_cache = cache + return ToolResult( + status="ok", + payload={"page_texts": page_texts, "page_lines": page_lines}, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={"page_count": len(page_texts)}, + ) diff --git a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py index b7683f76b..7a921d17a 100644 --- a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py @@ -43,6 +43,7 @@ def build_anatomy_map(ctx: ToolContext) -> PageAnatomyMap: skeleton_anchor=ctx.blackboard.skeleton_anchor, skeleton_nodes=ctx.blackboard.skeleton_nodes, pending_skeleton_anchors=list(ctx.blackboard.pending_skeleton_anchors), + page_full_text_cache=dict(ctx.blackboard.page_full_text_cache), global_signals=ctx.blackboard.global_signals, trace_summary={ "budget": ctx.budget.snapshot(), diff --git a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index fd0e82e08..6d254450b 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -1,8 +1,7 @@ -"""LLM-guided long-PDF shard planning from document profile evidence.""" +"""Deterministic long-PDF shard planning from TOC leaf boundaries.""" from __future__ import annotations -import json import os import time from typing import Any @@ -15,7 +14,6 @@ ) from app.services.document_agent.registry import has_doc_stats, has_toc_result, register_tool from app.services.document_agent.validators import single_shard_plan, validate_shard_plan -from shared.utils.token_estimate import estimate_tokens def derive_leaf_cut_pages( @@ -60,109 +58,6 @@ def derive_leaf_cut_pages( return sorted(set(all_pages)) -def derive_chapter_boundaries( - toc_hierarchies: list[dict[str, Any]] | None, - *, - offset_override: int | None = None, - page_count: int, -) -> list[dict[str, Any]]: - """Extract chapter entries with physical page ranges for shard planning. - - Returns a flat list sorted by page_start: - [{"title": str, "level": int, "page_start": int, "page_end": int, - "page_span": int, "sub_entries": [...]}, ...] - - Includes all L1 entries. For any L1 whose span exceeds 200 pages, - its direct L2 children are included as sub_entries so the LLM can - split within it. - - Requires calibrated ``offset_override``; otherwise returns []. - """ - if not toc_hierarchies or offset_override is None: - return [] - - all_entries: list[dict[str, Any]] = [] - for hier in toc_hierarchies: - if hier.get("toc_range_unit") != "page": - continue - toc_range = hier.get("toc_range") - entries = hier.get("toc_with_level") - if not toc_range or not entries: - continue - if isinstance(entries, str): - entries = _parse_toc_with_level_entries(entries) - if not entries: - continue - - offset = offset_override - - # Collect all entries with physical pages (integer printed labels only; - # roman/prefixed need regime-local offsets — shard plan uses primary). - phys_entries: list[dict[str, Any]] = [] - for entry in entries: - pn = entry.get("page_number") - if not isinstance(pn, int): - continue - physical = pn + offset - if physical < 1 or physical > page_count: - continue - phys_entries.append({ - "title": entry.get("heading", ""), - "level": entry.get("level", 1), - "page_start": physical, - }) - - if not phys_entries: - continue - - # Compute page_end for each entry: next entry's page_start - 1 - for i, item in enumerate(phys_entries): - if i + 1 < len(phys_entries): - item["page_end"] = phys_entries[i + 1]["page_start"] - 1 - else: - item["page_end"] = page_count - item["page_span"] = item["page_end"] - item["page_start"] + 1 - - all_entries.extend(phys_entries) - - if not all_entries: - return [] - - # Build chapter-level structure: group by L1 with L2 sub_entries - min_level = min(e["level"] for e in all_entries) - chapters: list[dict[str, Any]] = [] - current_l1: dict[str, Any] | None = None - - for entry in all_entries: - if entry["level"] == min_level: - if current_l1 is not None: - chapters.append(current_l1) - current_l1 = {**entry, "sub_entries": []} - elif current_l1 is not None and entry["level"] == min_level + 1: - current_l1["sub_entries"].append(entry) - - if current_l1 is not None: - chapters.append(current_l1) - - # Recompute L1 page_end from the next L1's page_start - 1 - for i, chapter in enumerate(chapters): - if i + 1 < len(chapters): - chapter["page_end"] = chapters[i + 1]["page_start"] - 1 - else: - chapter["page_end"] = page_count - chapter["page_span"] = chapter["page_end"] - chapter["page_start"] + 1 - # Recompute sub_entry page_end within the L1's range - subs = chapter["sub_entries"] - for j, sub in enumerate(subs): - if j + 1 < len(subs): - sub["page_end"] = subs[j + 1]["page_start"] - 1 - else: - sub["page_end"] = chapter["page_end"] - sub["page_span"] = sub["page_end"] - sub["page_start"] + 1 - - return chapters - - def split_toc_for_shard( toc_hierarchies: list[dict[str, Any]] | None, shard_page_start: int, @@ -341,219 +236,33 @@ def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> return shards -def _build_chapter_prompt( - *, - page_count: int, - max_pages: int, - chapters: list[dict[str, Any]], -) -> str: - """Build LLM prompt for TOC-based shard planning using chapter boundaries.""" - chapter_list = [] - for ch in chapters: - item: dict[str, Any] = { - "title": ch["title"], - "level": ch["level"], - "page_start": ch["page_start"], - "page_end": ch["page_end"], - "page_span": ch["page_span"], - } - if ch.get("sub_entries"): - item["sub_entries"] = [ - { - "title": s["title"], - "level": s["level"], - "page_start": s["page_start"], - "page_end": s["page_end"], - "page_span": s["page_span"], - } - for s in ch["sub_entries"] - ] - chapter_list.append(item) - - payload = { - "page_count": page_count, - "max_pages_per_shard": max_pages, - "chapters": chapter_list, - } - return ( - "You are a document splitting architect. Given a PDF's chapter structure, " - "decide how to split it into shards for downstream parsing.\n" - "Rules:\n" - "- Return strict JSON only.\n" - "- Each shard must be <= max_pages_per_shard pages.\n" - "- Group adjacent chapters into shards to fill each shard as evenly as possible.\n" - "- Cut points must align with chapter boundaries (use the page_end of the last " - "chapter in the shard as cut_after_page).\n" - "- If a single chapter exceeds max_pages_per_shard, split it at one of its " - "sub_entries boundaries (use that sub_entry's page_end as cut_after_page).\n" - "- Prefer fewer shards over many small ones.\n" - "- Keep each cut rationale under 120 characters.\n" - "- If the total page_count <= max_pages_per_shard, return enabled=false.\n" - "Output schema:\n" - "{\n" - ' "enabled": boolean,\n' - ' "cuts": [\n' - ' {"cut_after_page": number, "anchor_type": "toc_chapter_boundary", ' - '"confidence": number, "rationale": string}\n' - " ],\n" - ' "reason": "llm_boundary_decision" | "not_needed",\n' - ' "rationale": string\n' - "}\n" - "Payload:\n" - + json.dumps(payload, ensure_ascii=False) - ) - - -def _sanitize_rationale(text: str, max_length: int = 120) -> str: - # Truncate overlong rationales but preserve H1 title references - # which provide valuable semantic context for shard boundaries. - sanitized = (text or "").strip() - if len(sanitized) > max_length: - sanitized = sanitized[:max_length].rstrip() + "…" - return sanitized - - -def _validate_cut_lengths( - cuts: list[tuple[int, str, str, float]], - page_count: int, - min_pages: int, - max_pages: int, -) -> None: - previous = 0 - for cut_page, *_ in cuts: - if cut_page - previous < min_pages: - raise ValueError( - f"LLM cut plan creates shard length {cut_page - previous} < min_pages={min_pages}" - ) - if cut_page - previous > max_pages: - raise ValueError( - f"LLM cut plan creates shard length {cut_page - previous} > max_pages={max_pages}" - ) - previous = cut_page - if page_count - previous < min_pages and cuts: - raise ValueError( - f"LLM cut plan creates final shard length {page_count - previous} < min_pages={min_pages}" - ) - if page_count - previous > max_pages: - raise ValueError( - f"LLM cut plan creates final shard length {page_count - previous} > max_pages={max_pages}" - ) - - -def _parse_llm_plan( - raw: str, - page_count: int, - min_pages: int, - max_pages: int, -) -> tuple[bool, list[tuple[int, str, str, float]], str, str]: - data = json.loads(raw) - if not isinstance(data, dict): - raise ValueError("LLM shard plan is not an object") - enabled = bool(data.get("enabled")) - reason = str(data.get("reason") or ("llm_boundary_decision" if enabled else "not_needed")) - rationale = _sanitize_rationale(str(data.get("rationale") or "")) - cuts: list[tuple[int, str, str, float]] = [] - for item in data.get("cuts") or []: - if not isinstance(item, dict): - continue - raw_cut_page = item.get("cut_after_page") - if raw_cut_page is None: - continue - cut_page = int(raw_cut_page) - if not 1 <= cut_page < page_count: - continue - anchor_type = str(item.get("anchor_type") or "forced_max_size") - if anchor_type not in {"h1_boundary", "blank_separator", "forced_max_size", "toc_chapter_boundary"}: - anchor_type = "forced_max_size" - confidence = float(item.get("confidence") or 0.5) - cuts.append((cut_page, anchor_type, _sanitize_rationale(str(item.get("rationale") or rationale)), confidence)) - cuts = sorted({cut[0]: cut for cut in cuts}.values(), key=lambda cut: cut[0]) - if enabled: - _validate_cut_lengths(cuts, page_count, min_pages, max_pages) - return enabled, cuts, reason, rationale - - -def _deterministic_chapter_plan( - *, - chapters: list[dict[str, Any]], - max_pages: int, - page_count: int, - leaf_pages: list[int], -) -> tuple[list[tuple[int, str, str, float]], str]: - """Greedy chapter grouping when LLM is unavailable.""" - cuts: list[tuple[int, str, str, float]] = [] - shard_start = 0 - - for i, chapter in enumerate(chapters): - chapter_end = chapter["page_end"] - shard_span = chapter_end - shard_start - - if shard_span > max_pages: - # Current chapter alone exceeds max_pages; split within its sub_entries - subs = chapter.get("sub_entries") or [] - if subs: - for sub in subs: - sub_end = sub["page_end"] - if sub_end - shard_start > max_pages: - # cut before this sub_entry - cut_page = sub["page_start"] - 1 - if cut_page > shard_start: - cuts.append(( - cut_page, - "toc_chapter_boundary", - f"split within chapter at sub-entry: {sub['title'][:60]}", - 0.7, - )) - shard_start = cut_page - else: - # No sub_entries; fall back to leaf pages within this chapter - ch_leaf_pages = [ - p for p in leaf_pages - if chapter["page_start"] <= p <= chapter_end - ] - sub_previous = shard_start - while chapter_end - sub_previous > max_pages: - target = sub_previous + max_pages - eligible = [p for p in ch_leaf_pages if sub_previous + 20 < p <= target] - if eligible: - chosen = max(eligible) - cut_page = chosen - 1 - else: - cut_page = sub_previous + max_pages - cuts.append((cut_page, "forced_max_size", "oversized chapter, leaf fallback", 0.3)) - shard_start = cut_page - sub_previous = cut_page - - elif i + 1 < len(chapters): - next_chapter_end = chapters[i + 1]["page_end"] - next_shard_span = next_chapter_end - shard_start - if next_shard_span > max_pages: - # Adding next chapter would overflow; cut after current chapter - cuts.append(( - chapter_end, - "toc_chapter_boundary", - f"chapter boundary: {chapter['title'][:60]}", - 0.85, - )) - shard_start = chapter_end - - return cuts, "too_large" +def _finest_toc_ranges(leaf_pages: list[int], page_count: int) -> list[tuple[int, int]]: + starts = sorted({page for page in leaf_pages if 1 <= page <= page_count}) + if not starts: + return [] + ranges: list[tuple[int, int]] = [] + if starts[0] > 1: + ranges.append((1, starts[0] - 1)) + for index, start in enumerate(starts): + end = starts[index + 1] - 1 if index + 1 < len(starts) else page_count + if end >= start: + ranges.append((start, end)) + return ranges -def _deterministic_no_toc_plan( +def _pack_range_by_blanks( *, - page_count: int, + previous: int, + end: int, max_pages: int, blank_pages: list[int], -) -> tuple[list[tuple[int, str, str, float]], str]: - """Deterministic shard plan using blank-like pages as split candidates.""" +) -> list[tuple[int, str, str, float]]: cuts: list[tuple[int, str, str, float]] = [] - previous = 0 - while page_count - previous > max_pages: + while end - previous > max_pages: target = previous + max_pages - # Look for a blank-like page near the max boundary eligible = [ - p for p in blank_pages if previous + (max_pages - 20) < p <= target + page for page in blank_pages + if previous + (max_pages - 20) < page <= target ] if eligible: chosen = max(eligible) @@ -563,18 +272,51 @@ def _deterministic_no_toc_plan( cut_page = previous + max_pages cuts.append((cut_page, "forced_max_size", "no separator in range", 0.2)) previous = cut_page + return cuts + + +def _deterministic_leaf_plan( + *, + page_count: int, + max_pages: int, + leaf_pages: list[int], + blank_pages: list[int], +) -> tuple[list[tuple[int, str, str, float]], str]: + ranges = _finest_toc_ranges(leaf_pages, page_count) + cuts: list[tuple[int, str, str, float]] = [] + shard_start = 0 + index = 0 + while index < len(ranges): + start, end = ranges[index] + if end - shard_start <= max_pages: + index += 1 + continue + prior_end = start - 1 + if prior_end > shard_start: + cuts.append((prior_end, "toc_leaf_boundary", f"toc leaf at page {start}", 0.85)) + shard_start = prior_end + continue + range_cuts = _pack_range_by_blanks( + previous=shard_start, + end=end, + max_pages=max_pages, + blank_pages=blank_pages, + ) + cuts.extend(range_cuts) + if range_cuts: + shard_start = range_cuts[-1][0] + index += 1 return cuts, "too_large" def _get_blank_pages(ctx: ToolContext) -> list[int]: - """Extract blank-like page numbers from page features.""" features = ctx.blackboard.page_features or [] return sorted(feature.page for feature in features if feature.is_blank_like) @register_tool( name="propose.shard_plan", - description="Decide whether and where to split a long PDF using TOC chapter boundaries.", + description="Split a long PDF at TOC leaf boundaries, then blank pages, then max page size.", preconditions=(has_doc_stats, has_toc_result), ) def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: @@ -590,84 +332,27 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) - offset_hint = ctx.blackboard.toc_page_offset - - # Try TOC chapter-based planning first - chapters = derive_chapter_boundaries( + leaf_pages = derive_leaf_cut_pages( ctx.blackboard.toc_hierarchies, - offset_override=offset_hint, - page_count=page_count, - ) if ctx.blackboard.toc_hierarchies else [] - - warnings: list[str] = [] - raw_response = "" - rationale = "" - llm_attempted = False - leaf_pages = derive_leaf_cut_pages(ctx.blackboard.toc_hierarchies, offset_override=offset_hint) - - if chapters: - # Path A: TOC chapter-based LLM decision - model = ctx.settings.get("model") - prompt = _build_chapter_prompt( + offset_override=ctx.blackboard.toc_page_offset, + ) + blank_pages = _get_blank_pages(ctx) + if leaf_pages: + cuts, reason = _deterministic_leaf_plan( page_count=page_count, max_pages=max_pages, - chapters=chapters, + leaf_pages=leaf_pages, + blank_pages=blank_pages, ) - prompt_tokens_est = estimate_tokens(prompt) - - if model and ctx.budget.try_reserve("plan", prompt_tokens_est): - try: - llm_attempted = True - from shared.services.ai.llm_overrides import get_text_client - - client, model = get_text_client(requested_model=model) - raw_response, usage = client.chat_completion_with_usage( - messages=[{"role": "user", "content": prompt}], - model=model, - temperature=0.0, - max_tokens=1600, - response_format={"type": "json_object"}, - usage_task="document_agent.propose_shard_plan", - ) - ctx.budget.commit("plan", actual=usage.get("total_tokens", prompt_tokens_est), est=prompt_tokens_est) - enabled, cuts, reason, rationale = _parse_llm_plan(raw_response, page_count, min_pages, max_pages) - if not enabled: - cuts = [] - reason = "not_needed" - except Exception as exc: - ctx.budget.refund("plan", est=prompt_tokens_est) - warnings.append(f"LLM chapter shard decision failed; using deterministic plan: {exc}") - ctx.blackboard.global_signals.setdefault("degraded_reasons", []).append( - "shard_plan: llm_parse_failed" - ) - cuts, reason = _deterministic_chapter_plan( - chapters=chapters, - max_pages=max_pages, - page_count=page_count, - leaf_pages=leaf_pages, - ) - rationale = "Deterministic chapter plan after LLM failure." - else: - if not model: - warnings.append("No model configured; using deterministic chapter plan.") - ctx.blackboard.global_signals.setdefault("degraded_reasons", []).append( - "shard_plan: no model" - ) - cuts, reason = _deterministic_chapter_plan( - chapters=chapters, - max_pages=max_pages, - page_count=page_count, - leaf_pages=leaf_pages, - ) - rationale = "Deterministic chapter plan (no LLM)." + rationale = "Deterministic plan from TOC leaf boundaries." else: - # Path B: No TOC — purely deterministic using blank-like pages - blank_pages = _get_blank_pages(ctx) - cuts, reason = _deterministic_no_toc_plan( - page_count=page_count, + cuts = _pack_range_by_blanks( + previous=0, + end=page_count, max_pages=max_pages, blank_pages=blank_pages, ) + reason = "too_large" rationale = "Deterministic plan from blank-like page boundaries (no TOC)." shards = _cuts_to_shards(cuts, page_count) @@ -695,12 +380,9 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: "valid": plan.validation.valid, }, latency_ms=int((time.monotonic() - start) * 1000), - tokens_used=ctx.budget.snapshot()["plan"]["used"] if llm_attempted else 0, input_summary={ "page_count": page_count, - "chapter_count": len(chapters), "leaf_page_count": len(leaf_pages), - "model": ctx.settings.get("model"), }, output_summary={ "enabled": plan.enabled, @@ -708,9 +390,4 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: "rationale": rationale, "shards": [shard.to_dict() for shard in plan.shards], }, - warnings=warnings, - debug={ - "raw_response_excerpt": raw_response[:4000] if raw_response else "", - "llm_attempted": llm_attempted, - }, ) diff --git a/apps/worker/app/services/document_agent/visual.py b/apps/worker/app/services/document_agent/visual.py index f25bf70d6..14a39578f 100644 --- a/apps/worker/app/services/document_agent/visual.py +++ b/apps/worker/app/services/document_agent/visual.py @@ -21,6 +21,7 @@ "toc_pages", "verify_pages", "agent_visuals", + "ocr_pages", } _PAGE_MEMORY_VISUAL_DIRS = {"pages", "asset_annotate"} diff --git a/apps/worker/app/services/document_parser/profiling/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py index a0a3077bd..5d5a42028 100644 --- a/apps/worker/app/services/document_parser/profiling/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -129,9 +129,8 @@ def _profile_pdf_with_db( settings={ "planner_model": settings.IMAGE_MODEL, "vlm_model": settings.IMAGE_MODEL, - "model": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, "toc_profile_enabled": page_toc_enabled, - "toc_before_coarse": page_toc_enabled, + "model": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, }, ) agent_profile = coordinator.run_coarse() @@ -153,6 +152,7 @@ def _profile_pdf_with_db( {}, ), }, + page_full_text_cache=dict(coordinator.blackboard.page_full_text_cache), ) if profile.page_count > settings.MAX_PDF_PAGE_LIMIT: if oversized_policy != "page_memory": diff --git a/apps/worker/app/services/document_parser/profiling/profile_model.py b/apps/worker/app/services/document_parser/profiling/profile_model.py index 79df61468..2c557f56a 100644 --- a/apps/worker/app/services/document_parser/profiling/profile_model.py +++ b/apps/worker/app/services/document_parser/profiling/profile_model.py @@ -45,6 +45,7 @@ class ParserDocumentProfile: granularity: str = "page" anatomy: Any | None = None metrics: dict[str, Any] = field(default_factory=dict) + page_full_text_cache: dict[int, str] = field(default_factory=dict) @property def is_pdf(self) -> bool: @@ -60,6 +61,7 @@ def has_structural_anatomy(self) -> bool: def to_dict(self) -> dict[str, Any]: data = asdict(self) + data.pop("page_full_text_cache", None) data["routing_category"] = self.routing_category.value if self.anatomy is not None and hasattr(self.anatomy, "to_dict"): data["anatomy"] = self.anatomy.to_dict() diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index da5f8602b..30ca70088 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -8,7 +8,7 @@ import pandas as pd -from app.services.document_agent.pdf_text import read_page_texts +from app.services.document_agent.pdf_text import coerce_page_text_cache from app.services.document_agent.visual import ( purge_debug_visual_dirs, visual_debug_enabled, @@ -117,11 +117,13 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: if verdict == "whole_doc": parsed_df = _build_whole_doc_dataframe( - pdf_path=pdf_path, filename=request.filename, page_count=page_count, verdict=verdict, trace_recorder=trace_recorder, + page_texts=coerce_page_text_cache( + getattr(profile, "page_full_text_cache", None) + ), ) else: parsed_df = _build_page_dataframe( @@ -237,10 +239,7 @@ def _build_page_dataframe( asset_extraction_enabled = page_memory_config.asset_extraction_enabled # ── C4: skeleton (from profile anatomy) ─────────────────────────── - with stage_timer("page_memory.read_page_texts", page_count=page_count): - page_texts = read_page_texts( - pdf_path, list(range(1, page_count + 1)), timeout=300, - ) + page_texts = coerce_page_text_cache(getattr(profile, "page_full_text_cache", None)) with stage_timer("page_memory.skeleton", page_count=page_count): if anatomy is not None: skeletons = extract_section_skeletons( @@ -800,15 +799,15 @@ def _select_rendered_pages_with_assets( def _build_whole_doc_dataframe( *, - pdf_path: str, filename: str, page_count: int, verdict: str, trace_recorder: Any | None = None, + page_texts: dict[int, str] | None = None, ) -> pd.DataFrame: pages = list(range(1, page_count + 1)) if page_count > 0 else [1] - page_texts = read_page_texts(pdf_path, pages) - raw_text = "\n\n".join(page_texts.get(page, "") for page in pages).strip() + texts = coerce_page_text_cache(page_texts) + raw_text = "\n\n".join(texts.get(page, "") for page in pages).strip() summary = _build_summary(filename=filename, page_count=page_count, raw_text=raw_text) content = f"[SUMMARY]\n{summary}\n\n[RAW]\n{raw_text}".strip() know_id = gen_str_codes(f"wholedoc::{filename}::{content}") diff --git a/apps/worker/app/services/page_memory/page_renderer.py b/apps/worker/app/services/page_memory/page_renderer.py index 89a0a7567..3019cc2ee 100644 --- a/apps/worker/app/services/page_memory/page_renderer.py +++ b/apps/worker/app/services/page_memory/page_renderer.py @@ -1,8 +1,8 @@ """Page renderer: produce PNG and raw text for each page. -Wraps the existing ``document_agent/visual.render_pages`` and -``pdf_text.read_page_texts`` utilities, adding dimensions and -landscape detection from page_features. +Wraps the existing ``document_agent/visual.render_pages`` utility, +adding dimensions and landscape detection from page_features. +Page text comes from the PROFILE scan cache passed in as ``page_texts``. """ from __future__ import annotations @@ -12,7 +12,6 @@ from loguru import logger from app.services.document_agent.manifest import PageFeature, ToolContext -from app.services.document_agent.pdf_text import read_page_texts from app.services.document_agent.visual import render_pages @@ -27,7 +26,7 @@ class PageRenderResult: """Absolute path to full-resolution PNG (``pages/page-N.png``).""" raw_text: str - """PyMuPDF extracted text for this page.""" + """PROFILE scan text for this page.""" width: float """Page width in points.""" @@ -67,9 +66,8 @@ def render_document_pages( If available, dimensions are read from here (avoiding a second PyMuPDF open). Otherwise falls back to 0/0/False. page_texts: - Pre-read page texts ``{page_index: text}``. If provided, skips - the internal ``read_page_texts`` call (avoids a redundant - PyMuPDF subprocess). + PROFILE scan cache ``{page_index: text}``. Missing pages yield + empty ``raw_text``. ctx: An optional ``ToolContext`` forwarded to ``render_pages``. If *None*, a lightweight context is constructed internally. @@ -91,9 +89,7 @@ def render_document_pages( if not requested_pages: return [] - # ── raw text (reuse caller's data if provided) ──────────────────── - if page_texts is None: - page_texts = read_page_texts(pdf_path, requested_pages, timeout=timeout) + texts = page_texts or {} # ── full-resolution PNGs ────────────────────────────────────────── if ctx is not None: @@ -146,7 +142,7 @@ def render_document_pages( PageRenderResult( page_index=page, image_path=png_map.get(page, ""), - raw_text=page_texts.get(page, ""), + raw_text=texts.get(page, ""), width=feat.width if feat else 0.0, height=feat.height if feat else 0.0, is_landscape=( diff --git a/apps/worker/pyproject.toml b/apps/worker/pyproject.toml index 7c4e085d6..d6344112e 100644 --- a/apps/worker/pyproject.toml +++ b/apps/worker/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ "psycogreen>=1.0.2", "httpcore>=1.0.6", "tabula-py>=2.10.0", + "rapidocr-onnxruntime>=1.4.4", ] [dependency-groups] diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index ae11c7750..aeb1007f2 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -143,6 +143,7 @@ def test_run_lightweight_anatomy_builds_single_shard_without_planner_llm( routing_category=PdfRoutingCategory.GENERIC.value, ) coordinator.blackboard.toc_result = TocResult(method="none") + coordinator.blackboard.page_full_text_cache = {1: "hello", 2: "world"} anatomy = coordinator.run_lightweight_anatomy() @@ -156,6 +157,8 @@ def test_run_lightweight_anatomy_builds_single_shard_without_planner_llm( (output_dir / "anatomy_map.json").read_text(encoding="utf-8") ) assert list(anatomy_data)[:2] == ["version", "toc_hierarchies"] + assert anatomy.page_full_text_cache == {1: "hello", 2: "world"} + assert anatomy_data["page_full_text_cache"] == {"1": "hello", "2": "world"} assert "text_lines_preview" not in anatomy_data["page_features"][0] assert "asset_bboxes" not in anatomy_data["page_features"][0] trace_data = json.loads((output_dir / "trace.json").read_text(encoding="utf-8")) @@ -201,14 +204,24 @@ def fake_aggregate(ctx, _args): calls.append("aggregate.doc_stats") return ToolResult(status="ok", payload={}) + def fake_text_scan() -> None: + calls.append("text_scan") + coordinator.blackboard.page_full_text_cache = {1: "a", 2: "b"} + + def fake_toc() -> None: + calls.append("toc") + coordinator.blackboard.toc_result = TocResult(method="none") + monkeypatch.setattr(coordinator_module.ProfilePlanner, "propose", fake_propose) monkeypatch.setattr(coordinator_module, "probe_page_assets", fake_probe_page_assets) monkeypatch.setattr(coordinator_module, "aggregate_doc_stats", fake_aggregate) + monkeypatch.setattr(coordinator, "_run_text_scan", fake_text_scan) + monkeypatch.setattr(coordinator, "_run_toc_extraction_pipeline", fake_toc) profile = coordinator.run_coarse() assert profile.category == "Research Report" - assert calls == ["planner", "probe.page_assets", "aggregate.doc_stats"] + assert calls == ["planner", "text_scan", "toc", "probe.page_assets", "aggregate.doc_stats"] assert coordinator.blackboard.global_signals["assets_probed"] is True @@ -539,15 +552,15 @@ def run(self): assert anatomy.toc_result.toc_pages == [] -def test_run_coarse_runs_toc_before_planner_for_oversized_and_reuses_planner( +def test_run_coarse_runs_planner_then_text_scan_then_toc( monkeypatch, tmp_path: Path, ) -> None: coordinator = ProfileCoordinator( pdf_path=str(tmp_path / "oversized.pdf"), - job_id="job-toc-before-coarse", + job_id="job-planner-then-toc", output_dir=str(tmp_path / "profile"), - settings={"toc_before_coarse": True}, + settings={"toc_profile_enabled": True}, ) (tmp_path / "profile").mkdir() _seed_preprobed_pages(coordinator, page_count=3, pages=[1, 2]) @@ -561,10 +574,15 @@ def fake_toc_extraction() -> None: {"toc_range": [17, 17], "toc_range_unit": "page", "toc_tree": {}} ] + def fake_text_scan() -> None: + calls.append("text_scan") + coordinator.blackboard.page_full_text_cache = {1: "a", 2: "b", 3: "c"} + def fake_persist(_anatomy): calls.append("persist") monkeypatch.setattr(coordinator, "_run_toc_extraction_pipeline", fake_toc_extraction) + monkeypatch.setattr(coordinator, "_run_text_scan", fake_text_scan) monkeypatch.setattr(coordinator, "_persist_ready_anatomy", fake_persist) def fake_propose(_self): @@ -612,10 +630,62 @@ def run(self): coordinator.run_coarse() anatomy = coordinator.run_structural() - assert calls == ["toc", "planner", "persist"] + assert calls == ["planner", "text_scan", "toc", "persist"] assert anatomy.toc_result.toc_pages == [17] +def test_run_text_scan_native_uses_read_page_texts( + monkeypatch, + tmp_path: Path, +) -> None: + coordinator = ProfileCoordinator( + pdf_path=str(tmp_path / "doc.pdf"), + job_id="job-scan-native", + output_dir=str(tmp_path / "profile"), + ) + coordinator.blackboard.page_count = 2 + coordinator.blackboard.document_profile = DocumentProfile( + is_scanned=False, + category="Report", + routing_category=PdfRoutingCategory.GENERIC.value, + ) + + def fake_read(_pdf_path: str, pages: list[int], timeout: int = 300) -> dict[int, str]: + assert pages == [1, 2] + return {1: "a", 2: "b"} + + monkeypatch.setattr(coordinator_module, "read_page_texts", fake_read) + coordinator._run_text_scan() + assert coordinator.blackboard.page_full_text_cache == {1: "a", 2: "b"} + + +def test_run_text_scan_scanned_dispatches_ocr_pages( + monkeypatch, + tmp_path: Path, +) -> None: + coordinator = ProfileCoordinator( + pdf_path=str(tmp_path / "doc.pdf"), + job_id="job-scan-ocr", + output_dir=str(tmp_path / "profile"), + ) + coordinator.blackboard.page_count = 2 + coordinator.blackboard.document_profile = DocumentProfile( + is_scanned=True, + category="Report", + routing_category=PdfRoutingCategory.GENERIC.value, + ) + + def fake_dispatch(name: str, ctx, args): + assert name == "ocr.pages" + assert args == {"pages": [1, 2]} + ctx.blackboard.page_full_text_cache = {1: "ocr-1", 2: "ocr-2"} + return ToolResult(status="ok", payload={}) + + monkeypatch.setattr(coordinator_module.REGISTRY, "dispatch", fake_dispatch) + coordinator._run_text_scan() + assert coordinator.blackboard.page_full_text_cache == {1: "ocr-1", 2: "ocr-2"} + + def test_anchor_confirmation_failure_requires_one_strict_retry(tmp_path: Path) -> None: coordinator = ProfileCoordinator( pdf_path=str(tmp_path / "oversized.pdf"), @@ -682,6 +752,7 @@ def __init__(self, **kwargs) -> None: global_signals={}, toc_result=None, toc_hierarchies=None, + page_full_text_cache={}, ) fake_instances.append(self) @@ -714,7 +785,7 @@ def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): assert profile.anatomy is fake_anatomy assert fake_instances[0].calls == ["run_coarse", "run_lightweight_anatomy"] assert init_settings[0]["toc_profile_enabled"] is True - assert init_settings[0]["toc_before_coarse"] is True + assert "toc_before_coarse" not in init_settings[0] def test_standard_pdf_page_toc_kill_switch_builds_no_toc_anatomy( @@ -734,6 +805,7 @@ def __init__(self, **kwargs) -> None: global_signals={}, toc_result=None, toc_hierarchies=None, + page_full_text_cache={}, ) def run_coarse(self) -> DocumentProfile: @@ -768,7 +840,7 @@ def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): ) assert init_settings[0]["toc_profile_enabled"] is False - assert init_settings[0]["toc_before_coarse"] is False + assert "toc_before_coarse" not in init_settings[0] assert profile.toc.attempted is False assert profile.toc.has_toc is False assert profile.toc.notes == "TOC profiling disabled by PDF_PROFILE_TOC_ENABLED" @@ -792,6 +864,7 @@ def __init__(self, **kwargs) -> None: global_signals={}, toc_result=None, toc_hierarchies=None, + page_full_text_cache={}, ) def run_coarse(self) -> DocumentProfile: @@ -822,7 +895,7 @@ def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): ) assert init_settings[0]["toc_profile_enabled"] is True - assert init_settings[0]["toc_before_coarse"] is True + assert "toc_before_coarse" not in init_settings[0] assert profile.anatomy is fake_anatomy @@ -842,6 +915,7 @@ def __init__(self, **_kwargs) -> None: global_signals={}, toc_result=None, toc_hierarchies=None, + page_full_text_cache={}, ) fake_instances.append(self) @@ -891,6 +965,7 @@ def __init__(self, **_kwargs) -> None: global_signals={}, toc_result=None, toc_hierarchies=None, + page_full_text_cache={}, ) def run_coarse(self) -> DocumentProfile: @@ -918,7 +993,7 @@ def run_coarse(self) -> DocumentProfile: def run_toc(self) -> TocResult: self.calls.append("run_toc") - raise AssertionError("run_toc should be no-op after TOC-before-coarse") + raise AssertionError("run_toc should be no-op after coarse already ran TOC") def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): self.calls.append("run_lightweight_anatomy") @@ -966,6 +1041,7 @@ def __init__(self, **_kwargs) -> None: global_signals={}, toc_result=None, toc_hierarchies=None, + page_full_text_cache={}, ) def run_coarse(self) -> DocumentProfile: diff --git a/apps/worker/tests/contract/test_ocr_pages_contract.py b/apps/worker/tests/contract/test_ocr_pages_contract.py new file mode 100644 index 000000000..55db1c7ce --- /dev/null +++ b/apps/worker/tests/contract/test_ocr_pages_contract.py @@ -0,0 +1,78 @@ +"""ocr.pages writes RapidOCR page texts onto the profile blackboard.""" + +from __future__ import annotations + +import os +import sys +from types import ModuleType +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.budget import BudgetTracker +from app.services.document_agent.manifest import PageFeature, ToolContext +from app.services.document_agent.state import AgentBlackboard +from app.services.document_agent.tools.ocr_pages import ocr_pages + + +def _ctx() -> ToolContext: + blackboard = AgentBlackboard(page_count=2) + blackboard.page_features = [ + PageFeature( + page=1, + raw_text_length=0, + text_density=0.0, + image_coverage=1.0, + image_count=1, + table_count=0, + drawings_count=0, + orientation="portrait", + width=72.0, + height=72.0, + has_asset=True, + is_blank_like=True, + ) + ] + return ToolContext( + pdf_path="/tmp/doc.pdf", + job_id="job-ocr", + blackboard=blackboard, + budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), + trace=None, + settings={}, + ) + + +def test_ocr_pages_requires_pages() -> None: + result = ocr_pages(_ctx(), {}) + assert result.status == "error" + assert "requires pages" in (result.error or "") + + +def test_ocr_pages_writes_joined_text_to_blackboard() -> None: + ctx = _ctx() + + class FakeEngine: + def __call__(self, _image_path: str): + return [[[[0, 0], [1, 0], [1, 1], [0, 1]], "Hello", 0.9]], 0.01 + + fake_mod = ModuleType("rapidocr_onnxruntime") + fake_mod.RapidOCR = lambda: FakeEngine() # type: ignore[attr-defined] + + with ( + patch( + "app.services.document_agent.tools.ocr_pages.render_pages", + return_value=[{"page": 1, "png_path": "/tmp/ocr_page_1.png"}], + ), + patch.dict(sys.modules, {"rapidocr_onnxruntime": fake_mod}), + ): + result = ocr_pages(ctx, {"pages": [1]}) + + assert result.status == "ok" + assert ctx.blackboard.page_full_text_cache[1] == "Hello" + assert result.payload["page_lines"][1][0]["text"] == "Hello" diff --git a/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py b/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py index 388ca1d55..4ca5d3756 100644 --- a/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py +++ b/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py @@ -145,19 +145,14 @@ def test_profile_toc_anchoring_writes_skeleton_anchor() -> None: ctx = _ctx() ctx.blackboard.toc_hierarchies = _toc() ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1]) + ctx.blackboard.page_full_text_cache = {page: "Ch1" for page in range(1, 11)} def fake_anchor_hierarchy(**_kwargs): return [_node()], _anchor() - with ( - patch( - "app.services.document_agent.structure.toc_anchoring.read_page_texts", - return_value={page: "Ch1" for page in range(1, 11)}, - ), - patch( - "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", - side_effect=fake_anchor_hierarchy, - ), + with patch( + "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + side_effect=fake_anchor_hierarchy, ): run_toc_anchoring(ctx) @@ -256,6 +251,7 @@ def test_profile_classifies_pending_toc_before_finalize() -> None: ctx = _ctx(page_count=30) ctx.blackboard.toc_hierarchies = _pending_tocs() ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1, 20, 21]) + ctx.blackboard.page_full_text_cache = {page: "body" for page in range(1, 31)} pending_node = TitleNode(title="App", level=1, printed_page=22, children=[]) finalize_calls: list[object] = [] @@ -264,10 +260,6 @@ def fake_finalize(**kwargs): return [pending_node], _anchor(title="App", page=22), True with ( - patch( - "app.services.document_agent.structure.toc_anchoring.read_page_texts", - return_value={page: "body" for page in range(1, 31)}, - ), patch( "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", return_value=([_node()], _anchor()), @@ -313,15 +305,12 @@ def test_profile_skips_finalize_for_unresolvable_pending_toc() -> None: hierarchies[1]["toc_with_level"] = [{"heading": "App", "level": 1}] ctx.blackboard.toc_hierarchies = hierarchies ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1, 20, 21]) + ctx.blackboard.page_full_text_cache = {page: "body" for page in range(1, 31)} def _boom(*_args, **_kwargs): raise AssertionError("unresolvable pending TOC must not finalize") with ( - patch( - "app.services.document_agent.structure.toc_anchoring.read_page_texts", - return_value={page: "body" for page in range(1, 31)}, - ), patch( "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", return_value=([_node()], _anchor()), @@ -427,3 +416,33 @@ def test_classify_toc_relationship_is_not_on_c4_module() -> None: assert not hasattr(skeleton_extractor, "_classify_toc_relationship") assert callable(classify_toc_relationship) + + +def test_toc_anchoring_requires_page_text_cache() -> None: + ctx = _ctx() + ctx.blackboard.toc_hierarchies = _toc() + ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1]) + + try: + run_toc_anchoring(ctx) + raise AssertionError("expected missing cache to raise") + except ValueError as exc: + assert "page_full_text_cache" in str(exc) + + +def test_page_memory_and_toc_do_not_extract_page_texts() -> None: + import inspect + + from app.services.document_agent.structure import toc_anchoring + from app.services.document_agent.tools import find_toc_anchor_pages, grep_text + from app.services.page_memory import memory_service, page_renderer, skeleton_extractor + + for module in ( + memory_service, + page_renderer, + skeleton_extractor, + toc_anchoring, + find_toc_anchor_pages, + grep_text, + ): + assert "read_page_texts(" not in inspect.getsource(module) diff --git a/apps/worker/tests/contract/test_propose_shard_plan_contract.py b/apps/worker/tests/contract/test_propose_shard_plan_contract.py new file mode 100644 index 000000000..21d1e1e9a --- /dev/null +++ b/apps/worker/tests/contract/test_propose_shard_plan_contract.py @@ -0,0 +1,136 @@ +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 app.services.document_agent.budget import BudgetTracker +from app.services.document_agent.manifest import PageFeature, TocResult, ToolContext +from app.services.document_agent.state import AgentBlackboard +from app.services.document_agent.tools.propose_shard_plan import propose_shard_plan + + +def _feature(page: int, *, blank: bool = False) -> PageFeature: + return PageFeature( + page=page, + raw_text_length=0 if blank else 20, + text_density=0.1, + image_coverage=0.0, + image_count=0, + table_count=0, + drawings_count=0, + orientation="portrait", + width=72.0, + height=72.0, + has_asset=False, + is_blank_like=blank, + ) + + +def _ctx(*, page_count: int, blank_pages: list[int] | None = None) -> ToolContext: + blanks = set(blank_pages or []) + ctx = ToolContext( + pdf_path="/tmp/doc.pdf", + job_id="job-shard", + blackboard=AgentBlackboard(page_count=page_count), + budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), + trace=None, + settings={ + "shard_threshold": 200, + "max_pages_per_shard": 200, + "min_pages_per_shard": 20, + }, + ) + ctx.blackboard.doc_stats = {"page_count": page_count} + ctx.blackboard.toc_result = TocResult(method="vlm_batch") + ctx.blackboard.page_features = [ + _feature(page, blank=page in blanks) for page in range(1, page_count + 1) + ] + return ctx + + +def test_leaf_plan_cuts_at_finest_toc_boundary() -> None: + ctx = _ctx(page_count=250) + ctx.blackboard.toc_page_offset = 0 + ctx.blackboard.toc_hierarchies = [ + { + "toc_range": [1, 2], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 3}, + {"heading": "1.1", "level": 2, "page_number": 3}, + {"heading": "1.2", "level": 2, "page_number": 80}, + {"heading": "Ch2", "level": 1, "page_number": 120}, + ], + } + ] + + result = propose_shard_plan(ctx, {}) + plan = ctx.blackboard.shard_plan + + assert result.status == "ok" + assert plan is not None + assert plan.enabled is True + assert [(shard.page_start, shard.page_end, shard.anchor_type) for shard in plan.shards] == [ + (1, 119, "toc_leaf_boundary"), + (120, 250, "forced_max_size"), + ] + assert plan.validation.valid is True + + +def test_fat_leaf_uses_blank_page_in_window() -> None: + ctx = _ctx(page_count=450, blank_pages=[195]) + ctx.blackboard.toc_page_offset = 0 + ctx.blackboard.toc_hierarchies = [ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [{"heading": "Only", "level": 1, "page_number": 1}], + } + ] + + propose_shard_plan(ctx, {}) + plan = ctx.blackboard.shard_plan + assert plan is not None + assert plan.shards[0].page_end == 195 + assert plan.shards[0].anchor_type == "blank_separator" + assert all(shard.page_end - shard.page_start + 1 <= 200 for shard in plan.shards) + assert plan.validation.valid is True + + +def test_fat_leaf_without_blank_uses_forced_max_size() -> None: + ctx = _ctx(page_count=450) + ctx.blackboard.toc_page_offset = 0 + ctx.blackboard.toc_hierarchies = [ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [{"heading": "Only", "level": 1, "page_number": 1}], + } + ] + + propose_shard_plan(ctx, {}) + plan = ctx.blackboard.shard_plan + assert plan is not None + assert [(shard.page_start, shard.page_end, shard.anchor_type) for shard in plan.shards] == [ + (1, 200, "forced_max_size"), + (201, 400, "forced_max_size"), + (401, 450, "forced_max_size"), + ] + assert plan.validation.valid is True + + +def test_no_toc_leaves_uses_blank_pages_on_full_document() -> None: + ctx = _ctx(page_count=450, blank_pages=[195]) + + propose_shard_plan(ctx, {}) + plan = ctx.blackboard.shard_plan + assert plan is not None + assert plan.shards[0].page_end == 195 + assert plan.shards[0].anchor_type == "blank_separator" + assert plan.validation.valid is True diff --git a/uv.lock b/uv.lock index ea319a249..ff85b4332 100644 --- a/uv.lock +++ b/uv.lock @@ -1594,6 +1594,7 @@ dependencies = [ { name = "pypdf" }, { name = "python-docx" }, { name = "python-pptx" }, + { name = "rapidocr-onnxruntime" }, { name = "tabula-py" }, { name = "tqdm" }, ] @@ -1632,6 +1633,7 @@ requires-dist = [ { name = "pypdf", specifier = "==6.10.2" }, { name = "python-docx", specifier = "==1.2.0" }, { name = "python-pptx", specifier = "==1.0.2" }, + { name = "rapidocr-onnxruntime", specifier = ">=1.4.4" }, { name = "tabula-py", specifier = ">=2.10.0" }, { name = "tqdm", specifier = "==4.67.1" }, ] @@ -2335,6 +2337,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8b/b9/0df6351b25c6bd494c534d2a8191dc9460fb5bb09c88b1427775d49fde05/openai-1.93.3-py3-none-any.whl", hash = "sha256:41aaa7594c7d141b46eed0a58dcd75d20edcc809fdd2c931ecbb4957dc98a892", size = 755132, upload-time = "2025-07-09T14:08:25.533Z" }, ] +[[package]] +name = "opencv-python" +version = "5.0.0.93" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/79/4c/a438d23e09ce2033c09f7b784ad2fbdb0adf529e434101ed28f142226f98/opencv_python-5.0.0.93.tar.gz", hash = "sha256:66aac3e5b5faa48d4025816592f3af19e4bfc2c68dec067bae2dbb4ca10aa9e2", size = 81802749, upload-time = "2026-07-02T06:59:53.815Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/75/76f6ade78f6102c61034f828e2a22616708df2c9504bc8d6af9dd8f73dc5/opencv_python-5.0.0.93-cp37-abi3-macosx_13_0_arm64.whl", hash = "sha256:198a75138241810206a17c829dbcc40a7cb1841cda538ca86cbbfc6c7d95f898", size = 48322443, upload-time = "2026-07-02T05:50:25.466Z" }, + { url = "https://files.pythonhosted.org/packages/15/8c/bc1bda6aae69a32e9d84fc34153ba104cd25226861eb4aea33b2cea4860d/opencv_python-5.0.0.93-cp37-abi3-macosx_14_0_x86_64.whl", hash = "sha256:6bbc32f59e1b1a7db7b39c81f63d00625f041d333037fd8702f6da52cc39108b", size = 34782755, upload-time = "2026-07-02T05:51:30.556Z" }, + { url = "https://files.pythonhosted.org/packages/f4/8a/b04776ec45d2dea08a1b176f1829201db3515d4ed16c35f8fcc9fa7beb16/opencv_python-5.0.0.93-cp37-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e2b4272e736836f66c2d176e43ab8101f3a00d45654916399f52e150c58981ac", size = 50614064, upload-time = "2026-07-02T06:53:22.604Z" }, + { url = "https://files.pythonhosted.org/packages/95/54/eb47866b94f2b5b42dde17644b78055ef1ee05aae59962c7290e55270803/opencv_python-5.0.0.93-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f8b6d0a212253dd26ad338c812f1f23ca118fdf05a9c8c6b9444f161aa8c5881", size = 71064711, upload-time = "2026-07-02T06:54:13.148Z" }, + { url = "https://files.pythonhosted.org/packages/93/da/962579f1e703cbf8c5422fd1f576467dcb3b5b0b0b81c1471c979764353a/opencv_python-5.0.0.93-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:08d5d91d967b58d6db86073b2ad3eaef88ca4ebdfd45c9059bf59f5ded0c7ad2", size = 49798576, upload-time = "2026-07-02T06:54:33.781Z" }, + { url = "https://files.pythonhosted.org/packages/cf/4c/c73f828fdbcd37eaf21d08fa852544a3ca7c2dbb3ea76873d64f2ea413d1/opencv_python-5.0.0.93-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:c8de2dec111122a02e8beb28e16c31904992dfd6186560b142a92c71403c1039", size = 73783032, upload-time = "2026-07-02T06:55:03.415Z" }, + { url = "https://files.pythonhosted.org/packages/e2/4b/edaf83b996ca5a1a3d8ccad485706b9c6d4742b13b9c4586bf1c1e7d9423/opencv_python-5.0.0.93-cp37-abi3-win32.whl", hash = "sha256:4b4b1a34c79bf8d3738e3cfe9a9e67b51a79663f6b692cbdad8c31f570da4157", size = 35564734, upload-time = "2026-07-02T05:49:57.704Z" }, + { url = "https://files.pythonhosted.org/packages/21/f0/9fa6e85cb10c8eb36a0222d27e50fe381b86ce49a55446bf39f491727564/opencv_python-5.0.0.93-cp37-abi3-win_amd64.whl", hash = "sha256:f90ba04b8f73bc5c3814037699739f0156f597338a98f05956c684e7c3ca10d2", size = 44000345, upload-time = "2026-07-02T05:49:54.971Z" }, +] + [[package]] name = "openpyxl" version = "3.1.2" @@ -2972,6 +2993,43 @@ bcrypt = [ { name = "bcrypt" }, ] +[[package]] +name = "pyclipper" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/21/3c06205bb407e1f79b73b7b4dfb3950bd9537c4f625a68ab5cc41177f5bc/pyclipper-1.4.0.tar.gz", hash = "sha256:9882bd889f27da78add4dd6f881d25697efc740bf840274e749988d25496c8e1", size = 54489, upload-time = "2025-12-01T13:15:35.015Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/e3/64cf7794319b088c288706087141e53ac259c7959728303276d18adc665d/pyclipper-1.4.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:adcb7ca33c5bdc33cd775e8b3eadad54873c802a6d909067a57348bcb96e7a2d", size = 264281, upload-time = "2025-12-01T13:14:55.47Z" }, + { url = "https://files.pythonhosted.org/packages/34/cd/44ec0da0306fa4231e76f1c2cb1fa394d7bde8db490a2b24d55b39865f69/pyclipper-1.4.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:fd24849d2b94ec749ceac7c34c9f01010d23b6e9d9216cf2238b8481160e703d", size = 139426, upload-time = "2025-12-01T13:14:56.683Z" }, + { url = "https://files.pythonhosted.org/packages/ad/88/d8f6c6763ea622fe35e19c75d8b39ed6c55191ddc82d65e06bc46b26cb8e/pyclipper-1.4.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1b6c8d75ba20c6433c9ea8f1a0feb7e4d3ac06a09ad1fd6d571afc1ddf89b869", size = 989649, upload-time = "2025-12-01T13:14:58.28Z" }, + { url = "https://files.pythonhosted.org/packages/ff/e9/ea7d68c8c4af3842d6515bedcf06418610ad75f111e64c92c1d4785a1513/pyclipper-1.4.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:58e29d7443d7cc0e83ee9daf43927730386629786d00c63b04fe3b53ac01462c", size = 962842, upload-time = "2025-12-01T13:15:00.044Z" }, + { url = "https://files.pythonhosted.org/packages/4e/b7/0b4a272d8726e51ab05e2b933d8cc47f29757fb8212e38b619e170e6015c/pyclipper-1.4.0-cp311-cp311-win32.whl", hash = "sha256:a8d2b5fb75ebe57e21ce61e79a9131edec2622ff23cc665e4d1d1f201bc1a801", size = 95098, upload-time = "2025-12-01T13:15:01.359Z" }, + { url = "https://files.pythonhosted.org/packages/3a/76/4901de2919198bb2bd3d989f86d4a1dff363962425bb2d63e24e6c990042/pyclipper-1.4.0-cp311-cp311-win_amd64.whl", hash = "sha256:e9b973467d9c5fa9bc30bb6ac95f9f4d7c3d9fc25f6cf2d1cc972088e5955c01", size = 104362, upload-time = "2025-12-01T13:15:02.439Z" }, + { url = "https://files.pythonhosted.org/packages/90/1b/7a07b68e0842324d46c03e512d8eefa9cb92ba2a792b3b4ebf939dafcac3/pyclipper-1.4.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:222ac96c8b8281b53d695b9c4fedc674f56d6d4320ad23f1bdbd168f4e316140", size = 265676, upload-time = "2025-12-01T13:15:04.15Z" }, + { url = "https://files.pythonhosted.org/packages/6b/dd/8bd622521c05d04963420ae6664093f154343ed044c53ea260a310c8bb4d/pyclipper-1.4.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f3672dbafbb458f1b96e1ee3e610d174acb5ace5bd2ed5d1252603bb797f2fc6", size = 140458, upload-time = "2025-12-01T13:15:05.76Z" }, + { url = "https://files.pythonhosted.org/packages/7a/06/6e3e241882bf7d6ab23d9c69ba4e85f1ec47397cbbeee948a16cf75e21ed/pyclipper-1.4.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d1f807e2b4760a8e5c6d6b4e8c1d71ef52b7fe1946ff088f4fa41e16a881a5ca", size = 978235, upload-time = "2025-12-01T13:15:06.993Z" }, + { url = "https://files.pythonhosted.org/packages/cf/f4/3418c1cd5eea640a9fa2501d4bc0b3655fa8d40145d1a4f484b987990a75/pyclipper-1.4.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce1f83c9a4e10ea3de1959f0ae79e9a5bd41346dff648fee6228ba9eaf8b3872", size = 961388, upload-time = "2025-12-01T13:15:08.467Z" }, + { url = "https://files.pythonhosted.org/packages/ac/94/c85401d24be634af529c962dd5d781f3cb62a67cd769534df2cb3feee97a/pyclipper-1.4.0-cp312-cp312-win32.whl", hash = "sha256:3ef44b64666ebf1cb521a08a60c3e639d21b8c50bfbe846ba7c52a0415e936f4", size = 95169, upload-time = "2025-12-01T13:15:10.098Z" }, + { url = "https://files.pythonhosted.org/packages/97/77/dfea08e3b230b82ee22543c30c35d33d42f846a77f96caf7c504dd54fab1/pyclipper-1.4.0-cp312-cp312-win_amd64.whl", hash = "sha256:d1e5498d883b706a4ce636247f0d830c6eb34a25b843a1b78e2c969754ca9037", size = 104619, upload-time = "2025-12-01T13:15:11.592Z" }, + { url = "https://files.pythonhosted.org/packages/67/d0/cbce7d47de1e6458f66a4d999b091640134deb8f2c7351eab993b70d2e10/pyclipper-1.4.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:d49df13cbb2627ccb13a1046f3ea6ebf7177b5504ec61bdef87d6a704046fd6e", size = 264342, upload-time = "2025-12-01T13:15:12.697Z" }, + { url = "https://files.pythonhosted.org/packages/ce/cc/742b9d69d96c58ac156947e1b56d0f81cbacbccf869e2ac7229f2f86dc4e/pyclipper-1.4.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:37bfec361e174110cdddffd5ecd070a8064015c99383d95eb692c253951eee8a", size = 139839, upload-time = "2025-12-01T13:15:13.911Z" }, + { url = "https://files.pythonhosted.org/packages/db/48/dd301d62c1529efdd721b47b9e5fb52120fcdac5f4d3405cfc0d2f391414/pyclipper-1.4.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:14c8bdb5a72004b721c4e6f448d2c2262d74a7f0c9e3076aeff41e564a92389f", size = 972142, upload-time = "2025-12-01T13:15:15.477Z" }, + { url = "https://files.pythonhosted.org/packages/07/bf/d493fd1b33bb090fa64e28c1009374d5d72fa705f9331cd56517c35e381e/pyclipper-1.4.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f2a50c22c3a78cb4e48347ecf06930f61ce98cf9252f2e292aa025471e9d75b1", size = 952789, upload-time = "2025-12-01T13:15:17.042Z" }, + { url = "https://files.pythonhosted.org/packages/cf/88/b95ea8ea21ddca34aa14b123226a81526dd2faaa993f9aabd3ed21231604/pyclipper-1.4.0-cp313-cp313-win32.whl", hash = "sha256:c9a3faa416ff536cee93417a72bfb690d9dea136dc39a39dbbe1e5dadf108c9c", size = 94817, upload-time = "2025-12-01T13:15:18.724Z" }, + { url = "https://files.pythonhosted.org/packages/ba/42/0a1920d276a0e1ca21dc0d13ee9e3ba10a9a8aa3abac76cd5e5a9f503306/pyclipper-1.4.0-cp313-cp313-win_amd64.whl", hash = "sha256:d4b2d7c41086f1927d14947c563dfc7beed2f6c0d9af13c42fe3dcdc20d35832", size = 104007, upload-time = "2025-12-01T13:15:19.763Z" }, + { url = "https://files.pythonhosted.org/packages/1a/20/04d58c70f3ccd404f179f8dd81d16722a05a3bf1ab61445ee64e8218c1f8/pyclipper-1.4.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:7c87480fc91a5af4c1ba310bdb7de2f089a3eeef5fe351a3cedc37da1fcced1c", size = 265167, upload-time = "2025-12-01T13:15:20.844Z" }, + { url = "https://files.pythonhosted.org/packages/bd/2e/a570c1abe69b7260ca0caab4236ce6ea3661193ebf8d1bd7f78ccce537a5/pyclipper-1.4.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:81d8bb2d1fb9d66dc7ea4373b176bb4b02443a7e328b3b603a73faec088b952e", size = 139966, upload-time = "2025-12-01T13:15:22.036Z" }, + { url = "https://files.pythonhosted.org/packages/e8/3b/e0859e54adabdde8a24a29d3f525ebb31c71ddf2e8d93edce83a3c212ffc/pyclipper-1.4.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:773c0e06b683214dcfc6711be230c83b03cddebe8a57eae053d4603dd63582f9", size = 968216, upload-time = "2025-12-01T13:15:23.18Z" }, + { url = "https://files.pythonhosted.org/packages/f6/6b/e3c4febf0a35ae643ee579b09988dd931602b5bf311020535fd9e5b7e715/pyclipper-1.4.0-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9bc45f2463d997848450dbed91c950ca37c6cf27f84a49a5cad4affc0b469e39", size = 954198, upload-time = "2025-12-01T13:15:24.522Z" }, + { url = "https://files.pythonhosted.org/packages/fc/74/728efcee02e12acb486ce9d56fa037120c9bf5b77c54bbdbaa441c14a9d9/pyclipper-1.4.0-cp314-cp314-win32.whl", hash = "sha256:0b8c2105b3b3c44dbe1a266f64309407fe30bf372cf39a94dc8aaa97df00da5b", size = 96951, upload-time = "2025-12-01T13:15:25.79Z" }, + { url = "https://files.pythonhosted.org/packages/e3/d7/7f4354e69f10a917e5c7d5d72a499ef2e10945312f5e72c414a0a08d2ae4/pyclipper-1.4.0-cp314-cp314-win_amd64.whl", hash = "sha256:6c317e182590c88ec0194149995e3d71a979cfef3b246383f4e035f9d4a11826", size = 106782, upload-time = "2025-12-01T13:15:26.945Z" }, + { url = "https://files.pythonhosted.org/packages/63/60/fc32c7a3d7f61a970511ec2857ecd09693d8ac80d560ee7b8e67a6d268c9/pyclipper-1.4.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:f160a2c6ba036f7eaf09f1f10f4fbfa734234af9112fb5187877efed78df9303", size = 269880, upload-time = "2025-12-01T13:15:28.117Z" }, + { url = "https://files.pythonhosted.org/packages/49/df/c4a72d3f62f0ba03ec440c4fff56cd2d674a4334d23c5064cbf41c9583f6/pyclipper-1.4.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:a9f11ad133257c52c40d50de7a0ca3370a0cdd8e3d11eec0604ad3c34ba549e9", size = 141706, upload-time = "2025-12-01T13:15:30.134Z" }, + { url = "https://files.pythonhosted.org/packages/c5/0b/cf55df03e2175e1e2da9db585241401e0bc98f76bee3791bed39d0313449/pyclipper-1.4.0-cp314-cp314t-win32.whl", hash = "sha256:bbc827b77442c99deaeee26e0e7f172355ddb097a5e126aea206d447d3b26286", size = 105308, upload-time = "2025-12-01T13:15:31.225Z" }, + { url = "https://files.pythonhosted.org/packages/8f/dc/53df8b6931d47080b4fe4ee8450d42e660ee1c5c1556c7ab73359182b769/pyclipper-1.4.0-cp314-cp314t-win_amd64.whl", hash = "sha256:29dae3e0296dff8502eeb7639fcfee794b0eec8590ba3563aee28db269da6b04", size = 117608, upload-time = "2025-12-01T13:15:32.69Z" }, + { url = "https://files.pythonhosted.org/packages/18/59/81050abdc9e5b90ffc2c765738c5e40e9abd8e44864aaa737b600f16c562/pyclipper-1.4.0-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:98b2a40f98e1fc1b29e8a6094072e7e0c7dfe901e573bf6cfc6eb7ce84a7ae87", size = 126495, upload-time = "2025-12-01T13:15:33.743Z" }, +] + [[package]] name = "pycparser" version = "3.0" @@ -3536,6 +3594,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/aa/b5/363906b1064fc6fe611783a61764927bbd91919aaaabe8cba82151ca93ef/rapidfuzz-3.14.5-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:dfef96543ced67d9513a422755db422ae1dc34dade0a1485e0b43e7342ed3ebf", size = 1509889, upload-time = "2026-04-07T11:16:28.487Z" }, ] +[[package]] +name = "rapidocr-onnxruntime" +version = "1.4.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "onnxruntime" }, + { name = "opencv-python" }, + { name = "pillow" }, + { name = "pyclipper" }, + { name = "pyyaml" }, + { name = "shapely" }, + { name = "six" }, + { name = "tqdm" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/12/1e5497183bdbe782dbb91bad1d0d2297dba4d2831b2652657f7517bfc6df/rapidocr_onnxruntime-1.4.4-py3-none-any.whl", hash = "sha256:971d7d5f223a7a808662229df1ef69893809d8457d834e6373d3854bc1782cbf", size = 14915192, upload-time = "2025-01-17T01:48:25.104Z" }, +] + [[package]] name = "readerwriterlock" version = "1.0.9" @@ -3923,6 +4000,65 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/07/39/338d9219c4e87f3e708f18857ecd24d22a0c3094752393319553096b98af/scipy-1.17.1-cp314-cp314t-win_arm64.whl", hash = "sha256:200e1050faffacc162be6a486a984a0497866ec54149a01270adc8a59b7c7d21", size = 25489165, upload-time = "2026-02-23T00:22:29.563Z" }, ] +[[package]] +name = "shapely" +version = "2.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4d/bc/0989043118a27cccb4e906a46b7565ce36ca7b57f5a18b78f4f1b0f72d9d/shapely-2.1.2.tar.gz", hash = "sha256:2ed4ecb28320a433db18a5bf029986aa8afcfd740745e78847e330d5d94922a9", size = 315489, upload-time = "2025-09-24T13:51:41.432Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8f/8d/1ff672dea9ec6a7b5d422eb6d095ed886e2e523733329f75fdcb14ee1149/shapely-2.1.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:91121757b0a36c9aac3427a651a7e6567110a4a67c97edf04f8d55d4765f6618", size = 1820038, upload-time = "2025-09-24T13:50:15.628Z" }, + { url = "https://files.pythonhosted.org/packages/4f/ce/28fab8c772ce5db23a0d86bf0adaee0c4c79d5ad1db766055fa3dab442e2/shapely-2.1.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:16a9c722ba774cf50b5d4541242b4cce05aafd44a015290c82ba8a16931ff63d", size = 1626039, upload-time = "2025-09-24T13:50:16.881Z" }, + { url = "https://files.pythonhosted.org/packages/70/8b/868b7e3f4982f5006e9395c1e12343c66a8155c0374fdc07c0e6a1ab547d/shapely-2.1.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cc4f7397459b12c0b196c9efe1f9d7e92463cbba142632b4cc6d8bbbbd3e2b09", size = 3001519, upload-time = "2025-09-24T13:50:18.606Z" }, + { url = "https://files.pythonhosted.org/packages/13/02/58b0b8d9c17c93ab6340edd8b7308c0c5a5b81f94ce65705819b7416dba5/shapely-2.1.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:136ab87b17e733e22f0961504d05e77e7be8c9b5a8184f685b4a91a84efe3c26", size = 3110842, upload-time = "2025-09-24T13:50:21.77Z" }, + { url = "https://files.pythonhosted.org/packages/af/61/8e389c97994d5f331dcffb25e2fa761aeedfb52b3ad9bcdd7b8671f4810a/shapely-2.1.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:16c5d0fc45d3aa0a69074979f4f1928ca2734fb2e0dde8af9611e134e46774e7", size = 4021316, upload-time = "2025-09-24T13:50:23.626Z" }, + { url = "https://files.pythonhosted.org/packages/d3/d4/9b2a9fe6039f9e42ccf2cb3e84f219fd8364b0c3b8e7bbc857b5fbe9c14c/shapely-2.1.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:6ddc759f72b5b2b0f54a7e7cde44acef680a55019eb52ac63a7af2cf17cb9cd2", size = 4178586, upload-time = "2025-09-24T13:50:25.443Z" }, + { url = "https://files.pythonhosted.org/packages/16/f6/9840f6963ed4decf76b08fd6d7fed14f8779fb7a62cb45c5617fa8ac6eab/shapely-2.1.2-cp311-cp311-win32.whl", hash = "sha256:2fa78b49485391224755a856ed3b3bd91c8455f6121fee0db0e71cefb07d0ef6", size = 1543961, upload-time = "2025-09-24T13:50:26.968Z" }, + { url = "https://files.pythonhosted.org/packages/38/1e/3f8ea46353c2a33c1669eb7327f9665103aa3a8dfe7f2e4ef714c210b2c2/shapely-2.1.2-cp311-cp311-win_amd64.whl", hash = "sha256:c64d5c97b2f47e3cd9b712eaced3b061f2b71234b3fc263e0fcf7d889c6559dc", size = 1722856, upload-time = "2025-09-24T13:50:28.497Z" }, + { url = "https://files.pythonhosted.org/packages/24/c0/f3b6453cf2dfa99adc0ba6675f9aaff9e526d2224cbd7ff9c1a879238693/shapely-2.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:fe2533caae6a91a543dec62e8360fe86ffcdc42a7c55f9dfd0128a977a896b94", size = 1833550, upload-time = "2025-09-24T13:50:30.019Z" }, + { url = "https://files.pythonhosted.org/packages/86/07/59dee0bc4b913b7ab59ab1086225baca5b8f19865e6101db9ebb7243e132/shapely-2.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ba4d1333cc0bc94381d6d4308d2e4e008e0bd128bdcff5573199742ee3634359", size = 1643556, upload-time = "2025-09-24T13:50:32.291Z" }, + { url = "https://files.pythonhosted.org/packages/26/29/a5397e75b435b9895cd53e165083faed5d12fd9626eadec15a83a2411f0f/shapely-2.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0bd308103340030feef6c111d3eb98d50dc13feea33affc8a6f9fa549e9458a3", size = 2988308, upload-time = "2025-09-24T13:50:33.862Z" }, + { url = "https://files.pythonhosted.org/packages/b9/37/e781683abac55dde9771e086b790e554811a71ed0b2b8a1e789b7430dd44/shapely-2.1.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1e7d4d7ad262a48bb44277ca12c7c78cb1b0f56b32c10734ec9a1d30c0b0c54b", size = 3099844, upload-time = "2025-09-24T13:50:35.459Z" }, + { url = "https://files.pythonhosted.org/packages/d8/f3/9876b64d4a5a321b9dc482c92bb6f061f2fa42131cba643c699f39317cb9/shapely-2.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e9eddfe513096a71896441a7c37db72da0687b34752c4e193577a145c71736fc", size = 3988842, upload-time = "2025-09-24T13:50:37.478Z" }, + { url = "https://files.pythonhosted.org/packages/d1/a0/704c7292f7014c7e74ec84eddb7b109e1fbae74a16deae9c1504b1d15565/shapely-2.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:980c777c612514c0cf99bc8a9de6d286f5e186dcaf9091252fcd444e5638193d", size = 4152714, upload-time = "2025-09-24T13:50:39.9Z" }, + { url = "https://files.pythonhosted.org/packages/53/46/319c9dc788884ad0785242543cdffac0e6530e4d0deb6c4862bc4143dcf3/shapely-2.1.2-cp312-cp312-win32.whl", hash = "sha256:9111274b88e4d7b54a95218e243282709b330ef52b7b86bc6aaf4f805306f454", size = 1542745, upload-time = "2025-09-24T13:50:41.414Z" }, + { url = "https://files.pythonhosted.org/packages/ec/bf/cb6c1c505cb31e818e900b9312d514f381fbfa5c4363edfce0fcc4f8c1a4/shapely-2.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:743044b4cfb34f9a67205cee9279feaf60ba7d02e69febc2afc609047cb49179", size = 1722861, upload-time = "2025-09-24T13:50:43.35Z" }, + { url = "https://files.pythonhosted.org/packages/c3/90/98ef257c23c46425dc4d1d31005ad7c8d649fe423a38b917db02c30f1f5a/shapely-2.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:b510dda1a3672d6879beb319bc7c5fd302c6c354584690973c838f46ec3e0fa8", size = 1832644, upload-time = "2025-09-24T13:50:44.886Z" }, + { url = "https://files.pythonhosted.org/packages/6d/ab/0bee5a830d209adcd3a01f2d4b70e587cdd9fd7380d5198c064091005af8/shapely-2.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8cff473e81017594d20ec55d86b54bc635544897e13a7cfc12e36909c5309a2a", size = 1642887, upload-time = "2025-09-24T13:50:46.735Z" }, + { url = "https://files.pythonhosted.org/packages/2d/5e/7d7f54ba960c13302584c73704d8c4d15404a51024631adb60b126a4ae88/shapely-2.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fe7b77dc63d707c09726b7908f575fc04ff1d1ad0f3fb92aec212396bc6cfe5e", size = 2970931, upload-time = "2025-09-24T13:50:48.374Z" }, + { url = "https://files.pythonhosted.org/packages/f2/a2/83fc37e2a58090e3d2ff79175a95493c664bcd0b653dd75cb9134645a4e5/shapely-2.1.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7ed1a5bbfb386ee8332713bf7508bc24e32d24b74fc9a7b9f8529a55db9f4ee6", size = 3082855, upload-time = "2025-09-24T13:50:50.037Z" }, + { url = "https://files.pythonhosted.org/packages/44/2b/578faf235a5b09f16b5f02833c53822294d7f21b242f8e2d0cf03fb64321/shapely-2.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a84e0582858d841d54355246ddfcbd1fce3179f185da7470f41ce39d001ee1af", size = 3979960, upload-time = "2025-09-24T13:50:51.74Z" }, + { url = "https://files.pythonhosted.org/packages/4d/04/167f096386120f692cc4ca02f75a17b961858997a95e67a3cb6a7bbd6b53/shapely-2.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dc3487447a43d42adcdf52d7ac73804f2312cbfa5d433a7d2c506dcab0033dfd", size = 4142851, upload-time = "2025-09-24T13:50:53.49Z" }, + { url = "https://files.pythonhosted.org/packages/48/74/fb402c5a6235d1c65a97348b48cdedb75fb19eca2b1d66d04969fc1c6091/shapely-2.1.2-cp313-cp313-win32.whl", hash = "sha256:9c3a3c648aedc9f99c09263b39f2d8252f199cb3ac154fadc173283d7d111350", size = 1541890, upload-time = "2025-09-24T13:50:55.337Z" }, + { url = "https://files.pythonhosted.org/packages/41/47/3647fe7ad990af60ad98b889657a976042c9988c2807cf322a9d6685f462/shapely-2.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:ca2591bff6645c216695bdf1614fca9c82ea1144d4a7591a466fef64f28f0715", size = 1722151, upload-time = "2025-09-24T13:50:57.153Z" }, + { url = "https://files.pythonhosted.org/packages/3c/49/63953754faa51ffe7d8189bfbe9ca34def29f8c0e34c67cbe2a2795f269d/shapely-2.1.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2d93d23bdd2ed9dc157b46bc2f19b7da143ca8714464249bef6771c679d5ff40", size = 1834130, upload-time = "2025-09-24T13:50:58.49Z" }, + { url = "https://files.pythonhosted.org/packages/7f/ee/dce001c1984052970ff60eb4727164892fb2d08052c575042a47f5a9e88f/shapely-2.1.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:01d0d304b25634d60bd7cf291828119ab55a3bab87dc4af1e44b07fb225f188b", size = 1642802, upload-time = "2025-09-24T13:50:59.871Z" }, + { url = "https://files.pythonhosted.org/packages/da/e7/fc4e9a19929522877fa602f705706b96e78376afb7fad09cad5b9af1553c/shapely-2.1.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8d8382dd120d64b03698b7298b89611a6ea6f55ada9d39942838b79c9bc89801", size = 3018460, upload-time = "2025-09-24T13:51:02.08Z" }, + { url = "https://files.pythonhosted.org/packages/a1/18/7519a25db21847b525696883ddc8e6a0ecaa36159ea88e0fef11466384d0/shapely-2.1.2-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:19efa3611eef966e776183e338b2d7ea43569ae99ab34f8d17c2c054d3205cc0", size = 3095223, upload-time = "2025-09-24T13:51:04.472Z" }, + { url = "https://files.pythonhosted.org/packages/48/de/b59a620b1f3a129c3fecc2737104a0a7e04e79335bd3b0a1f1609744cf17/shapely-2.1.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:346ec0c1a0fcd32f57f00e4134d1200e14bf3f5ae12af87ba83ca275c502498c", size = 4030760, upload-time = "2025-09-24T13:51:06.455Z" }, + { url = "https://files.pythonhosted.org/packages/96/b3/c6655ee7232b417562bae192ae0d3ceaadb1cc0ffc2088a2ddf415456cc2/shapely-2.1.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6305993a35989391bd3476ee538a5c9a845861462327efe00dd11a5c8c709a99", size = 4170078, upload-time = "2025-09-24T13:51:08.584Z" }, + { url = "https://files.pythonhosted.org/packages/a0/8e/605c76808d73503c9333af8f6cbe7e1354d2d238bda5f88eea36bfe0f42a/shapely-2.1.2-cp313-cp313t-win32.whl", hash = "sha256:c8876673449f3401f278c86eb33224c5764582f72b653a415d0e6672fde887bf", size = 1559178, upload-time = "2025-09-24T13:51:10.73Z" }, + { url = "https://files.pythonhosted.org/packages/36/f7/d317eb232352a1f1444d11002d477e54514a4a6045536d49d0c59783c0da/shapely-2.1.2-cp313-cp313t-win_amd64.whl", hash = "sha256:4a44bc62a10d84c11a7a3d7c1c4fe857f7477c3506e24c9062da0db0ae0c449c", size = 1739756, upload-time = "2025-09-24T13:51:12.105Z" }, + { url = "https://files.pythonhosted.org/packages/fc/c4/3ce4c2d9b6aabd27d26ec988f08cb877ba9e6e96086eff81bfea93e688c7/shapely-2.1.2-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:9a522f460d28e2bf4e12396240a5fc1518788b2fcd73535166d748399ef0c223", size = 1831290, upload-time = "2025-09-24T13:51:13.56Z" }, + { url = "https://files.pythonhosted.org/packages/17/b9/f6ab8918fc15429f79cb04afa9f9913546212d7fb5e5196132a2af46676b/shapely-2.1.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1ff629e00818033b8d71139565527ced7d776c269a49bd78c9df84e8f852190c", size = 1641463, upload-time = "2025-09-24T13:51:14.972Z" }, + { url = "https://files.pythonhosted.org/packages/a5/57/91d59ae525ca641e7ac5551c04c9503aee6f29b92b392f31790fcb1a4358/shapely-2.1.2-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f67b34271dedc3c653eba4e3d7111aa421d5be9b4c4c7d38d30907f796cb30df", size = 2970145, upload-time = "2025-09-24T13:51:16.961Z" }, + { url = "https://files.pythonhosted.org/packages/8a/cb/4948be52ee1da6927831ab59e10d4c29baa2a714f599f1f0d1bc747f5777/shapely-2.1.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:21952dc00df38a2c28375659b07a3979d22641aeb104751e769c3ee825aadecf", size = 3073806, upload-time = "2025-09-24T13:51:18.712Z" }, + { url = "https://files.pythonhosted.org/packages/03/83/f768a54af775eb41ef2e7bec8a0a0dbe7d2431c3e78c0a8bdba7ab17e446/shapely-2.1.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:1f2f33f486777456586948e333a56ae21f35ae273be99255a191f5c1fa302eb4", size = 3980803, upload-time = "2025-09-24T13:51:20.37Z" }, + { url = "https://files.pythonhosted.org/packages/9f/cb/559c7c195807c91c79d38a1f6901384a2878a76fbdf3f1048893a9b7534d/shapely-2.1.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:cf831a13e0d5a7eb519e96f58ec26e049b1fad411fc6fc23b162a7ce04d9cffc", size = 4133301, upload-time = "2025-09-24T13:51:21.887Z" }, + { url = "https://files.pythonhosted.org/packages/80/cd/60d5ae203241c53ef3abd2ef27c6800e21afd6c94e39db5315ea0cbafb4a/shapely-2.1.2-cp314-cp314-win32.whl", hash = "sha256:61edcd8d0d17dd99075d320a1dd39c0cb9616f7572f10ef91b4b5b00c4aeb566", size = 1583247, upload-time = "2025-09-24T13:51:23.401Z" }, + { url = "https://files.pythonhosted.org/packages/74/d4/135684f342e909330e50d31d441ace06bf83c7dc0777e11043f99167b123/shapely-2.1.2-cp314-cp314-win_amd64.whl", hash = "sha256:a444e7afccdb0999e203b976adb37ea633725333e5b119ad40b1ca291ecf311c", size = 1773019, upload-time = "2025-09-24T13:51:24.873Z" }, + { url = "https://files.pythonhosted.org/packages/a3/05/a44f3f9f695fa3ada22786dc9da33c933da1cbc4bfe876fe3a100bafe263/shapely-2.1.2-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:5ebe3f84c6112ad3d4632b1fd2290665aa75d4cef5f6c5d77c4c95b324527c6a", size = 1834137, upload-time = "2025-09-24T13:51:26.665Z" }, + { url = "https://files.pythonhosted.org/packages/52/7e/4d57db45bf314573427b0a70dfca15d912d108e6023f623947fa69f39b72/shapely-2.1.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5860eb9f00a1d49ebb14e881f5caf6c2cf472c7fd38bd7f253bbd34f934eb076", size = 1642884, upload-time = "2025-09-24T13:51:28.029Z" }, + { url = "https://files.pythonhosted.org/packages/5a/27/4e29c0a55d6d14ad7422bf86995d7ff3f54af0eba59617eb95caf84b9680/shapely-2.1.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b705c99c76695702656327b819c9660768ec33f5ce01fa32b2af62b56ba400a1", size = 3018320, upload-time = "2025-09-24T13:51:29.903Z" }, + { url = "https://files.pythonhosted.org/packages/9f/bb/992e6a3c463f4d29d4cd6ab8963b75b1b1040199edbd72beada4af46bde5/shapely-2.1.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a1fd0ea855b2cf7c9cddaf25543e914dd75af9de08785f20ca3085f2c9ca60b0", size = 3094931, upload-time = "2025-09-24T13:51:32.699Z" }, + { url = "https://files.pythonhosted.org/packages/9c/16/82e65e21070e473f0ed6451224ed9fa0be85033d17e0c6e7213a12f59d12/shapely-2.1.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:df90e2db118c3671a0754f38e36802db75fe0920d211a27481daf50a711fdf26", size = 4030406, upload-time = "2025-09-24T13:51:34.189Z" }, + { url = "https://files.pythonhosted.org/packages/7c/75/c24ed871c576d7e2b64b04b1fe3d075157f6eb54e59670d3f5ffb36e25c7/shapely-2.1.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:361b6d45030b4ac64ddd0a26046906c8202eb60d0f9f53085f5179f1d23021a0", size = 4169511, upload-time = "2025-09-24T13:51:36.297Z" }, + { url = "https://files.pythonhosted.org/packages/b1/f7/b3d1d6d18ebf55236eec1c681ce5e665742aab3c0b7b232720a7d43df7b6/shapely-2.1.2-cp314-cp314t-win32.whl", hash = "sha256:b54df60f1fbdecc8ebc2c5b11870461a6417b3d617f555e5033f1505d36e5735", size = 1602607, upload-time = "2025-09-24T13:51:37.757Z" }, + { url = "https://files.pythonhosted.org/packages/9a/f6/f09272a71976dfc138129b8faf435d064a811ae2f708cb147dccdf7aacdb/shapely-2.1.2-cp314-cp314t-win_amd64.whl", hash = "sha256:0036ac886e0923417932c2e6369b6c52e38e0ff5d9120b90eef5cd9a5fc5cae9", size = 1796682, upload-time = "2025-09-24T13:51:39.233Z" }, +] + [[package]] name = "six" version = "1.17.0" From b3a8bcf8c5ae1af60583f11cfef2a0f9461e0364 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 13 Aug 2026 16:56:35 +0800 Subject: [PATCH 3/9] refactor: remove TOC-related methods and page full text cache - Removed `run_toc` and `_run_toc` methods from `ProfileCoordinator` to streamline TOC processing. - Eliminated `page_full_text_cache` from `PageAnatomyMap` and related logic to simplify data handling. - Updated tests to reflect changes in TOC extraction and anatomy map structure, ensuring consistency in functionality. --- .../services/document_agent/coordinator.py | 35 ---------- .../app/services/document_agent/manifest.py | 4 -- .../document_agent/planner/planner.py | 22 ++---- .../document_agent/planner/prompts.py | 4 +- .../tools/persist_anatomy_map.py | 1 - .../services/page_memory/node_assembler.py | 60 ++-------------- .../test_doc_profile_anatomy_contract.py | 17 ++--- ...est_page_memory_node_assembler_contract.py | 69 +------------------ 8 files changed, 18 insertions(+), 194 deletions(-) diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index b1be64991..c0ce49daa 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -113,23 +113,6 @@ def run_structural(self, *, skip_shard_plan: bool = False) -> PageAnatomyMap: self._record_failure(exc) raise - def run_toc(self) -> TocResult: - try: - return self._run_toc() - except Exception as exc: - logger.warning( - "[document_agent] TOC profiling failed, degrading to empty TOC: {}", - exc, - ) - self.blackboard.toc_result = TocResult( - method="none", - notes=f"degraded: {type(exc).__name__}: {exc}", - failure_kind="degraded", - ) - self.blackboard.toc_hierarchies = None - self._clear_toc_anchor_state() - return self.blackboard.toc_result - def run_lightweight_anatomy( self, *, skip_shard_plan: bool = False ) -> PageAnatomyMap: @@ -185,24 +168,6 @@ def _run_structural(self, *, skip_shard_plan: bool = False) -> PageAnatomyMap: self._persist_ready_anatomy(anatomy) return anatomy - def _run_toc(self) -> TocResult: - self.state = DocumentAgentState.RUNNING - if not self.blackboard.page_features: - self._run_bootstrap() - if not self._toc_profile_enabled(): - self._ensure_disabled_toc_placeholder() - toc_result = self.blackboard.toc_result - if toc_result is None: - raise RuntimeError("TOC placeholder was not initialized") - return toc_result - self._ensure_toc_profile(strict=False) - if self.blackboard.toc_result is None: - self.blackboard.toc_result = TocResult( - method="none", - notes="TOC extraction completed without a result", - ) - return self.blackboard.toc_result - def _run_lightweight_anatomy( self, *, skip_shard_plan: bool = False ) -> PageAnatomyMap: diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index fda87f8b0..b416badee 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -240,7 +240,6 @@ class PageAnatomyMap: skeleton_anchor: dict[str, Any] | None = None skeleton_nodes: list[dict[str, Any]] | None = None pending_skeleton_anchors: list[dict[str, Any]] = field(default_factory=list) - page_full_text_cache: dict[int, str] = field(default_factory=dict) global_signals: dict[str, Any] = field(default_factory=dict) trace_summary: dict[str, Any] = field(default_factory=dict) created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) @@ -265,9 +264,6 @@ def to_dict(self) -> dict[str, Any]: "skeleton_anchor": self.skeleton_anchor, "skeleton_nodes": self.skeleton_nodes, "pending_skeleton_anchors": list(self.pending_skeleton_anchors), - "page_full_text_cache": { - str(page): text for page, text in self.page_full_text_cache.items() - }, "global_signals": dict(self.global_signals), "trace_summary": dict(self.trace_summary), "created_at": self.created_at.isoformat(), diff --git a/apps/worker/app/services/document_agent/planner/planner.py b/apps/worker/app/services/document_agent/planner/planner.py index daeff47fa..6cb266fda 100644 --- a/apps/worker/app/services/document_agent/planner/planner.py +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -73,7 +73,6 @@ def _segment_sample(candidates: list[int], count: int) -> list[int]: def _sample_pages( page_count: int, extrema_pages: list[int], - exclude_pages: set[int] | None = None, *, random_extra: int = 0, rng: random.Random | None = None, @@ -93,18 +92,14 @@ def _sample_pages( page_count: Total number of pages. extrema_pages: Pages with text extrema (min/max raw_text_length / text_density). Low-text extrema already surface chart/asset pages. - exclude_pages: Pages to skip entirely (e.g. TOC pages already detected - by the TOC pipeline). These inflate text-density metrics without - adding profiling value. random_extra: Extra pages to draw uniformly from pages not already - selected (and not excluded). + selected. rng: Optional RNG for deterministic tests. """ if page_count <= 0: return [] - skip = exclude_pages or set() - extrema = [page for page in extrema_pages if 1 <= page <= page_count and page not in skip] - pool = [page for page in range(1, page_count + 1) if page not in set(extrema) and page not in skip] + extrema = [page for page in extrema_pages if 1 <= page <= page_count] + pool = [page for page in range(1, page_count + 1) if page not in set(extrema)] if not pool: return sorted(set(extrema))[:_COARSE_SAMPLE_CAP] third = max(len(pool) // 3, 1) @@ -127,7 +122,7 @@ def _sample_pages( leftover = [ page for page in range(1, page_count + 1) - if page not in ordered and page not in skip + if page not in ordered ] if leftover: picker = rng if rng is not None else random.Random() @@ -220,11 +215,6 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: or self.ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") ) - toc_pages = set( - self.ctx.blackboard.toc_result.toc_pages - if self.ctx.blackboard.toc_result - else [] - ) text_max = float( ( ((self.ctx.blackboard.doc_stats or {}).get("raw_text_length") or {}).get( @@ -245,7 +235,6 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: pages = _sample_pages( self.ctx.blackboard.page_count, extrema_pages, - exclude_pages=toc_pages, random_extra=random_extra, ) if not model: @@ -291,9 +280,6 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: [], ), "sampled_page_features": feature_summary, - "toc_pages": self.ctx.blackboard.toc_result.toc_pages - if self.ctx.blackboard.toc_result - else [], "h1_pages": [ {"title": item.title, "page": item.page} for item in ( diff --git a/apps/worker/app/services/document_agent/planner/prompts.py b/apps/worker/app/services/document_agent/planner/prompts.py index ae4a8c334..0cd83f03c 100644 --- a/apps/worker/app/services/document_agent/planner/prompts.py +++ b/apps/worker/app/services/document_agent/planner/prompts.py @@ -1,8 +1,8 @@ """Prompts for the document profile planner.""" PLANNER_INSTRUCTIONS = ( - "You are a document profile agent. Use page-feature statistics, optional " - "TOC/H1 evidence, and the provided page screenshots to classify the PDF. " + "You are a document profile agent. Use page-feature statistics " + "and the provided page screenshots to classify the PDF. " "Return strict JSON only with keys: is_scanned, category, routing_category, " "category_rationale, language, rationale, header_y, footer_y, next_action, " "grep_query. " diff --git a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py index 7a921d17a..b7683f76b 100644 --- a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py @@ -43,7 +43,6 @@ def build_anatomy_map(ctx: ToolContext) -> PageAnatomyMap: skeleton_anchor=ctx.blackboard.skeleton_anchor, skeleton_nodes=ctx.blackboard.skeleton_nodes, pending_skeleton_anchors=list(ctx.blackboard.pending_skeleton_anchors), - page_full_text_cache=dict(ctx.blackboard.page_full_text_cache), global_signals=ctx.blackboard.global_signals, trace_summary={ "budget": ctx.budget.snapshot(), diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index 7f23d4998..a68d55c23 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -38,7 +38,7 @@ ) from app.services.page_memory.page_tagger import PageTagResult from app.services.page_memory.skeleton_extractor import SectionSkeleton -from shared.services.ai.summary.engine import summarize, transcribe +from shared.services.ai.summary.engine import summarize SAME_AS_PREFIX = "SAME-AS" @@ -243,32 +243,6 @@ def pages_by_leaf_count(views: list[NodePageView]) -> dict[int, list[LeafNode]]: # ── VLM-backed helpers ─────────────────────────────────────────────── -def resolve_page_text( - *, - page: int, - raw_text: str, - image_path: str | None, - vlm_model: str | None, - budget: Any | None = None, -) -> str: - """Body text for an owned page: PyMuPDF text, or VLM OCR for scanned pages. - - Electronic PDFs already have PyMuPDF text; scanned pages have (near) empty - text and fall back to the shared ``transcribe()`` OCR primitive (§4.2). - """ - text = (raw_text or "").strip() - if text: - return text - if not vlm_model or not image_path or not os.path.exists(image_path): - return "" - return transcribe( - image_paths=[image_path], - model=vlm_model, - max_tokens=1500, - usage_task="page_memory.node_ocr", - ) - - def compute_node_summary( *, view: NodePageView, @@ -426,34 +400,10 @@ def build_node_rows( page_to_leaves = pages_by_leaf_count(views) resolved_concurrency = max(1, node_assembly_concurrency) - # Resolve body text once per owned page (PyMuPDF, OCR fallback for scanned). - resolved_text: dict[int, str] = {} - owned_pages = sorted({page for view in views for page in view.owned_pages}) - if owned_pages: - import gevent - from gevent.pool import Pool as GeventPool - - def _resolve_one(page: int) -> tuple[int, str]: - return page, resolve_page_text( - page=page, - raw_text=raw_text_by_page.get(page, ""), - image_path=image_path_by_page.get(page), - vlm_model=vlm_model, - ) - - with stage_timer( - "page_memory.node_ocr", - page_count=len(owned_pages), - concurrency=resolved_concurrency, - ): - pool = GeventPool(size=min(resolved_concurrency, len(owned_pages))) - greenlets = [pool.spawn(_resolve_one, page) for page in owned_pages] - gevent.joinall(greenlets, raise_error=True) - resolved_pairs = [ - cast(tuple[int, str], greenlet.value) - for greenlet in greenlets - ] - resolved_text = {page: text for page, text in resolved_pairs} + resolved_text = { + page: (raw_text_by_page.get(page, "") or "").strip() + for page in sorted({page for view in views for page in view.owned_pages}) + } summaries: dict[int, tuple[str, list[str], list[dict[str, str]]]] = {} if views: diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index aeb1007f2..600432750 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -103,7 +103,7 @@ def test_toc_anchor_text_scan_matches_full_page_and_cross_line_keywords() -> Non assert split_matches[0]["match_kind"] == "cross_line:tableofcontents" -def test_run_toc_degrades_to_empty_result_on_standard_failure(tmp_path: Path) -> None: +def test_toc_extraction_degrades_to_empty_result_on_failure(tmp_path: Path) -> None: coordinator = ProfileCoordinator( pdf_path=str(tmp_path / "standard.pdf"), job_id="job-toc-fail-soft", @@ -117,8 +117,10 @@ def _fail_toc_extraction() -> None: coordinator._run_toc_extraction_pipeline = _fail_toc_extraction # type: ignore[method-assign] - toc_result = coordinator.run_toc() + coordinator._ensure_toc_profile(strict=False) + toc_result = coordinator.blackboard.toc_result + assert toc_result is not None assert toc_result.method == "none" assert toc_result.toc_pages == [] assert toc_result.failure_kind == "degraded" @@ -143,7 +145,6 @@ def test_run_lightweight_anatomy_builds_single_shard_without_planner_llm( routing_category=PdfRoutingCategory.GENERIC.value, ) coordinator.blackboard.toc_result = TocResult(method="none") - coordinator.blackboard.page_full_text_cache = {1: "hello", 2: "world"} anatomy = coordinator.run_lightweight_anatomy() @@ -157,8 +158,7 @@ def test_run_lightweight_anatomy_builds_single_shard_without_planner_llm( (output_dir / "anatomy_map.json").read_text(encoding="utf-8") ) assert list(anatomy_data)[:2] == ["version", "toc_hierarchies"] - assert anatomy.page_full_text_cache == {1: "hello", 2: "world"} - assert anatomy_data["page_full_text_cache"] == {"1": "hello", "2": "world"} + assert "page_full_text_cache" not in anatomy_data assert "text_lines_preview" not in anatomy_data["page_features"][0] assert "asset_bboxes" not in anatomy_data["page_features"][0] trace_data = json.loads((output_dir / "trace.json").read_text(encoding="utf-8")) @@ -816,9 +816,6 @@ def run_coarse(self) -> DocumentProfile: routing_category=PdfRoutingCategory.GENERIC.value, ) - def run_toc(self) -> TocResult: - raise AssertionError("kill switch should not call TOC profiling") - def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): self.calls.append("run_lightweight_anatomy") self.blackboard.toc_result = TocResult( @@ -991,10 +988,6 @@ def run_coarse(self) -> DocumentProfile: routing_category=PdfRoutingCategory.GENERIC.value, ) - def run_toc(self) -> TocResult: - self.calls.append("run_toc") - raise AssertionError("run_toc should be no-op after coarse already ran TOC") - def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): self.calls.append("run_lightweight_anatomy") return fake_anatomy 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 a5e2edb26..3b98debf6 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 @@ -142,26 +142,16 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: assert leaf_b["extra_metadata"] == {} -def test_build_node_rows_preserves_order_under_ocr_and_summary_concurrency( +def test_build_node_rows_preserves_order_under_summary_concurrency( monkeypatch, ) -> None: import gevent - def _fake_resolve_page_text(**kwargs) -> str: - page = int(kwargs["page"]) - gevent.sleep(0.01 * (4 - page)) - return f"text-{page}" - def _fake_compute_node_summary(**kwargs): view = kwargs["view"] gevent.sleep(0.01 * view.leaf.start_page) return f"summary-{view.leaf.start_page}", [f"k{view.leaf.start_page}"], [] - monkeypatch.setattr( - node_assembler, - "resolve_page_text", - _fake_resolve_page_text, - ) monkeypatch.setattr( node_assembler, "compute_node_summary", @@ -170,7 +160,7 @@ def _fake_compute_node_summary(**kwargs): rows = node_assembler.build_node_rows( skeletons=_ordered_page_skeletons(), - raw_text_by_page={1: "", 2: "", 3: ""}, + raw_text_by_page={1: "text-1", 2: "text-2", 3: "text-3"}, image_path_by_page={}, kind_by_page={}, tag_by_page={}, @@ -194,61 +184,6 @@ def _fake_compute_node_summary(**kwargs): ] -def test_build_node_rows_failed_ocr_greenlet_fails_stage(monkeypatch) -> None: - def _fake_resolve_page_text(**kwargs) -> str: - if int(kwargs["page"]) == 2: - raise RuntimeError("ocr failed") - return "ok" - - monkeypatch.setattr( - node_assembler, - "resolve_page_text", - _fake_resolve_page_text, - ) - - with pytest.raises(RuntimeError): - node_assembler.build_node_rows( - skeletons=_ordered_page_skeletons()[:2], - raw_text_by_page={1: "", 2: ""}, - image_path_by_page={}, - kind_by_page={}, - tag_by_page={}, - filename="demo.pdf", - verdict="page", - budget=None, - vlm_model="fake-vlm", - node_assembly_concurrency=2, - ) - - -def test_build_node_rows_unavailable_propagates_from_ocr(monkeypatch) -> None: - def _fake_resolve_page_text(**kwargs) -> str: - raise UnavailableException( - internal_message="ocr capacity busy", - retry_after=5, - ) - - monkeypatch.setattr( - node_assembler, - "resolve_page_text", - _fake_resolve_page_text, - ) - - with pytest.raises(UnavailableException): - node_assembler.build_node_rows( - skeletons=_ordered_page_skeletons()[:1], - raw_text_by_page={1: ""}, - image_path_by_page={}, - kind_by_page={}, - tag_by_page={}, - filename="demo.pdf", - verdict="page", - budget=None, - vlm_model="fake-vlm", - node_assembly_concurrency=1, - ) - - def test_build_node_rows_unavailable_propagates_from_node_summary( monkeypatch, ) -> None: From a6c4772f63c52a9ec3c54ae38e0d1db4759f3508 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 13 Aug 2026 21:35:48 +0800 Subject: [PATCH 4/9] refactor: streamline shard plan validation and TOC handling - Removed `min_pages` parameter from `validate_shard_plan` and related functions to simplify validation logic. - Updated `ProfileCoordinator` to clarify comments regarding single-shard placeholder usage. - Refactored `propose_shard_plan` and `validate_anatomy_map` to enhance clarity and maintainability. - Adjusted tests to reflect the removal of `min_pages` and ensure consistency with updated validation logic. --- .../services/document_agent/coordinator.py | 5 +- .../app/services/document_agent/manifest.py | 5 - .../tools/propose_shard_plan.py | 351 ++++++++++++----- .../tools/validate_anatomy_map.py | 12 +- .../app/services/document_agent/validators.py | 14 +- .../document_parser/formats/pdf/parser.py | 4 - .../formats/pdf/shard_merger.py | 33 +- .../formats/pdf/shard_splitter.py | 6 +- .../test_doc_profile_anatomy_contract.py | 1 - .../contract/test_parse_task_contract.py | 4 +- ...profile_skeleton_anchor_wiring_contract.py | 35 +- .../test_propose_shard_plan_contract.py | 368 ++++++++++++++++-- 12 files changed, 638 insertions(+), 200 deletions(-) diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index c0ce49daa..8b62b07d9 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -186,9 +186,8 @@ def _run_lightweight_anatomy( if skip_shard_plan: # Page-based track processes pages individually via VLM and never # consumes the shard plan; only build_anatomy_map's invariant needs - # it. Populate a single-shard placeholder to skip the LLM shard - # decision + H2 refinement (kept global for chunk-track oversized - # MinerU sharding). + # it. Populate a single-shard placeholder to skip propose.shard_plan + # (kept for chunk-track oversized MinerU sharding). self._apply_single_shard_placeholder() else: result = REGISTRY.dispatch("propose.shard_plan", self.ctx, {}) diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index b416badee..a8c136024 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -187,10 +187,8 @@ class Shard: page_end: int page_offset: int anchor_type: Literal[ - "h1_boundary", "blank_separator", "forced_max_size", - "toc_chapter_boundary", "toc_leaf_boundary", ] anchor_evidence: str @@ -206,9 +204,6 @@ class ShardPlan: reason: Literal[ "too_large", "not_needed", - "parser_stability", - "hierarchy_isolation", - "llm_boundary_decision", ] shards: list[Shard] = field(default_factory=list) validation: ValidationReport = field( diff --git a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index 6d254450b..9c65bffb5 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -1,9 +1,10 @@ -"""Deterministic long-PDF shard planning from TOC leaf boundaries.""" +"""Deterministic long-PDF shard planning from anchored TOC hierarchy ranges.""" from __future__ import annotations import os import time +from dataclasses import dataclass from typing import Any from app.services.document_agent.manifest import ( @@ -13,51 +14,23 @@ ToolResult, ) from app.services.document_agent.registry import has_doc_stats, has_toc_result, register_tool +from app.services.document_agent.structure.anchoring_primitives import ( + deserialize_skeleton_anchor, + deserialize_title_node, + toc_range_start, +) +from app.services.document_agent.structure.hierarchy_locator import ( + ResolvedHierarchyRange, + resolve_hierarchy_page_ranges, +) +from app.services.document_agent.structure.toc_anchoring import ( + body_pages_excluding_toc, + pending_toc_body_scope, + select_global_toc_hierarchies, +) from app.services.document_agent.validators import single_shard_plan, validate_shard_plan -def derive_leaf_cut_pages( - toc_hierarchies: list[dict[str, Any]] | None, - *, - offset_override: int | None = None, -) -> list[int]: - """Derive physical page numbers of TOC leaf nodes for shard splitting. - - Leaf nodes are entries in toc_with_level whose next sibling has level <= theirs - (i.e. they have no children). Requires a calibrated ``offset_override``; - without it this returns [] and the caller falls back to non-TOC planning. - """ - if not toc_hierarchies or offset_override is None: - return [] - - all_pages: list[int] = [] - for hier in toc_hierarchies: - if hier.get("toc_range_unit") != "page": - continue - toc_range = hier.get("toc_range") - entries = hier.get("toc_with_level") - if not toc_range or not entries: - continue - if isinstance(entries, str): - entries = _parse_toc_with_level_entries(entries) - if not entries: - continue - - offset = offset_override - for i, entry in enumerate(entries): - pn = entry.get("page_number") - if not isinstance(pn, int): - continue - is_leaf = ( - i == len(entries) - 1 - or entries[i + 1].get("level", 1) <= entry.get("level", 1) - ) - if is_leaf: - all_pages.append(pn + offset) - - return sorted(set(all_pages)) - - def split_toc_for_shard( toc_hierarchies: list[dict[str, Any]] | None, shard_page_start: int, @@ -187,20 +160,16 @@ def _safe_int(value: Any) -> int | None: return None -def _thresholds(ctx: ToolContext) -> tuple[int, int, int]: +def _thresholds(ctx: ToolContext) -> tuple[int, int]: threshold = int( ctx.settings.get("shard_threshold") or os.environ.get("PARSE_AGENT_SHARD_THRESHOLD", "200") ) - min_pages = int( - ctx.settings.get("min_pages_per_shard") - or os.environ.get("PARSE_AGENT_MIN_PAGES_PER_SHARD", "20") - ) max_pages = int( ctx.settings.get("max_pages_per_shard") or os.environ.get("PARSE_AGENT_MAX_PAGES_PER_SHARD", "200") ) - return threshold, min_pages, max_pages + return threshold, max_pages def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> list[Shard]: @@ -236,18 +205,209 @@ def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> return shards -def _finest_toc_ranges(leaf_pages: list[int], page_count: int) -> list[tuple[int, int]]: - starts = sorted({page for page in leaf_pages if 1 <= page <= page_count}) - if not starts: +@dataclass(frozen=True) +class _PackUnit: + start: int + end: int + path: tuple[str, ...] + + +def _page_span(start: int, end: int) -> int: + return end - start + 1 + + +def _resolve_hierarchy_forests( + ctx: ToolContext, +) -> list[list[ResolvedHierarchyRange]]: + hierarchies = list(ctx.blackboard.toc_hierarchies or []) + anchor_raw = ctx.blackboard.skeleton_anchor + nodes_raw = ctx.blackboard.skeleton_nodes + if ( + not hierarchies + or not isinstance(anchor_raw, dict) + or not isinstance(nodes_raw, list) + ): return [] - ranges: list[tuple[int, int]] = [] - if starts[0] > 1: - ranges.append((1, starts[0] - 1)) - for index, start in enumerate(starts): - end = starts[index + 1] - 1 if index + 1 < len(starts) else page_count + + filename = os.path.basename(ctx.pdf_path) + _primary, pending, _summary = select_global_toc_hierarchies( + hierarchies=hierarchies, + filename=filename, + ) + nodes = [ + deserialize_title_node(node) + for node in nodes_raw + if isinstance(node, dict) + ] + if not nodes: + return [] + + page_count = ctx.blackboard.page_count + page_texts = dict(ctx.blackboard.page_full_text_cache or {}) + toc_result = ctx.blackboard.toc_result + body_pages = body_pages_excluding_toc( + getattr(toc_result, "toc_pages", None) if toc_result else None, + page_count, + ) + records_by_range: dict[tuple[Any, ...], dict[str, Any]] = {} + for record in ctx.blackboard.pending_skeleton_anchors or []: + toc = record.get("toc") + if isinstance(toc, dict): + records_by_range[tuple(toc.get("toc_range") or [])] = record + + parallel_pending: list[tuple[dict[str, Any], dict[str, Any]]] = [] + for pending_toc in pending: + record = records_by_range.get(tuple(pending_toc.get("toc_range") or [])) + if record is None: + continue + relationship = record.get("relationship") + if relationship in {"unresolvable", "contained"}: + continue + if relationship != "parallel": + raise ValueError( + "pending TOC relationship missing after PROFILE classify" + ) + parallel_pending.append((pending_toc, record)) + + parallel_tocs = [toc for toc, _record in parallel_pending] + pending_starts: list[int] = [] + for toc in parallel_tocs: + start = toc_range_start(toc) + if start is not None: + pending_starts.append(start) + primary_page_count = page_count + primary_body_pages = body_pages + if pending_starts: + primary_page_count = min(pending_starts) - 1 + primary_body_pages = [ + page for page in body_pages if page <= primary_page_count + ] + + forests: list[list[ResolvedHierarchyRange]] = [] + anchor = deserialize_skeleton_anchor(anchor_raw) + primary_ranges = resolve_hierarchy_page_ranges( + nodes, + page_count=primary_page_count, + page_texts=page_texts, + body_pages=primary_body_pages, + match_overrides=anchor.match_overrides, + ) + if primary_ranges: + forests.append(primary_ranges) + + for index, (_pending_toc, record) in enumerate(parallel_pending): + resolve_nodes = [ + deserialize_title_node(node) + for node in (record.get("nodes") or []) + if isinstance(node, dict) + ] + pending_anchor_raw = record.get("skeleton_anchor") + if not isinstance(pending_anchor_raw, dict) or not resolve_nodes: + raise ValueError( + "pending TOC skeleton_anchor/nodes missing after PROFILE classify" + ) + pending_anchor = deserialize_skeleton_anchor(pending_anchor_raw) + toc_scope_end, toc_body_pages = pending_toc_body_scope( + pending_tocs=parallel_tocs, + index=index, + page_count=page_count, + body_pages=body_pages, + ) + pending_ranges = resolve_hierarchy_page_ranges( + resolve_nodes, + page_count=toc_scope_end, + page_texts=page_texts, + body_pages=toc_body_pages, + match_overrides=pending_anchor.match_overrides, + ) + if pending_ranges: + forests.append(pending_ranges) + return forests + + +def _greedy_pack_siblings(units: list[_PackUnit], max_pages: int) -> list[_PackUnit]: + if not units: + return [] + parent = units[0].path[:-1] + packed: list[_PackUnit] = [] + current_start, current_end = units[0].start, units[0].end + for unit in units[1:]: + merged_start = min(current_start, unit.start) + merged_end = max(current_end, unit.end) + if _page_span(merged_start, merged_end) <= max_pages: + current_start, current_end = merged_start, merged_end + else: + packed.append(_PackUnit(current_start, current_end, parent)) + current_start, current_end = unit.start, unit.end + packed.append(_PackUnit(current_start, current_end, parent)) + return packed + + +def _pack_forest( + ranges: list[ResolvedHierarchyRange], + max_pages: int, +) -> list[_PackUnit]: + units = [ + _PackUnit(item.start_page, item.end_page, tuple(item.path_titles)) + for item in ranges + if item.start_page <= item.end_page + ] + while units and any(unit.path for unit in units): + max_depth = max(len(unit.path) for unit in units) + next_units: list[_PackUnit] = [] + index = 0 + while index < len(units): + unit = units[index] + if len(unit.path) < max_depth: + next_units.append(unit) + index += 1 + continue + parent = unit.path[:-1] + group = [unit] + index += 1 + while ( + index < len(units) + and len(units[index].path) == max_depth + and units[index].path[:-1] == parent + ): + group.append(units[index]) + index += 1 + next_units.extend(_greedy_pack_siblings(group, max_pages)) + units = next_units + return units + + +def _pack_forests( + forests: list[list[ResolvedHierarchyRange]], + max_pages: int, +) -> list[_PackUnit]: + packed_units: list[_PackUnit] = [] + for index, forest in enumerate(forests): + packed = _pack_forest(forest, max_pages) + if index == 0 and packed and packed[0].start > 1: + packed = _greedy_pack_siblings( + [_PackUnit(1, packed[0].start - 1, ()), *packed], + max_pages, + ) + packed_units.extend(packed) + return packed_units + + +def _exclusive_pieces( + units: list[_PackUnit], +) -> list[tuple[int, int]]: + ordered = sorted(units, key=lambda unit: (unit.start, unit.end)) + pieces: list[tuple[int, int]] = [] + for index, unit in enumerate(ordered): + start = unit.start + end = unit.end + if index + 1 < len(ordered): + next_start = ordered[index + 1].start + if next_start <= end: + end = next_start - 1 if end >= start: - ranges.append((start, end)) - return ranges + pieces.append((start, end)) + return pieces def _pack_range_by_blanks( @@ -275,38 +435,36 @@ def _pack_range_by_blanks( return cuts -def _deterministic_leaf_plan( +def _hierarchy_plan( *, + units: list[_PackUnit], page_count: int, max_pages: int, - leaf_pages: list[int], blank_pages: list[int], -) -> tuple[list[tuple[int, str, str, float]], str]: - ranges = _finest_toc_ranges(leaf_pages, page_count) +) -> list[tuple[int, str, str, float]]: cuts: list[tuple[int, str, str, float]] = [] - shard_start = 0 - index = 0 - while index < len(ranges): - start, end = ranges[index] - if end - shard_start <= max_pages: - index += 1 + previous = 0 + for start, end in _exclusive_pieces(units): + if end <= previous: continue - prior_end = start - 1 - if prior_end > shard_start: - cuts.append((prior_end, "toc_leaf_boundary", f"toc leaf at page {start}", 0.85)) - shard_start = prior_end + if end - previous > max_pages: + range_cuts = _pack_range_by_blanks( + previous=previous, + end=end, + max_pages=max_pages, + blank_pages=blank_pages, + ) + cuts.extend(range_cuts) + if range_cuts: + previous = range_cuts[-1][0] + if previous < end and end < page_count: + cuts.append((end, "toc_leaf_boundary", f"toc leaf at page {end + 1}", 0.85)) + previous = end continue - range_cuts = _pack_range_by_blanks( - previous=shard_start, - end=end, - max_pages=max_pages, - blank_pages=blank_pages, - ) - cuts.extend(range_cuts) - if range_cuts: - shard_start = range_cuts[-1][0] - index += 1 - return cuts, "too_large" + if end < page_count: + cuts.append((end, "toc_leaf_boundary", f"toc leaf at page {end + 1}", 0.85)) + previous = end + return cuts def _get_blank_pages(ctx: ToolContext) -> list[int]: @@ -316,13 +474,16 @@ def _get_blank_pages(ctx: ToolContext) -> list[int]: @register_tool( name="propose.shard_plan", - description="Split a long PDF at TOC leaf boundaries, then blank pages, then max page size.", + description=( + "Split a long PDF by packing anchored TOC hierarchy ranges, " + "then blank pages, then max page size." + ), preconditions=(has_doc_stats, has_toc_result), ) def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() page_count = ctx.blackboard.page_count - threshold, min_pages, max_pages = _thresholds(ctx) + threshold, max_pages = _thresholds(ctx) if page_count <= threshold: plan = single_shard_plan(page_count) ctx.blackboard.shard_plan = plan @@ -332,19 +493,18 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) - leaf_pages = derive_leaf_cut_pages( - ctx.blackboard.toc_hierarchies, - offset_override=ctx.blackboard.toc_page_offset, - ) + forests = _resolve_hierarchy_forests(ctx) blank_pages = _get_blank_pages(ctx) - if leaf_pages: - cuts, reason = _deterministic_leaf_plan( + packed_units = _pack_forests(forests, max_pages) if forests else [] + if packed_units: + cuts = _hierarchy_plan( + units=packed_units, page_count=page_count, max_pages=max_pages, - leaf_pages=leaf_pages, blank_pages=blank_pages, ) - rationale = "Deterministic plan from TOC leaf boundaries." + reason = "too_large" + rationale = "Deterministic plan from anchored TOC hierarchy ranges." else: cuts = _pack_range_by_blanks( previous=0, @@ -366,7 +526,6 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: validation=validate_shard_plan( ShardPlan(enabled=enabled, reason=reason, shards=shards), # type: ignore[arg-type] page_count=page_count, - min_pages=min_pages, max_pages=max_pages, ), ) @@ -382,7 +541,7 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), input_summary={ "page_count": page_count, - "leaf_page_count": len(leaf_pages), + "forest_count": len(forests), }, output_summary={ "enabled": plan.enabled, diff --git a/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py index 2e75594c4..595dc8122 100644 --- a/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py @@ -15,16 +15,11 @@ from app.services.document_agent.validators import validate_anatomy_map -def _thresholds(ctx: ToolContext) -> tuple[int, int]: - min_pages = int( - ctx.settings.get("min_pages_per_shard") - or os.environ.get("PARSE_AGENT_MIN_PAGES_PER_SHARD", "20") - ) - max_pages = int( +def _max_pages(ctx: ToolContext) -> int: + return int( ctx.settings.get("max_pages_per_shard") or os.environ.get("PARSE_AGENT_MAX_PAGES_PER_SHARD", "200") ) - return min_pages, max_pages @register_tool( @@ -56,8 +51,7 @@ def validate_current_anatomy(ctx: ToolContext, _args: dict[str, Any]) -> ToolRes global_signals=ctx.blackboard.global_signals, trace_summary={}, ) - min_pages, max_pages = _thresholds(ctx) - report = validate_anatomy_map(anatomy, min_pages=min_pages, max_pages=max_pages) + report = validate_anatomy_map(anatomy, max_pages=_max_pages(ctx)) ctx.blackboard.validation_report = report.to_dict() if ctx.blackboard.shard_plan: ctx.blackboard.shard_plan.validation = report diff --git a/apps/worker/app/services/document_agent/validators.py b/apps/worker/app/services/document_agent/validators.py index fe26b4328..7d6210021 100644 --- a/apps/worker/app/services/document_agent/validators.py +++ b/apps/worker/app/services/document_agent/validators.py @@ -14,7 +14,6 @@ def validate_shard_plan( plan: ShardPlan, *, page_count: int, - min_pages: int, max_pages: int, ) -> ValidationReport: errors: list[str] = [] @@ -24,8 +23,7 @@ def validate_shard_plan( return ValidationReport(valid=False, errors=errors, warnings=warnings) sorted_shards = sorted(plan.shards, key=lambda item: item.shard_index) expected_start = 1 - for idx, shard in enumerate(sorted_shards): - is_last = idx == len(sorted_shards) - 1 + for shard in sorted_shards: if shard.page_start != expected_start: errors.append( f"shard {shard.shard_index} starts at {shard.page_start}, expected {expected_start}" @@ -37,14 +35,6 @@ def validate_shard_plan( length = shard.page_end - shard.page_start + 1 if length > max_pages: errors.append(f"shard {shard.shard_index} exceeds max_pages={max_pages}") - if plan.enabled and length < min_pages: - if is_last: - warnings.append( - f"shard {shard.shard_index} (final) shorter than min_pages={min_pages} " - f"({length} pages)" - ) - else: - errors.append(f"shard {shard.shard_index} shorter than min_pages={min_pages}") expected_start = shard.page_end + 1 if expected_start != page_count + 1: errors.append("shard_plan does not cover full document") @@ -72,7 +62,6 @@ def single_shard_plan(page_count: int) -> ShardPlan: def validate_anatomy_map( anatomy: PageAnatomyMap, *, - min_pages: int, max_pages: int, ) -> ValidationReport: errors: list[str] = [] @@ -95,7 +84,6 @@ def validate_anatomy_map( shard_report = validate_shard_plan( anatomy.shard_plan, page_count=page_count, - min_pages=min_pages, max_pages=max_pages, ) errors.extend(shard_report.errors) diff --git a/apps/worker/app/services/document_parser/formats/pdf/parser.py b/apps/worker/app/services/document_parser/formats/pdf/parser.py index fb4b8e031..7bfb39e16 100755 --- a/apps/worker/app/services/document_parser/formats/pdf/parser.py +++ b/apps/worker/app/services/document_parser/formats/pdf/parser.py @@ -321,12 +321,8 @@ def _predict_shard_headings( raise RuntimeError(f"Missing heading result for shard_{index}") complete_heading_results.append(result) - # No level offsets needed — leaf-node splitting produces self-contained shards - shard_offsets: list[int] = [0] * len(complete_heading_results) - all_lines_with_heading: list[str] = merge_shard_lines( [result.lines_with_heading for result in complete_heading_results], - shard_offsets=shard_offsets, ) total_headings = sum( 1 for line in all_lines_with_heading if line.startswith("#") diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py index 748413e5b..78ed71a0a 100644 --- a/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py +++ b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py @@ -17,33 +17,11 @@ def _extract_heading_key(line: str) -> tuple[int, str] | None: return len(m.group(1)), m.group(2).strip() -def _apply_level_offset(lines: list[str], offset: int) -> list[str]: - """Shift all markdown heading levels by *offset* (e.g. ## → ### when offset=1).""" - if offset <= 0: - return lines - result: list[str] = [] - for line in lines: - key = _extract_heading_key(line) - if key is not None: - level, text = key - new_level = level + offset - result.append(f"{'#' * new_level} {text}") - else: - result.append(line) - return result - - def merge_shard_lines( shard_lines_list: list[list[str]], - shard_offsets: list[int] | None = None, ) -> list[str]: """Concatenate per-shard lines_with_heading in order, removing boundary duplicates. - When ``shard_offsets`` is provided, each shard's heading levels are shifted - by the corresponding offset. This is used for continuation shards (from - H2+ splitting) whose heading predictor starts from L1 but should be deeper - in the global hierarchy. - When a PDF section-divider page falls at the end of shard N and the same heading opens shard N+1, each shard independently identifies it as a heading, resulting in two consecutive identical headings after naïve concatenation. @@ -61,25 +39,17 @@ def merge_shard_lines( if not lines: continue - # Apply level offset for continuation shards - offset = shard_offsets[shard_idx] if shard_offsets else 0 - lines = _apply_level_offset(lines, offset) - - # Determine next shard's first heading (if any) next_first_heading: tuple[int, str] | None = None for future_idx in range(shard_idx + 1, len(shard_lines_list)): future_lines = shard_lines_list[future_idx] - future_offset = shard_offsets[future_idx] if shard_offsets else 0 for next_line in future_lines: key = _extract_heading_key(next_line) if key is not None: - # Compare with the offset-adjusted level - next_first_heading = (key[0] + future_offset, key[1]) + next_first_heading = key break if next_first_heading is not None: break - # Find this shard's last heading and its position lines_to_add = list(lines) if next_first_heading is not None: last_heading_pos: int | None = None @@ -95,7 +65,6 @@ def merge_shard_lines( and last_heading_key is not None and last_heading_key == next_first_heading ): - # Truncate from the last (duplicate) heading onward logger.info( f"🔗 shard_{shard_idx}: removing trailing boundary heading " f"'{last_heading_key[1]}' (L{last_heading_key[0]}) duplicated " diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py index 6fe1308b8..accc1dde6 100644 --- a/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py +++ b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py @@ -1,4 +1,4 @@ -"""PDF shard splitting: bin-packing + physical split.""" +"""PDF shard splitting: 1:1 agent-shard mapping + physical split.""" from __future__ import annotations @@ -37,8 +37,8 @@ def bin_pack_shards( ) -> list[MergedShard]: """1:1 mapping: each agent shard becomes its own MinerU shard. - Agent shards are cut at semantic boundaries (H1/H2) by the document - agent. Merging them would cross those boundaries and degrade heading + Agent shards are cut at TOC hierarchy pack boundaries by the document + agent. Merging them would cross those boundaries and degrade heading prediction quality, so we preserve them as-is. """ return [ diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index 600432750..512af8ca6 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -726,7 +726,6 @@ def test_oversized_single_shard_plan_is_invalid() -> None: ], ), page_count=407, - min_pages=20, max_pages=200, ) diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 696c5e766..66267f357 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -829,7 +829,7 @@ def delete_upload_file(self, storage_key: str) -> bool: page_start=1, page_end=2, page_offset=0, - anchor_type="h1_boundary", + anchor_type="toc_leaf_boundary", anchor_evidence="Chapter 1", confidence=0.9, ), @@ -838,7 +838,7 @@ def delete_upload_file(self, storage_key: str) -> bool: page_start=3, page_end=3, page_offset=2, - anchor_type="h1_boundary", + anchor_type="toc_leaf_boundary", anchor_evidence="Chapter 2", confidence=0.9, ), diff --git a/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py b/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py index 4ca5d3756..096494f52 100644 --- a/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py +++ b/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py @@ -208,11 +208,44 @@ def test_shard_plan_reads_offset_and_does_not_calibrate() -> None: } ] ctx.blackboard.toc_page_offset = 0 + ctx.blackboard.skeleton_nodes = [ + serialize_title_node(TitleNode(title="Ch1", level=1, printed_page=3)), + serialize_title_node(TitleNode(title="Ch2", level=1, printed_page=120)), + ] + ctx.blackboard.skeleton_anchor = serialize_skeleton_anchor( + SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides={ + ("Ch1",): TitleMatch( + page=3, + confidence=1.0, + source="anchored", + matched_line="Ch1", + score=1.0, + candidates=[3], + evidence={}, + ), + ("Ch2",): TitleMatch( + page=120, + confidence=1.0, + source="anchored", + matched_line="Ch2", + score=1.0, + candidates=[120], + evidence={}, + ), + }, + null_page_report=[], + bulk_count=2, + pruned_count=0, + locate_agent="offset_guided_bulk", + ) + ) ctx.blackboard.toc_result = TocResult(method="vlm_batch") ctx.blackboard.doc_stats = {"page_count": 250} ctx.settings["shard_threshold"] = 200 ctx.settings["max_pages_per_shard"] = 200 - ctx.settings["min_pages_per_shard"] = 20 def _boom(*_args, **_kwargs): raise AssertionError("shard plan must not recalibrate") diff --git a/apps/worker/tests/contract/test_propose_shard_plan_contract.py b/apps/worker/tests/contract/test_propose_shard_plan_contract.py index 21d1e1e9a..5953f341a 100644 --- a/apps/worker/tests/contract/test_propose_shard_plan_contract.py +++ b/apps/worker/tests/contract/test_propose_shard_plan_contract.py @@ -12,6 +12,12 @@ from app.services.document_agent.budget import BudgetTracker from app.services.document_agent.manifest import PageFeature, TocResult, ToolContext from app.services.document_agent.state import AgentBlackboard +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.tools.propose_shard_plan import propose_shard_plan @@ -43,7 +49,6 @@ def _ctx(*, page_count: int, blank_pages: list[int] | None = None) -> ToolContex settings={ "shard_threshold": 200, "max_pages_per_shard": 200, - "min_pages_per_shard": 20, }, ) ctx.blackboard.doc_stats = {"page_count": page_count} @@ -54,21 +59,80 @@ def _ctx(*, page_count: int, blank_pages: list[int] | None = None) -> ToolContex return ctx -def test_leaf_plan_cuts_at_finest_toc_boundary() -> None: +def _match(title: str, page: int) -> TitleMatch: + return TitleMatch( + page=page, + confidence=1.0, + source="anchored", + matched_line=title, + score=1.0, + candidates=[page], + evidence={}, + ) + + +def _seed_skeleton( + ctx: ToolContext, + *, + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], int], + hierarchies: list[dict[str, object]], + pending_records: list[dict[str, object]] | None = None, +) -> None: + ctx.blackboard.toc_hierarchies = hierarchies + ctx.blackboard.skeleton_nodes = [serialize_title_node(node) for node in nodes] + ctx.blackboard.skeleton_anchor = serialize_skeleton_anchor( + SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides={ + path: _match(path[-1], page) for path, page in overrides.items() + }, + null_page_report=[], + bulk_count=len(overrides), + pruned_count=0, + locate_agent="offset_guided_bulk", + ) + ) + if pending_records is not None: + ctx.blackboard.pending_skeleton_anchors = pending_records + + +def test_hierarchy_pack_cuts_before_next_chapter() -> None: ctx = _ctx(page_count=250) - ctx.blackboard.toc_page_offset = 0 - ctx.blackboard.toc_hierarchies = [ - { - "toc_range": [1, 2], - "toc_range_unit": "page", - "toc_with_level": [ - {"heading": "Ch1", "level": 1, "page_number": 3}, - {"heading": "1.1", "level": 2, "page_number": 3}, - {"heading": "1.2", "level": 2, "page_number": 80}, - {"heading": "Ch2", "level": 1, "page_number": 120}, - ], - } - ] + _seed_skeleton( + ctx, + nodes=[ + TitleNode( + title="Ch1", + level=1, + printed_page=3, + children=[ + TitleNode(title="1.1", level=2, printed_page=3), + TitleNode(title="1.2", level=2, printed_page=80), + ], + ), + TitleNode(title="Ch2", level=1, printed_page=120), + ], + overrides={ + ("Ch1",): 3, + ("Ch1", "1.1"): 3, + ("Ch1", "1.2"): 80, + ("Ch2",): 120, + }, + hierarchies=[ + { + "toc_range": [1, 2], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 3}, + {"heading": "1.1", "level": 2, "page_number": 3}, + {"heading": "1.2", "level": 2, "page_number": 80}, + {"heading": "Ch2", "level": 1, "page_number": 120}, + ], + } + ], + ) result = propose_shard_plan(ctx, {}) plan = ctx.blackboard.shard_plan @@ -83,16 +147,254 @@ def test_leaf_plan_cuts_at_finest_toc_boundary() -> None: assert plan.validation.valid is True +def test_hierarchy_pack_keeps_same_parent_siblings_together() -> None: + ctx = _ctx(page_count=250) + _seed_skeleton( + ctx, + nodes=[ + TitleNode( + title="Ch1", + level=1, + printed_page=1, + children=[ + TitleNode(title="1.1", level=2, printed_page=1), + TitleNode(title="1.2", level=2, printed_page=120), + ], + ), + TitleNode( + title="Ch2", + level=1, + printed_page=150, + children=[ + TitleNode(title="2.1", level=2, printed_page=150), + TitleNode(title="2.2", level=2, printed_page=180), + ], + ), + ], + overrides={ + ("Ch1",): 1, + ("Ch1", "1.1"): 1, + ("Ch1", "1.2"): 120, + ("Ch2",): 150, + ("Ch2", "2.1"): 150, + ("Ch2", "2.2"): 180, + }, + hierarchies=[ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 1}, + {"heading": "1.1", "level": 2, "page_number": 1}, + {"heading": "1.2", "level": 2, "page_number": 120}, + {"heading": "Ch2", "level": 1, "page_number": 150}, + {"heading": "2.1", "level": 2, "page_number": 150}, + {"heading": "2.2", "level": 2, "page_number": 180}, + ], + } + ], + ) + + propose_shard_plan(ctx, {}) + plan = ctx.blackboard.shard_plan + assert plan is not None + assert [(shard.page_start, shard.page_end) for shard in plan.shards] == [ + (1, 149), + (150, 250), + ] + + +def test_match_override_keeps_non_int_printed_page() -> None: + ctx = _ctx(page_count=250) + _seed_skeleton( + ctx, + nodes=[ + TitleNode( + title="Ch1", + level=1, + printed_page=1, + children=[ + TitleNode(title="A", level=2, printed_page=1), + TitleNode( + title="B", + level=2, + printed_page=None, + printed_label="iv", + page_kind="roman", + ), + TitleNode(title="C", level=2, printed_page=201), + ], + ), + ], + overrides={ + ("Ch1",): 1, + ("Ch1", "A"): 1, + ("Ch1", "B"): 101, + ("Ch1", "C"): 201, + }, + hierarchies=[ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 1}, + {"heading": "A", "level": 2, "page_number": 1}, + {"heading": "B", "level": 2, "page_number": "iv"}, + {"heading": "C", "level": 2, "page_number": 201}, + ], + } + ], + ) + + propose_shard_plan(ctx, {}) + plan = ctx.blackboard.shard_plan + assert plan is not None + assert [(shard.page_start, shard.page_end) for shard in plan.shards] == [ + (1, 100), + (101, 250), + ] + + +def test_pending_toc_forest_is_packed_separately() -> None: + ctx = _ctx(page_count=350) + pending_toc = { + "toc_range": [200, 201], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "App", "level": 1, "page_number": 210}, + ], + } + _seed_skeleton( + ctx, + nodes=[ + TitleNode(title="Ch1", level=1, printed_page=3), + TitleNode(title="Ch2", level=1, printed_page=50), + ], + overrides={ + ("Ch1",): 3, + ("Ch2",): 50, + }, + hierarchies=[ + { + "toc_range": [1, 2], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 3}, + {"heading": "Ch2", "level": 1, "page_number": 50}, + ], + }, + pending_toc, + ], + pending_records=[ + { + "toc": pending_toc, + "relationship": "parallel", + "nodes": [ + serialize_title_node( + TitleNode(title="App", level=1, printed_page=210) + ) + ], + "skeleton_anchor": serialize_skeleton_anchor( + SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides={("App",): _match("App", 210)}, + null_page_report=[], + bulk_count=1, + pruned_count=0, + locate_agent="offset_guided_bulk", + ) + ), + } + ], + ) + + propose_shard_plan(ctx, {}) + plan = ctx.blackboard.shard_plan + assert plan is not None + assert [(shard.page_start, shard.page_end) for shard in plan.shards] == [ + (1, 199), + (200, 350), + ] + + +def test_contained_pending_toc_does_not_cut() -> None: + ctx = _ctx(page_count=250) + pending_toc = { + "toc_range": [200, 201], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Inner", "level": 1, "page_number": 210}, + ], + } + _seed_skeleton( + ctx, + nodes=[ + TitleNode(title="Ch1", level=1, printed_page=3), + TitleNode(title="Ch2", level=1, printed_page=120), + ], + overrides={ + ("Ch1",): 3, + ("Ch2",): 120, + }, + hierarchies=[ + { + "toc_range": [1, 2], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 3}, + {"heading": "Ch2", "level": 1, "page_number": 120}, + ], + }, + pending_toc, + ], + pending_records=[ + { + "toc": pending_toc, + "relationship": "contained", + "nodes": [ + serialize_title_node( + TitleNode(title="Inner", level=1, printed_page=210) + ) + ], + "skeleton_anchor": serialize_skeleton_anchor( + SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides={("Inner",): _match("Inner", 210)}, + null_page_report=[], + bulk_count=1, + pruned_count=0, + locate_agent="offset_guided_bulk", + ) + ), + } + ], + ) + + propose_shard_plan(ctx, {}) + plan = ctx.blackboard.shard_plan + assert plan is not None + assert [(shard.page_start, shard.page_end) for shard in plan.shards] == [ + (1, 119), + (120, 250), + ] + + def test_fat_leaf_uses_blank_page_in_window() -> None: ctx = _ctx(page_count=450, blank_pages=[195]) - ctx.blackboard.toc_page_offset = 0 - ctx.blackboard.toc_hierarchies = [ - { - "toc_range": [1, 1], - "toc_range_unit": "page", - "toc_with_level": [{"heading": "Only", "level": 1, "page_number": 1}], - } - ] + _seed_skeleton( + ctx, + nodes=[TitleNode(title="Only", level=1, printed_page=1)], + overrides={("Only",): 1}, + hierarchies=[ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [{"heading": "Only", "level": 1, "page_number": 1}], + } + ], + ) propose_shard_plan(ctx, {}) plan = ctx.blackboard.shard_plan @@ -105,14 +407,18 @@ def test_fat_leaf_uses_blank_page_in_window() -> None: def test_fat_leaf_without_blank_uses_forced_max_size() -> None: ctx = _ctx(page_count=450) - ctx.blackboard.toc_page_offset = 0 - ctx.blackboard.toc_hierarchies = [ - { - "toc_range": [1, 1], - "toc_range_unit": "page", - "toc_with_level": [{"heading": "Only", "level": 1, "page_number": 1}], - } - ] + _seed_skeleton( + ctx, + nodes=[TitleNode(title="Only", level=1, printed_page=1)], + overrides={("Only",): 1}, + hierarchies=[ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [{"heading": "Only", "level": 1, "page_number": 1}], + } + ], + ) propose_shard_plan(ctx, {}) plan = ctx.blackboard.shard_plan From 0d4dd2519ea5131c398499cfbe546892291ac85c Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 14 Aug 2026 23:14:55 +0800 Subject: [PATCH 5/9] feat: enhance TOC processing and link enrichment - Introduced `_attach_toc_page_links` method to attach hyperlinks to TOC entries before calibration. - Updated `run_toc_anchoring` to conditionally skip anchoring based on settings. - Refactored `ProfileCoordinator` to ensure asset probes are executed in the correct order relative to TOC processing. - Enhanced `collect_toc_page_links` to normalize page link destinations and improve link matching logic. - Updated `profile_document` to include a new parameter for skipping TOC anchoring during profiling. - Adjusted tests to validate the new TOC processing flow and ensure correct asset probing order. --- .../services/document_agent/coordinator.py | 52 +- .../document_agent/structure/toc_anchoring.py | 107 +++- .../document_agent/structure/toc_graft.py | 384 ++++++++++++ .../structure/toc_link_enrichment.py | 144 ++--- .../document_parser/profiling/doc_profiler.py | 10 + .../page_memory/skeleton_extractor.py | 99 +-- .../test_doc_profile_anatomy_contract.py | 8 +- .../tests/contract/test_toc_graft_contract.py | 583 ++++++++++++++++++ .../test_toc_link_attach_wiring_contract.py | 190 ++++++ .../contract/test_toc_link_match_contract.py | 104 ++++ 10 files changed, 1509 insertions(+), 172 deletions(-) create mode 100644 apps/worker/app/services/document_agent/structure/toc_graft.py create mode 100644 apps/worker/tests/contract/test_toc_graft_contract.py create mode 100644 apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py create mode 100644 apps/worker/tests/contract/test_toc_link_match_contract.py diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 8b62b07d9..04beae484 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -28,6 +28,9 @@ from app.services.document_agent.registry import REGISTRY from app.services.document_agent.state import AgentBlackboard, DocumentAgentState from app.services.document_agent.structure.toc_anchoring import run_toc_anchoring +from app.services.document_agent.structure.toc_link_enrichment import ( + enrich_toc_hierarchies_with_links, +) from app.services.document_agent import tools as _registered_tools # noqa: F401 from app.services.document_agent.trace import ParseRunRecorder from app.services.document_agent.validators import single_shard_plan @@ -130,17 +133,27 @@ def _run_coarse(self) -> DocumentProfile: actor="planner:coarse" ) self._run_text_scan() + # Asset coarse probe is independent of TOC; run it before TOC so + # PROFILE / debug Stage-0 share the same order as later anatomy/shard + # consumers of ``page_features.has_asset``. + self._ensure_asset_probe() + if self.ctx.settings.get("stop_after_asset_probe"): + # Debug Stage-0: bootstrap → coarse VLM → text scan → asset probe. + return profile if self._toc_profile_enabled(): self._ensure_toc_profile(strict=False) else: self._ensure_disabled_toc_placeholder() - self._ensure_asset_probe() return profile def _run_structural(self, *, skip_shard_plan: bool = False) -> PageAnatomyMap: self.state = DocumentAgentState.RUNNING if not self.blackboard.page_features: self._run_bootstrap() + # Prefer assets before TOC so cold structural matches coarse order. + # After ``run_coarse`` this is a no-op (assets_probed already set with + # coarse header/footer margins). + self._ensure_asset_probe() if self._toc_profile_enabled(): self._ensure_toc_profile(strict=True) else: @@ -148,7 +161,6 @@ def _run_structural(self, *, skip_shard_plan: bool = False) -> PageAnatomyMap: profile, initial_decision, _planner_result = self._propose_profile( actor="planner" ) - self._ensure_asset_probe() if skip_shard_plan: # Page-memory oversized path never consumes shard_plan; only # build_anatomy_map's invariant needs a non-empty plan. @@ -174,6 +186,8 @@ def _run_lightweight_anatomy( self.state = DocumentAgentState.RUNNING if not self.blackboard.page_features: self._run_bootstrap() + # Same relative order as coarse: assets before any TOC placeholder. + # After ``run_coarse`` this is a no-op. self._ensure_asset_probe() if self.blackboard.toc_result is None: if self._toc_profile_enabled(): @@ -421,5 +435,39 @@ def _run_toc_extraction_pipeline(self) -> None: tool_name=tool_name, actor=f"toc:{tool_name}", ) + self._attach_toc_page_links() + if self.ctx.settings.get("skip_toc_anchoring"): + # Debug Stage-1: stop after TOC extract + link attach. + self._clear_toc_anchor_state() + logger.info( + "[document_agent] skip_toc_anchoring=True; " + "leaving calibration to a later stage" + ) + return run_toc_anchoring(self.ctx) + def _attach_toc_page_links(self) -> None: + """Attach TOC-page hyperlinks onto VLM entries before calibration.""" + hierarchies = list(self.blackboard.toc_hierarchies or []) + if not hierarchies: + return + try: + enriched, stats = enrich_toc_hierarchies_with_links( + pdf_path=self.ctx.pdf_path, + toc_hierarchies=hierarchies, + ) + except Exception as exc: + logger.warning( + "[document_agent] TOC link attach failed, " + "continuing without links: {}", + exc, + ) + return + self.blackboard.toc_hierarchies = enriched + logger.info( + "[document_agent] TOC link attach: matched={}/{} skipped_no_links={}", + stats.entries_matched, + stats.entries_total, + stats.skipped_no_links, + ) + 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 2f7fe1b70..ba9995261 100644 --- a/apps/worker/app/services/document_agent/structure/toc_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/toc_anchoring.py @@ -2,19 +2,17 @@ from __future__ import annotations +from dataclasses import replace from pathlib import Path from typing import Any from loguru import logger -from app.services.document_agent.agents.calibration import calibrate_offset -from app.services.document_agent.agents.calibration.orchestrator import anchor_hierarchy -from app.services.document_agent.agents.calibration.procedure import ( - finalize_calibration_result, - pick_primary_offset, -) from app.services.document_agent.manifest import ToolContext from app.services.document_agent.structure.anchoring_primitives import ( + SkeletonAnchor, + deserialize_skeleton_anchor, + deserialize_title_node, serialize_skeleton_anchor, serialize_title_node, toc_range_end, @@ -42,7 +40,9 @@ def run_toc_anchoring(ctx: ToolContext) -> None: page_texts = dict(ctx.blackboard.page_full_text_cache) if not page_texts: - raise ValueError("page_full_text_cache missing; run text scan before TOC anchoring") + raise ValueError( + "page_full_text_cache missing; run text scan before TOC anchoring" + ) filename = Path(ctx.pdf_path).name primary, pending, _summary = select_global_toc_hierarchies( hierarchies=hierarchies, @@ -58,37 +58,22 @@ def run_toc_anchoring(ctx: ToolContext) -> None: getattr(toc_result, "toc_pages", None), page_count, ) - pending_starts: list[int] = [] - for toc in pending: - start = toc_range_start(toc) - if start is not None: - pending_starts.append(start) - primary_page_count = page_count - primary_body_pages = body_pages - if pending_starts: - primary_page_count = min(pending_starts) - 1 - primary_body_pages = [ - page for page in body_pages if page <= primary_page_count - ] resolve_nodes, skeleton_anchor = anchor_hierarchy( nodes=nodes, toc_hierarchies=primary, page_texts=page_texts, - body_pages=primary_body_pages, + body_pages=body_pages, page_count=page_count, ctx=ctx, ) - ctx.blackboard.skeleton_anchor = serialize_skeleton_anchor(skeleton_anchor) - ctx.blackboard.skeleton_nodes = [serialize_title_node(node) for node in resolve_nodes] - ctx.blackboard.toc_page_offset = skeleton_anchor.offset pending_records: list[dict[str, Any]] = [] if pending: primary_ranges = resolve_hierarchy_page_ranges( resolve_nodes, - page_count=primary_page_count, + page_count=page_count, page_texts=page_texts, - body_pages=primary_body_pages, + body_pages=body_pages, match_overrides=skeleton_anchor.match_overrides, ) pending_records = _anchor_pending_tocs( @@ -99,6 +84,19 @@ def run_toc_anchoring(ctx: ToolContext) -> None: body_pages=body_pages, primary_ranges=primary_ranges, ) + resolve_nodes, skeleton_anchor = _graft_contained_pending( + resolve_nodes=resolve_nodes, + skeleton_anchor=skeleton_anchor, + pending_records=pending_records, + page_count=page_count, + page_texts=page_texts, + body_pages=body_pages, + ) + ctx.blackboard.skeleton_anchor = serialize_skeleton_anchor(skeleton_anchor) + ctx.blackboard.skeleton_nodes = [ + serialize_title_node(node) for node in resolve_nodes + ] + ctx.blackboard.toc_page_offset = skeleton_anchor.offset ctx.blackboard.pending_skeleton_anchors = pending_records @@ -213,7 +211,9 @@ def classify_toc_relationship( section's explicitly-anchored range. """ leaves = [ - node for _, node in iter_leaf_title_nodes(nodes) if node.printed_page is not None + node + for _, node in iter_leaf_title_nodes(nodes) + if node.printed_page is not None ] if not leaves: return "unresolvable" @@ -251,6 +251,21 @@ def classify_toc_relationship( return "parallel" +from app.services.document_agent.agents.calibration.orchestrator import ( # noqa: E402 + anchor_hierarchy, +) +from app.services.document_agent.agents.calibration.procedure import ( # noqa: E402 + finalize_calibration_result, + pick_primary_offset, +) +from app.services.document_agent.agents.calibration.service import ( # noqa: E402 + calibrate_offset, +) +from app.services.document_agent.structure.toc_graft import ( # noqa: E402 + graft_contained_toc, +) + + def _anchor_pending_tocs( *, pending_tocs: list[dict[str, Any]], @@ -325,3 +340,43 @@ def _anchor_pending_tocs( } ) return records + + +def _graft_contained_pending( + *, + resolve_nodes: list[TitleNode], + skeleton_anchor: SkeletonAnchor, + pending_records: list[dict[str, Any]], + page_count: int, + page_texts: dict[int, str], + body_pages: list[int], +) -> tuple[list[TitleNode], SkeletonAnchor]: + nodes = resolve_nodes + overrides = dict(skeleton_anchor.match_overrides) + for record in pending_records: + if record.get("relationship") != "contained": + continue + nodes_raw = record.get("nodes") or [] + anchor_raw = record.get("skeleton_anchor") + if not isinstance(anchor_raw, dict) or not nodes_raw: + continue + contained_nodes = [ + deserialize_title_node(node) for node in nodes_raw if isinstance(node, dict) + ] + if not contained_nodes: + continue + contained_anchor = deserialize_skeleton_anchor(anchor_raw) + grafted = graft_contained_toc( + primary_nodes=nodes, + primary_overrides=overrides, + contained_nodes=contained_nodes, + contained_overrides=contained_anchor.match_overrides, + page_count=page_count, + page_texts=page_texts, + body_pages=body_pages, + ) + nodes = grafted.nodes + overrides = grafted.match_overrides + record["grafted"] = True + record["graft"] = grafted.events + return nodes, replace(skeleton_anchor, match_overrides=overrides) diff --git a/apps/worker/app/services/document_agent/structure/toc_graft.py b/apps/worker/app/services/document_agent/structure/toc_graft.py new file mode 100644 index 000000000..efebf6955 --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/toc_graft.py @@ -0,0 +1,384 @@ +"""Graft contained pending TOC trees into the primary hierarchy by physical start page.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any + +from loguru import logger + +from app.services.document_agent.structure.hierarchy_locator import ( + ResolvedHierarchyRange, + TitleMatch, + TitleNode, + resolve_hierarchy_page_ranges, +) + +_LOG_PREFIX = "[profile.toc_graft]" + + +@dataclass(frozen=True) +class ContainedGraftResult: + nodes: list[TitleNode] + match_overrides: dict[tuple[str, ...], TitleMatch] + events: list[dict[str, Any]] + + +def graft_contained_toc( + *, + primary_nodes: list[TitleNode], + primary_overrides: dict[tuple[str, ...], TitleMatch], + contained_nodes: list[TitleNode], + contained_overrides: dict[tuple[str, ...], TitleMatch], + page_count: int, + page_texts: dict[int, str], + body_pages: list[int], +) -> ContainedGraftResult: + """Merge one contained TOC forest into the current primary tree.""" + overrides = dict(primary_overrides) + events: list[dict[str, Any]] = [] + ranges = resolve_hierarchy_page_ranges( + primary_nodes, + page_count=page_count, + page_texts=page_texts, + body_pages=body_pages, + match_overrides=overrides, + ) + coverage = _coverage_by_path(ranges) + nodes = _graft_children( + parent=None, + parent_path=(), + primary_children=list(primary_nodes), + contained_children=list(contained_nodes), + contained_prefix=(), + primary_overrides=overrides, + contained_overrides=contained_overrides, + coverage=coverage, + events=events, + ) + return ContainedGraftResult(nodes=nodes, match_overrides=overrides, events=events) + + +def _coverage_by_path( + ranges: list[ResolvedHierarchyRange], +) -> dict[tuple[str, ...], tuple[int, int]]: + coverage: dict[tuple[str, ...], tuple[int, int]] = {} + for item in ranges: + for depth in range(1, len(item.path_titles) + 1): + path = item.path_titles[:depth] + span = coverage.get(path) + if span is None: + coverage[path] = (item.start_page, item.end_page) + else: + coverage[path] = ( + min(span[0], item.start_page), + max(span[1], item.end_page), + ) + return coverage + + +def _graft_children( + *, + parent: TitleNode | None, + parent_path: tuple[str, ...], + primary_children: list[TitleNode], + contained_children: list[TitleNode], + contained_prefix: tuple[str, ...], + primary_overrides: dict[tuple[str, ...], TitleMatch], + contained_overrides: dict[tuple[str, ...], TitleMatch], + coverage: dict[tuple[str, ...], tuple[int, int]], + events: list[dict[str, Any]], +) -> list[TitleNode]: + working = list(primary_children) + for child in contained_children: + working = _graft_one_child( + parent=parent, + parent_path=parent_path, + primary_children=working, + child=child, + contained_prefix=contained_prefix, + primary_overrides=primary_overrides, + contained_overrides=contained_overrides, + coverage=coverage, + events=events, + ) + return working + + +def _graft_one_child( + *, + parent: TitleNode | None, + parent_path: tuple[str, ...], + primary_children: list[TitleNode], + child: TitleNode, + contained_prefix: tuple[str, ...], + primary_overrides: dict[tuple[str, ...], TitleMatch], + contained_overrides: dict[tuple[str, ...], TitleMatch], + coverage: dict[tuple[str, ...], tuple[int, int]], + events: list[dict[str, Any]], +) -> list[TitleNode]: + contained_path = (*contained_prefix, child.title) + contained_match = contained_overrides.get(contained_path) + if contained_match is None: + return _graft_children( + parent=parent, + parent_path=parent_path, + primary_children=primary_children, + contained_children=list(child.children), + contained_prefix=contained_path, + primary_overrides=primary_overrides, + contained_overrides=contained_overrides, + coverage=coverage, + events=events, + ) + + start = contained_match.page + hit_indexes = _sibling_hits( + siblings=primary_children, + parent_path=parent_path, + start=start, + primary_overrides=primary_overrides, + ) + if hit_indexes: + if len(hit_indexes) > 1: + logger.info( + "{} collision parent_path={} start={} sibling_titles={}", + _LOG_PREFIX, + parent_path, + start, + [primary_children[index].title for index in hit_indexes], + ) + index = hit_indexes[0] + matched = primary_children[index] + matched_path = (*parent_path, matched.title) + events.append( + { + "action": "dedup", + "contained_path": contained_path, + "primary_path": matched_path, + "start": start, + "title_equal": matched.title == child.title, + } + ) + _remap_overrides( + contained_overrides=contained_overrides, + primary_overrides=primary_overrides, + old_prefix=contained_path, + new_prefix=matched_path, + drop_root=True, + ) + new_matched = replace( + matched, + children=_graft_children( + parent=matched, + parent_path=matched_path, + primary_children=list(matched.children), + contained_children=list(child.children), + contained_prefix=contained_path, + primary_overrides=primary_overrides, + contained_overrides=contained_overrides, + coverage=coverage, + events=events, + ), + ) + updated = list(primary_children) + updated[index] = new_matched + return updated + + if parent is not None: + span = coverage.get(parent_path) + if span is not None and span[0] <= start <= span[1]: + return _attach_new_child( + parent=parent, + parent_path=parent_path, + primary_children=primary_children, + child=child, + contained_path=contained_path, + start=start, + primary_overrides=primary_overrides, + contained_overrides=contained_overrides, + events=events, + ) + return primary_children + + covering = _longest_covering_path(coverage, start) + if covering is None: + events.append( + { + "action": "skip", + "contained_path": contained_path, + "start": start, + } + ) + return primary_children + covering_node = _node_at_path(primary_children, covering) + if covering_node is None: + events.append( + { + "action": "skip", + "contained_path": contained_path, + "start": start, + } + ) + return primary_children + updated_parent = replace( + covering_node, + children=_graft_one_child( + parent=covering_node, + parent_path=covering, + primary_children=list(covering_node.children), + child=child, + contained_prefix=contained_prefix, + primary_overrides=primary_overrides, + contained_overrides=contained_overrides, + coverage=coverage, + events=events, + ), + ) + return _replace_node_at_path(primary_children, covering, updated_parent) + + +def _attach_new_child( + *, + parent: TitleNode, + parent_path: tuple[str, ...], + primary_children: list[TitleNode], + child: TitleNode, + contained_path: tuple[str, ...], + start: int, + primary_overrides: dict[tuple[str, ...], TitleMatch], + contained_overrides: dict[tuple[str, ...], TitleMatch], + events: list[dict[str, Any]], +) -> list[TitleNode]: + grafted = _rebase_levels(child, (parent.level + 1) - child.level) + new_path = (*parent_path, grafted.title) + events.append( + { + "action": "attach", + "contained_path": contained_path, + "primary_path": new_path, + "start": start, + } + ) + _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( + siblings=primary_children, + new_node=grafted, + start=start, + parent_path=parent_path, + primary_overrides=primary_overrides, + ) + + +def _sibling_hits( + *, + siblings: list[TitleNode], + parent_path: tuple[str, ...], + start: int, + primary_overrides: dict[tuple[str, ...], TitleMatch], +) -> list[int]: + hits: list[int] = [] + for index, sibling in enumerate(siblings): + match = primary_overrides.get((*parent_path, sibling.title)) + if match is not None and match.page == start: + hits.append(index) + return hits + + +def _remap_overrides( + *, + contained_overrides: dict[tuple[str, ...], TitleMatch], + primary_overrides: dict[tuple[str, ...], TitleMatch], + old_prefix: tuple[str, ...], + new_prefix: tuple[str, ...], + drop_root: bool, +) -> None: + prefix_len = len(old_prefix) + for path, match in contained_overrides.items(): + if path[:prefix_len] != old_prefix: + continue + if drop_root and path == old_prefix: + continue + new_path = new_prefix + path[prefix_len:] + if new_path not in primary_overrides: + primary_overrides[new_path] = match + + +def _rebase_levels(node: TitleNode, delta: int) -> TitleNode: + return replace( + node, + level=node.level + delta, + children=[_rebase_levels(child, delta) for child in node.children], + ) + + +def _insert_by_start( + *, + siblings: list[TitleNode], + new_node: TitleNode, + start: int, + parent_path: tuple[str, ...], + primary_overrides: dict[tuple[str, ...], TitleMatch], +) -> list[TitleNode]: + insert_at = len(siblings) + for index, sibling in enumerate(siblings): + match = primary_overrides.get((*parent_path, sibling.title)) + if match is not None and match.page > start: + insert_at = index + break + updated = list(siblings) + updated.insert(insert_at, new_node) + return updated + + +def _longest_covering_path( + coverage: dict[tuple[str, ...], tuple[int, int]], + start: int, +) -> tuple[str, ...] | None: + covering: tuple[str, ...] | None = None + for path, span in coverage.items(): + if span[0] <= start <= span[1] and ( + covering is None or len(path) > len(covering) + ): + covering = path + return covering + + +def _node_at_path(nodes: list[TitleNode], path: tuple[str, ...]) -> TitleNode | None: + current: list[TitleNode] = nodes + node: TitleNode | None = None + for title in path: + node = next((item for item in current if item.title == title), None) + if node is None: + return None + current = list(node.children) + return node + + +def _replace_node_at_path( + nodes: list[TitleNode], + path: tuple[str, ...], + new_node: TitleNode, +) -> list[TitleNode]: + title, *rest = path + updated: list[TitleNode] = [] + for node in nodes: + if node.title != title: + updated.append(node) + continue + if not rest: + updated.append(new_node) + continue + updated.append( + replace( + 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_link_enrichment.py b/apps/worker/app/services/document_agent/structure/toc_link_enrichment.py index e332f1eb0..f82e97ded 100644 --- a/apps/worker/app/services/document_agent/structure/toc_link_enrichment.py +++ b/apps/worker/app/services/document_agent/structure/toc_link_enrichment.py @@ -1,38 +1,28 @@ -"""Optional TOC hyperlink enrichment after VLM title extraction. - -Only runs when TOC pages actually contain internal links. A link is attached to -a ``toc_with_level`` entry only when anchor text character-matches the extracted -heading. Unmatched entries are left unchanged (no ``link`` field). +"""Attach TOC-page hyperlinks onto VLM entries before calibration. + +PROFILE calls this after ``extract.toc_with_boundaries`` and before +``run_toc_anchoring``. Only runs when TOC pages actually contain internal +page hyperlinks (``page.get_links()``). + +Page-number convention (all 1-based after normalize): + - ``get_links()`` dest ``page``: already 1-based — do not add 1. + - VLM / probe pages: already 1-based. + - TODO(bookmarks): ``get_toc()`` ``meta.page`` is 0-based — add 1 when wired. + +Matching (strict): + - Walk VLM ``toc_with_level`` entries in order, once each. + - ``heading.strip() in anchor_text.strip()``. + - Attach only when exactly one link hits; zero or many → leave unmatched. + - Cross-line / truncated anchors are not special-cased. """ from __future__ import annotations -import re -import unicodedata from dataclasses import dataclass from typing import Any from loguru import logger -_APOSTROPHE_TRANS = str.maketrans( - { - "\u2018": "'", - "\u2019": "'", - "\u201b": "'", - "\u2032": "'", - "\u00b4": "'", - "\u0060": "'", - } -) -_NON_ALNUM = re.compile(r"[^a-z0-9]+") - - -def normalize_toc_heading(text: str) -> str: - """Normalize heading / anchor text for strict character matching.""" - raw = unicodedata.normalize("NFKC", str(text or "")).translate(_APOSTROPHE_TRANS) - raw = raw.replace("\xa0", " ").strip().lower() - return _NON_ALNUM.sub(" ", raw).strip() - @dataclass(frozen=True) class TocPageLink: @@ -83,8 +73,22 @@ def _anchor_text_for_rect(page: Any, rect: Any) -> str: return " ".join(part for _, _, part in hit).strip() +def _link_dest_physical_page(raw_page: Any) -> int: + """Normalize ``page.get_links()`` destination to 1-based physical page. + + PyMuPDF page hyperlinks expose ``link["page"]`` already as 1-based (int or + digit string). Do **not** add 1 here — that off-by-one sent every TOC link + one page past the real click target. + + TODO(bookmarks): ``doc.get_toc()`` outline ``meta.page`` is 0-based. When + bookmark signal is wired into calibration, convert that field with ``+1`` + (or use get_toc's 1-based display page) before merging with links / VLM. + """ + return int(raw_page) + + def collect_toc_page_links(pdf_path: str, toc_pages: list[int]) -> list[TocPageLink]: - """Collect internal goto links on TOC pages with nearby anchor text.""" + """Collect internal page hyperlinks on TOC pages with nearby anchor text.""" import fitz if not toc_pages: @@ -98,26 +102,23 @@ def collect_toc_page_links(pdf_path: str, toc_pages: list[int]) -> list[TocPageL continue page = doc[toc_page - 1] for link in page.get_links() or []: - # PyMuPDF: LINK_GOTO=1, LINK_NAMED=4 commonly used for TOC. - kind = link.get("kind") - dest_idx = link.get("page") - if dest_idx is None: + # Page hyperlinks only (NAMED/GOTO via get_links). Not bookmarks. + dest_raw = link.get("page") + if dest_raw is None: continue try: - dest_physical = int(dest_idx) + 1 + dest_physical = _link_dest_physical_page(dest_raw) except (TypeError, ValueError): continue if dest_physical < 1 or dest_physical > doc.page_count: continue - # Skip obvious self / header "back to TOC" loops to same/near page. - if abs(dest_physical - toc_page) <= 1: - continue rect = link.get("from") if rect is None: continue anchor = _anchor_text_for_rect(page, rect) if not anchor: continue + kind = link.get("kind") out.append( TocPageLink( toc_page=toc_page, @@ -131,52 +132,13 @@ def collect_toc_page_links(pdf_path: str, toc_pages: list[int]) -> list[TocPageL return out -def _is_page_number_label(text: str) -> bool: - t = str(text or "").strip() - if not t: - return False - if t.isdigit(): - return True - # Roman / folio labels: iv, xii, F-1 - if re.fullmatch(r"[ivxlcdm]+", t.lower()): - return True - if re.fullmatch(r"[A-Za-z]-?\d+", t): - return True - return False - - -def _headings_match(heading: str, anchor: str) -> bool: - h = normalize_toc_heading(heading) - a = normalize_toc_heading(anchor) - if not h or not a: - return False - if h == a: - return True - # Anchor sometimes truncates long titles; require substantial prefix/containment. - if len(h) >= 12 and (a.startswith(h) or h.startswith(a)): - shorter, longer = (a, h) if len(a) <= len(h) else (h, a) - if len(shorter) >= 12 and shorter in longer: - return True - return False - - def match_toc_entries_to_links( entries: list[dict[str, Any]], links: list[TocPageLink], ) -> tuple[list[dict[str, Any]], int]: - """Return new entry dicts; only matched ones gain a ``link`` object.""" - # Index title-like anchors (skip pure page-number chips). - title_links = [ - link - for link in links - if not _is_page_number_label(link.anchor_text) - and normalize_toc_heading(link.anchor_text) - and normalize_toc_heading(link.anchor_text) not in {"table of contents", "contents"} - ] - + """Attach ``link`` when heading.strip() is in exactly one link anchor.""" matched = 0 enriched: list[dict[str, Any]] = [] - used_dest_for_heading: set[str] = set() for entry in entries: if not isinstance(entry, dict): @@ -186,43 +148,27 @@ def match_toc_entries_to_links( "level": entry.get("level"), "page_number": entry.get("page_number"), } - # Preserve unknown keys except stale link from a prior run. for key, value in entry.items(): if key in new_entry or key == "link": continue new_entry[key] = value heading = str(entry.get("heading") or "").strip() - if not heading or not title_links: - enriched.append(new_entry) - continue - - hits = [link for link in title_links if _headings_match(heading, link.anchor_text)] - if not hits: - enriched.append(new_entry) - continue - - # Prefer unique dest; if multiple dests, refuse (ambiguous). - dests = {link.dest_physical_page for link in hits} - if len(dests) != 1: - logger.info( - "[toc_link_enrich] ambiguous link for heading={!r} dests={}", - heading, - sorted(dests), - ) + if not heading or not links: enriched.append(new_entry) continue - chosen = hits[0] - heading_key = normalize_toc_heading(heading) - # One heading → one link attachment (first wins if duplicates). - if heading_key in used_dest_for_heading: + hits = [ + link + for link in links + if heading in str(link.anchor_text or "").strip() + ] + if len(hits) != 1: enriched.append(new_entry) continue - used_dest_for_heading.add(heading_key) new_entry["link"] = { - "physical_page": chosen.dest_physical_page, + "physical_page": hits[0].dest_physical_page, } matched += 1 enriched.append(new_entry) diff --git a/apps/worker/app/services/document_parser/profiling/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py index 5d5a42028..0f3a2e31e 100644 --- a/apps/worker/app/services/document_parser/profiling/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -33,6 +33,7 @@ def profile_document( output_dir: str | None = None, skip_shard_plan: bool = False, oversized_policy: Literal["chunk", "page_memory"] = "chunk", + skip_toc_anchoring: bool = False, ) -> ParserDocumentProfile: """ General document profiling entry point. @@ -50,6 +51,10 @@ def profile_document( oversized_policy: Controls oversized PDF admission. ``chunk`` applies the MinerU shard gate, while ``page_memory`` lets the page-memory track continue to structural profiling. + skip_toc_anchoring: When True, TOC find/extract/link-attach still run, + but ``run_toc_anchoring`` is skipped. Used by the staged debug + path (Stage-1 TOC after Stage-0 bootstrap) so calibration stays + in Stage-2. Returns: ParserDocumentProfile @@ -67,6 +72,7 @@ def profile_document( output_dir=output_dir, skip_shard_plan=skip_shard_plan, oversized_policy=oversized_policy, + skip_toc_anchoring=skip_toc_anchoring, ) finally: if not visual_debug_enabled(): @@ -88,6 +94,7 @@ def _profile_pdf( output_dir: str | None, skip_shard_plan: bool = False, oversized_policy: Literal["chunk", "page_memory"] = "chunk", + skip_toc_anchoring: bool = False, ) -> ParserDocumentProfile: with _profile_db_context(enabled=bool(job_id)) as db: return _profile_pdf_with_db( @@ -98,6 +105,7 @@ def _profile_pdf( db=db, skip_shard_plan=skip_shard_plan, oversized_policy=oversized_policy, + skip_toc_anchoring=skip_toc_anchoring, ) @@ -110,6 +118,7 @@ def _profile_pdf_with_db( db: Any | None, skip_shard_plan: bool = False, oversized_policy: Literal["chunk", "page_memory"] = "chunk", + skip_toc_anchoring: bool = False, ) -> ParserDocumentProfile: profile_job_id = job_id or filename agent_output_dir = os.path.join(output_dir, "_doc_agent") if output_dir else None @@ -130,6 +139,7 @@ def _profile_pdf_with_db( "planner_model": settings.IMAGE_MODEL, "vlm_model": settings.IMAGE_MODEL, "toc_profile_enabled": page_toc_enabled, + "skip_toc_anchoring": bool(skip_toc_anchoring), "model": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, }, ) diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 976c9ecfe..4b382e738 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -34,8 +34,6 @@ ) - - @dataclass(frozen=True) class SectionSkeleton: section_path: str @@ -87,7 +85,9 @@ def extract_section_skeletons( if not isinstance(skeleton_anchor_raw, dict) or not isinstance( skeleton_nodes_raw, list ): - raise ValueError("anatomy.skeleton_anchor/skeleton_nodes missing after TOC extract") + raise ValueError( + "anatomy.skeleton_anchor/skeleton_nodes missing after TOC extract" + ) skeleton_anchor = deserialize_skeleton_anchor(skeleton_anchor_raw) resolve_nodes = [ @@ -110,17 +110,22 @@ def extract_section_skeletons( getattr(toc_result, "toc_pages", None), page_count, ) + pending_records = list(getattr(anatomy, "pending_skeleton_anchors", None) or []) + parallel_pending = _parallel_pending_tocs( + pending_tocs=pending_tocs, + pending_records=pending_records, + ) + parallel_tocs = [toc for toc, _record in parallel_pending] primary_page_count = page_count primary_body_pages = body_pages - if pending_tocs: - pending_starts: list[int] = [] - for pending_toc in pending_tocs: - start = toc_range_start(pending_toc) - if start is not None: - pending_starts.append(start) - if pending_starts: - primary_page_count = min(pending_starts) - 1 - primary_body_pages = [page for page in body_pages if page <= primary_page_count] + pending_starts: list[int] = [] + for toc in parallel_tocs: + start = toc_range_start(toc) + if start is not None: + pending_starts.append(start) + if pending_starts: + primary_page_count = min(pending_starts) - 1 + primary_body_pages = [page for page in body_pages if page <= primary_page_count] match_overrides = skeleton_anchor.match_overrides null_page_report = skeleton_anchor.null_page_report @@ -178,11 +183,9 @@ def extract_section_skeletons( for item in ranges ] - pending_records = list(getattr(anatomy, "pending_skeleton_anchors", None) or []) - if pending_tocs and pending_records: + if parallel_pending: secondary_skeletons = _resolve_pending_tocs( - pending_tocs=pending_tocs, - pending_records=pending_records, + parallel_pending=parallel_pending, page_texts=page_texts, page_count=page_count, filename=filename, @@ -205,7 +208,9 @@ def _range_to_skeleton( start_page = _clamp_page(item.start_page, page_count) end_page = _clamp_page(item.end_page, page_count) # Keep original TOC titles (incl. numbering) in section_path / HIERARCHY. - path_titles = [str(title).strip() for title in item.path_titles if str(title).strip()] + path_titles = [ + str(title).strip() for title in item.path_titles if str(title).strip() + ] section_path = join_document_path([filename, *path_titles]) parent_path = ( join_document_path([filename, *path_titles[:-1]]) @@ -254,45 +259,49 @@ def _page_count(anatomy: Any | None) -> int: return max(int(getattr(anatomy, "page_count", 0) or 0), 0) -def _resolve_pending_tocs( +def _parallel_pending_tocs( *, pending_tocs: list[dict[str, Any]], pending_records: list[dict[str, Any]], +) -> list[tuple[dict[str, Any], dict[str, Any]]]: + records_by_range: dict[tuple[Any, ...], dict[str, Any]] = {} + for record in pending_records: + toc = record.get("toc") + if isinstance(toc, dict): + records_by_range[tuple(toc.get("toc_range") or [])] = record + + parallel_pending: list[tuple[dict[str, Any], dict[str, Any]]] = [] + for pending_toc in pending_tocs: + record = records_by_range.get(tuple(pending_toc.get("toc_range") or [])) + if record is None: + continue + relationship = record.get("relationship") + if relationship in {"unresolvable", "contained"}: + continue + if relationship != "parallel": + raise ValueError("pending TOC relationship missing after PROFILE classify") + parallel_pending.append((pending_toc, record)) + return parallel_pending + + +def _resolve_pending_tocs( + *, + parallel_pending: list[tuple[dict[str, Any], dict[str, Any]]], page_texts: dict[int, str], page_count: int, filename: str, body_pages: list[int], ) -> list[SectionSkeleton]: - """Graft pending TOCs from PROFILE-persisted skeleton anchors.""" - if not pending_tocs or not pending_records: + """Resolve parallel pending TOCs as secondary forests.""" + if not parallel_pending: return [] - records_by_range: dict[tuple[Any, ...], dict[str, Any]] = {} - for record in pending_records: - toc = record.get("toc") - if not isinstance(toc, dict): - continue - key = tuple(toc.get("toc_range") or []) - records_by_range[key] = record - + parallel_tocs = [toc for toc, _record in parallel_pending] all_secondary_skeletons: list[SectionSkeleton] = [] - for i, pending_toc in enumerate(pending_tocs): + for i, (pending_toc, record) in enumerate(parallel_pending): toc_range = pending_toc.get("toc_range") - record = records_by_range.get(tuple(toc_range or [])) - if record is None: - continue relationship = record.get("relationship") - if relationship == "unresolvable": - logger.info( - "[page_memory.skeleton] pending TOC toc_range={}: unresolvable, skipping", - toc_range, - ) - continue - if relationship not in {"parallel", "contained"}: - raise ValueError( - "pending TOC relationship missing after PROFILE classify" - ) resolve_nodes_raw = record.get("nodes") or [] resolve_nodes = [ deserialize_title_node(node) @@ -310,7 +319,7 @@ def _resolve_pending_tocs( raise ValueError("pending TOC offset missing after PROFILE classify") toc_scope_end, toc_body_pages = pending_toc_body_scope( - pending_tocs=pending_tocs, + pending_tocs=parallel_tocs, index=i, page_count=page_count, body_pages=body_pages, @@ -326,7 +335,9 @@ def _resolve_pending_tocs( } locate_summary["null_page_parent_locate"] = { "attempted": len(null_page_report), - "located": sum(1 for row in null_page_report if row.get("page") is not None), + "located": sum( + 1 for row in null_page_report if row.get("page") is not None + ), "unresolved": sum( 1 for row in null_page_report if row.get("result") == "unresolved" ), diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index 512af8ca6..5a0d8b49e 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -221,7 +221,13 @@ def fake_toc() -> None: profile = coordinator.run_coarse() assert profile.category == "Research Report" - assert calls == ["planner", "text_scan", "toc", "probe.page_assets", "aggregate.doc_stats"] + assert calls == [ + "planner", + "text_scan", + "probe.page_assets", + "aggregate.doc_stats", + "toc", + ] assert coordinator.blackboard.global_signals["assets_probed"] is True diff --git a/apps/worker/tests/contract/test_toc_graft_contract.py b/apps/worker/tests/contract/test_toc_graft_contract.py new file mode 100644 index 000000000..f355bc265 --- /dev/null +++ b/apps/worker/tests/contract/test_toc_graft_contract.py @@ -0,0 +1,583 @@ +"""Contract tests for contained TOC graft. Synthetic trees only; no PDF.""" + +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.budget import BudgetTracker +from app.services.document_agent.manifest import ( + PageAnatomyMap, + PageFeature, + PageLabel, + Shard, + ShardPlan, + TocResult, + ToolContext, +) +from app.services.document_agent.state import AgentBlackboard +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_graft import graft_contained_toc +from app.services.page_memory.skeleton_extractor import extract_section_skeletons + + +def _match(title: str, page: int) -> TitleMatch: + return TitleMatch( + page=page, + confidence=1.0, + source="anchored", + matched_line=title, + score=1.0, + 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 _graft( + *, + primary: list[TitleNode], + primary_pages: dict[tuple[str, ...], int], + contained: list[TitleNode], + contained_pages: dict[tuple[str, ...], int], + page_count: int = 50, +) -> object: + body_pages = list(range(1, page_count + 1)) + page_texts = {page: "" for page in body_pages} + return graft_contained_toc( + primary_nodes=primary, + primary_overrides=_overrides(primary_pages), + contained_nodes=contained, + contained_overrides=_overrides(contained_pages), + page_count=page_count, + page_texts=page_texts, + body_pages=body_pages, + ) + + +def _node_at(nodes: list[TitleNode], path: tuple[str, ...]) -> TitleNode | None: + current = nodes + found: TitleNode | None = None + for title in path: + found = next((node for node in current if node.title == title), None) + if found is None: + return None + current = list(found.children) + return found + + +def test_dedup_keeps_primary_title_and_hangs_child() -> None: + result = _graft( + primary=[TitleNode(title="第一章", level=1, printed_page=10)], + primary_pages={("第一章",): 10}, + contained=[ + TitleNode( + title="第一章", + level=1, + printed_page=10, + children=[TitleNode(title="1.2", level=2, printed_page=12)], + ) + ], + contained_pages={("第一章",): 10, ("第一章", "1.2"): 12}, + ) + + assert [node.title for node in result.nodes] == ["第一章"] + child = _node_at(result.nodes, ("第一章", "1.2")) + assert child is not None + assert child.level == 2 + assert result.match_overrides[("第一章",)].page == 10 + assert result.match_overrides[("第一章", "1.2")].page == 12 + dedup = next(event for event in result.events if event["action"] == "dedup") + assert dedup["title_equal"] is True + + +def test_dedup_ignores_title_and_keeps_primary_override() -> None: + primary_match = _match("第一章", 10) + body_pages = list(range(1, 51)) + result = graft_contained_toc( + primary_nodes=[TitleNode(title="第一章", level=1, printed_page=10)], + primary_overrides={("第一章",): primary_match}, + contained_nodes=[ + TitleNode( + title="Chapter 1", + level=1, + printed_page=10, + children=[TitleNode(title="1.2", level=2, printed_page=12)], + ) + ], + contained_overrides=_overrides({("Chapter 1",): 10, ("Chapter 1", "1.2"): 12}), + page_count=50, + page_texts={page: "" for page in body_pages}, + body_pages=body_pages, + ) + + assert [node.title for node in result.nodes] == ["第一章"] + assert _node_at(result.nodes, ("第一章", "1.2")) is not None + assert result.match_overrides[("第一章",)] is primary_match + assert ("Chapter 1",) not in result.match_overrides + dedup = next(event for event in result.events if event["action"] == "dedup") + assert dedup["title_equal"] is False + + +def test_attach_new_child_rebases_level() -> None: + result = _graft( + primary=[TitleNode(title="第一章", level=1, printed_page=10)], + primary_pages={("第一章",): 10}, + contained=[TitleNode(title="1.2", level=1, printed_page=12)], + contained_pages={("1.2",): 12}, + ) + + child = _node_at(result.nodes, ("第一章", "1.2")) + assert child is not None + assert child.level == 2 + assert result.match_overrides[("第一章", "1.2")].page == 12 + assert any(event["action"] == "attach" for event in result.events) + + +def test_same_start_parent_and_child_do_not_collapse() -> None: + result = _graft( + primary=[TitleNode(title="1.1", level=1, printed_page=10)], + primary_pages={("1.1",): 10}, + contained=[ + TitleNode( + title="1.1", + level=1, + printed_page=10, + children=[TitleNode(title="1.1.1", level=2, printed_page=10)], + ) + ], + contained_pages={("1.1",): 10, ("1.1", "1.1.1"): 10}, + ) + + parent = _node_at(result.nodes, ("1.1",)) + child = _node_at(result.nodes, ("1.1", "1.1.1")) + assert parent is not None + assert child is not None + assert parent.title == "1.1" + assert child.title == "1.1.1" + assert [node.title for node in result.nodes] == ["1.1"] + + +def test_two_contained_tocs_graft_in_order() -> None: + first = _graft( + primary=[TitleNode(title="第一章", level=1, printed_page=10)], + primary_pages={("第一章",): 10}, + contained=[ + TitleNode( + title="第一章", + level=1, + printed_page=10, + children=[TitleNode(title="1.2", level=2, printed_page=12)], + ) + ], + contained_pages={("第一章",): 10, ("第一章", "1.2"): 12}, + ) + second = graft_contained_toc( + primary_nodes=first.nodes, + primary_overrides=first.match_overrides, + contained_nodes=[ + TitleNode( + title="第一章", + level=1, + printed_page=10, + children=[TitleNode(title="1.3", level=2, printed_page=20)], + ) + ], + contained_overrides=_overrides({("第一章",): 10, ("第一章", "1.3"): 20}), + page_count=50, + page_texts={page: "" for page in range(1, 51)}, + body_pages=list(range(1, 51)), + ) + + titles = [child.title for child in second.nodes[0].children] + assert titles == ["1.2", "1.3"] + + +def _ctx(*, page_count: int) -> ToolContext: + return ToolContext( + pdf_path="/tmp/doc.pdf", + job_id="job-graft", + blackboard=AgentBlackboard(page_count=page_count), + budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), + trace=None, + settings={}, + ) + + +def _anchor(pages: dict[tuple[str, ...], int]) -> SkeletonAnchor: + return SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides=_overrides(pages), + null_page_report=[], + bulk_count=len(pages), + pruned_count=0, + locate_agent="offset_guided_bulk", + ) + + +def test_parallel_pending_is_not_grafted() -> None: + ctx = _ctx(page_count=30) + ctx.blackboard.toc_hierarchies = [ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [{"heading": "Ch1", "level": 1, "page_number": 2}], + }, + { + "toc_range": [20, 21], + "toc_range_unit": "page", + "toc_with_level": [{"heading": "App", "level": 1, "page_number": 22}], + }, + ] + ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1, 20, 21]) + ctx.blackboard.page_full_text_cache = {page: "body" for page in range(1, 31)} + primary = TitleNode(title="Ch1", level=1, printed_page=2) + pending = TitleNode(title="App", level=1, printed_page=22) + captured: dict[str, object] = {} + + def fake_anchor_hierarchy(**kwargs): + captured["body_pages"] = kwargs["body_pages"] + return [primary], _anchor({("Ch1",): 2}) + + with ( + patch( + "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + side_effect=fake_anchor_hierarchy, + ), + patch( + "app.services.document_agent.structure.toc_anchoring.calibrate_offset", + return_value=object(), + ), + patch( + "app.services.document_agent.structure.toc_anchoring.pick_primary_offset", + return_value=0, + ), + patch( + "app.services.document_agent.structure.toc_anchoring.classify_toc_relationship", + return_value="parallel", + ), + patch( + "app.services.document_agent.structure.toc_anchoring.finalize_calibration_result", + return_value=([pending], _anchor({("App",): 22}), True), + ), + ): + run_toc_anchoring(ctx) + + assert 22 in captured["body_pages"] + assert [node["title"] for node in ctx.blackboard.skeleton_nodes] == ["Ch1"] + record = ctx.blackboard.pending_skeleton_anchors[0] + assert record["relationship"] == "parallel" + assert "grafted" not in record + + +def test_profile_grafts_contained_and_keeps_original_pending() -> None: + ctx = _ctx(page_count=50) + ctx.blackboard.toc_hierarchies = [ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [{"heading": "第一章", "level": 1, "page_number": 10}], + }, + { + "toc_range": [20, 21], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "第一章", "level": 1, "page_number": 10}, + {"heading": "1.2", "level": 2, "page_number": 12}, + ], + }, + ] + ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1, 20, 21]) + ctx.blackboard.page_full_text_cache = {page: "body" for page in range(1, 51)} + primary = TitleNode(title="第一章", level=1, printed_page=10) + contained = TitleNode( + title="第一章", + level=1, + printed_page=10, + children=[TitleNode(title="1.2", level=2, printed_page=12)], + ) + + with ( + patch( + "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + return_value=([primary], _anchor({("第一章",): 10})), + ), + patch( + "app.services.document_agent.structure.toc_anchoring.calibrate_offset", + return_value=object(), + ), + patch( + "app.services.document_agent.structure.toc_anchoring.pick_primary_offset", + return_value=0, + ), + patch( + "app.services.document_agent.structure.toc_anchoring.classify_toc_relationship", + return_value="contained", + ), + patch( + "app.services.document_agent.structure.toc_anchoring.finalize_calibration_result", + return_value=( + [contained], + _anchor({("第一章",): 10, ("第一章", "1.2"): 12}), + True, + ), + ), + ): + run_toc_anchoring(ctx) + + assert ctx.blackboard.skeleton_nodes[0]["title"] == "第一章" + assert ctx.blackboard.skeleton_nodes[0]["children"][0]["title"] == "1.2" + record = ctx.blackboard.pending_skeleton_anchors[0] + assert record["grafted"] is True + assert record["nodes"][0]["title"] == "第一章" + assert "第一章 / 1.2" in ctx.blackboard.skeleton_anchor["match_overrides"] + + +def _feature(page: int) -> PageFeature: + return PageFeature( + page=page, + raw_text_length=20, + text_density=0.1, + image_coverage=0.0, + image_count=0, + table_count=0, + drawings_count=0, + orientation="portrait", + width=72.0, + height=72.0, + has_asset=False, + is_blank_like=False, + ) + + +def _anatomy( + *, + page_count: int, + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], int], + hierarchies: list[dict[str, object]], + pending_records: list[dict[str, object]], + toc_pages: list[int], +) -> PageAnatomyMap: + return PageAnatomyMap( + job_id="job-graft", + file_path="/tmp/doc.pdf", + page_count=page_count, + page_features=[_feature(page) for page in range(1, page_count + 1)], + page_labels=[ + PageLabel(page=page, kind="normal", confidence=1.0) + for page in range(1, page_count + 1) + ], + toc_result=TocResult(method="vlm_batch", toc_pages=toc_pages), + shard_plan=ShardPlan( + enabled=False, + reason="not_needed", + shards=[ + Shard( + shard_index=0, + page_start=1, + page_end=page_count, + page_offset=0, + anchor_type="forced_max_size", + anchor_evidence="test", + confidence=1.0, + ) + ], + ), + toc_hierarchies=hierarchies, + toc_page_offset=0, + skeleton_anchor=serialize_skeleton_anchor(_anchor(overrides)), + skeleton_nodes=[serialize_title_node(node) for node in nodes], + pending_skeleton_anchors=pending_records, + ) + + +def _pending_record( + *, + toc: dict[str, object], + relationship: str, + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], int], + grafted: bool = False, +) -> dict[str, object]: + record: dict[str, object] = { + "toc": toc, + "relationship": relationship, + "nodes": [serialize_title_node(node) for node in nodes], + "skeleton_anchor": serialize_skeleton_anchor(_anchor(overrides)), + } + if grafted: + record["grafted"] = True + record["graft"] = [] + return record + + +def test_page_contained_does_not_cut_primary_window() -> None: + pending_toc = { + "toc_range": [200, 201], + "toc_range_unit": "page", + "toc_with_level": [{"heading": "Inner", "level": 1, "page_number": 210}], + } + anatomy = _anatomy( + page_count=250, + nodes=[ + TitleNode(title="Ch1", level=1, printed_page=3), + TitleNode(title="Ch2", level=1, printed_page=120), + ], + overrides={("Ch1",): 3, ("Ch2",): 120}, + hierarchies=[ + { + "toc_range": [1, 2], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 3}, + {"heading": "Ch2", "level": 1, "page_number": 120}, + ], + }, + pending_toc, + ], + pending_records=[ + _pending_record( + toc=pending_toc, + relationship="contained", + nodes=[TitleNode(title="Inner", level=1, printed_page=210)], + overrides={("Inner",): 210}, + grafted=True, + ) + ], + toc_pages=[1, 2, 200, 201], + ) + page_texts = {page: "body" for page in range(1, 251)} + skeletons = extract_section_skeletons( + anatomy=anatomy, + filename="doc.pdf", + page_texts=page_texts, + ) + ch2 = next(item for item in skeletons if item.title == "Ch2") + assert ch2.end_page > 201 + assert all(item.title != "Inner" for item in skeletons) + + +def test_page_parallel_still_cuts_primary_window() -> None: + pending_toc = { + "toc_range": [200, 201], + "toc_range_unit": "page", + "toc_with_level": [{"heading": "App", "level": 1, "page_number": 210}], + } + anatomy = _anatomy( + page_count=250, + nodes=[ + TitleNode(title="Ch1", level=1, printed_page=3), + TitleNode(title="Ch2", level=1, printed_page=120), + ], + overrides={("Ch1",): 3, ("Ch2",): 120}, + hierarchies=[ + { + "toc_range": [1, 2], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 3}, + {"heading": "Ch2", "level": 1, "page_number": 120}, + ], + }, + pending_toc, + ], + pending_records=[ + _pending_record( + toc=pending_toc, + relationship="parallel", + nodes=[TitleNode(title="App", level=1, printed_page=210)], + overrides={("App",): 210}, + ) + ], + toc_pages=[1, 2, 200, 201], + ) + page_texts = {page: "body" for page in range(1, 251)} + page_texts[210] = "App" + skeletons = extract_section_skeletons( + anatomy=anatomy, + filename="doc.pdf", + page_texts=page_texts, + ) + ch2 = next(item for item in skeletons if item.title == "Ch2") + assert ch2.end_page == 199 + assert any(item.title == "App" for item in skeletons) + + +def test_page_grafted_contained_is_not_flattened() -> None: + pending_toc = { + "toc_range": [20, 21], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "第一章", "level": 1, "page_number": 10}, + {"heading": "1.2", "level": 2, "page_number": 12}, + ], + } + anatomy = _anatomy( + page_count=50, + nodes=[ + TitleNode( + title="第一章", + level=1, + printed_page=10, + children=[TitleNode(title="1.2", level=2, printed_page=12)], + ) + ], + overrides={("第一章",): 10, ("第一章", "1.2"): 12}, + hierarchies=[ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "第一章", "level": 1, "page_number": 10} + ], + }, + pending_toc, + ], + pending_records=[ + _pending_record( + toc=pending_toc, + relationship="contained", + nodes=[ + TitleNode( + title="第一章", + level=1, + printed_page=10, + children=[TitleNode(title="1.2", level=2, printed_page=12)], + ) + ], + overrides={("第一章",): 10, ("第一章", "1.2"): 12}, + grafted=True, + ) + ], + toc_pages=[1, 20, 21], + ) + page_texts = {page: "body" for page in range(1, 51)} + skeletons = extract_section_skeletons( + anatomy=anatomy, + filename="doc.pdf", + page_texts=page_texts, + ) + titled = [item for item in skeletons if item.title == "1.2"] + assert len(titled) == 1 + assert "第一章" in titled[0].section_path diff --git a/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py b/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py new file mode 100644 index 000000000..ed48f979f --- /dev/null +++ b/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py @@ -0,0 +1,190 @@ +"""PROFILE attaches TOC-page links before calibration.""" + +from __future__ import annotations + +import os +from types import SimpleNamespace +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.coordinator import ProfileCoordinator +from app.services.document_agent.manifest import ToolResult + + +def _hierarchy() -> list[dict[str, object]]: + return [ + { + "toc_range": [2, 5], + "toc_with_level": [ + {"heading": "Ch1", "level": 1, "page_number": 2}, + ], + } + ] + + +def _coordinator() -> ProfileCoordinator: + coordinator = ProfileCoordinator(pdf_path="/tmp/doc.pdf", job_id="job-link-order") + coordinator.blackboard.toc_hierarchies = _hierarchy() + return coordinator + + +def test_profile_attaches_toc_links_before_anchoring() -> None: + coordinator = _coordinator() + seen: dict[str, object] = {} + + def fake_enrich(*, pdf_path: str, toc_hierarchies: list[dict[str, object]]): + assert pdf_path == "/tmp/doc.pdf" + attached = [ + { + **toc_hierarchies[0], + "toc_with_level": [ + { + **toc_hierarchies[0]["toc_with_level"][0], # type: ignore[index] + "link": {"physical_page": 8}, + } + ], + } + ] + return attached, SimpleNamespace( + entries_matched=1, + entries_total=1, + skipped_no_links=False, + ) + + def fake_anchor(ctx) -> None: + entry = ctx.blackboard.toc_hierarchies[0]["toc_with_level"][0] + seen["physical_page"] = (entry.get("link") or {}).get("physical_page") + + with ( + patch.object( + ProfileCoordinator, + "_dispatch_profile_tool", + return_value=ToolResult(status="ok"), + ), + patch( + "app.services.document_agent.coordinator.enrich_toc_hierarchies_with_links", + side_effect=fake_enrich, + ), + patch( + "app.services.document_agent.coordinator.run_toc_anchoring", + side_effect=fake_anchor, + ), + ): + coordinator._run_toc_extraction_pipeline() + + assert seen["physical_page"] == 8 + assert ( + coordinator.blackboard.toc_hierarchies[0]["toc_with_level"][0]["link"][ + "physical_page" + ] + == 8 + ) + + +def test_link_attach_failure_keeps_hierarchies_and_still_anchors() -> None: + coordinator = _coordinator() + seen = {"anchored": False} + + def fake_anchor(ctx) -> None: + seen["anchored"] = True + entry = ctx.blackboard.toc_hierarchies[0]["toc_with_level"][0] + assert entry.get("link") is None + + with ( + patch.object( + ProfileCoordinator, + "_dispatch_profile_tool", + return_value=ToolResult(status="ok"), + ), + patch( + "app.services.document_agent.coordinator.enrich_toc_hierarchies_with_links", + side_effect=RuntimeError("pymupdf failed"), + ), + patch( + "app.services.document_agent.coordinator.run_toc_anchoring", + side_effect=fake_anchor, + ), + ): + coordinator._run_toc_extraction_pipeline() + + assert seen["anchored"] is True + assert coordinator.blackboard.toc_hierarchies == _hierarchy() + + +def test_skip_toc_anchoring_stops_after_link_attach() -> None: + coordinator = _coordinator() + coordinator.ctx.settings["skip_toc_anchoring"] = True + seen = {"anchored": False} + + def fake_enrich(*, pdf_path: str, toc_hierarchies: list[dict[str, object]]): + return toc_hierarchies, SimpleNamespace( + entries_matched=0, + entries_total=1, + skipped_no_links=True, + ) + + def fake_anchor(_ctx) -> None: + seen["anchored"] = True + + with ( + patch.object( + ProfileCoordinator, + "_dispatch_profile_tool", + return_value=ToolResult(status="ok"), + ), + patch( + "app.services.document_agent.coordinator.enrich_toc_hierarchies_with_links", + side_effect=fake_enrich, + ), + patch( + "app.services.document_agent.coordinator.run_toc_anchoring", + side_effect=fake_anchor, + ), + ): + coordinator._run_toc_extraction_pipeline() + + assert seen["anchored"] is False + assert coordinator.blackboard.skeleton_anchor is None + assert coordinator.blackboard.skeleton_nodes is None + assert coordinator.blackboard.toc_page_offset is None + + +def test_stop_after_asset_probe_skips_toc() -> None: + coordinator = ProfileCoordinator(pdf_path="/tmp/doc.pdf", job_id="job-stage0") + coordinator.ctx.settings["stop_after_asset_probe"] = True + seen = {"toc": False, "assets": False} + + profile = SimpleNamespace( + category="spec", + routing_category="generic", + is_scanned=False, + ) + + def fake_toc(self, *, strict: bool) -> None: + seen["toc"] = True + + def fake_assets(self) -> None: + seen["assets"] = True + + with ( + patch.object(ProfileCoordinator, "_run_bootstrap", return_value=None), + patch.object( + ProfileCoordinator, + "_propose_profile", + return_value=(profile, None, ToolResult(status="ok")), + ), + patch.object(ProfileCoordinator, "_run_text_scan", return_value=None), + patch.object(ProfileCoordinator, "_ensure_toc_profile", fake_toc), + patch.object(ProfileCoordinator, "_ensure_asset_probe", fake_assets), + ): + out = coordinator._run_coarse() + + assert out is profile + assert seen["assets"] is True + assert seen["toc"] is False diff --git a/apps/worker/tests/contract/test_toc_link_match_contract.py b/apps/worker/tests/contract/test_toc_link_match_contract.py new file mode 100644 index 000000000..54b35c78c --- /dev/null +++ b/apps/worker/tests/contract/test_toc_link_match_contract.py @@ -0,0 +1,104 @@ +"""Contract tests for TOC heading → link containment matching.""" + +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 app.services.document_agent.structure.toc_link_enrichment import ( + TocPageLink, + _link_dest_physical_page, + match_toc_entries_to_links, +) + + +def _link(anchor: str, dest: int, *, toc_page: int = 2) -> TocPageLink: + return TocPageLink( + toc_page=toc_page, + dest_physical_page=dest, + anchor_text=anchor, + kind=4, + ) + + +def test_get_links_dest_page_is_already_one_based() -> None: + """``get_links()['page']`` must not be shifted with +1.""" + assert _link_dest_physical_page(6) == 6 + assert _link_dest_physical_page("6") == 6 + assert _link_dest_physical_page("7") == 7 + + +def test_match_exact_one_hit_attaches_physical_page() -> None: + entries = [ + {"heading": " 1.投标人营业执照扫描件; ", "level": 2, "page_number": 2}, + {"heading": "二、施工组织设计", "level": 1, "page_number": 26}, + ] + links = [ + _link("1.投标人营业执照扫描件;...............................", 7), + _link("二、施工组织设计.............................................", 31), + _link("无关导航链接", 99), + ] + + enriched, matched = match_toc_entries_to_links(entries, links) + + assert matched == 2 + assert enriched[0]["link"] == {"physical_page": 7} + assert enriched[1]["link"] == {"physical_page": 31} + assert enriched[0]["heading"] == " 1.投标人营业执照扫描件; " + assert enriched[0]["page_number"] == 2 + + +def test_match_zero_or_many_hits_leaves_entry_unmatched() -> None: + entries = [ + {"heading": "一、资格复审资料", "level": 1, "page_number": 1}, + {"heading": "共用标题", "level": 2, "page_number": 3}, + ] + links = [ + _link("共用标题..............2", 10), + _link("共用标题..............9", 20), + ] + + enriched, matched = match_toc_entries_to_links(entries, links) + + assert matched == 0 + assert "link" not in enriched[0] + assert "link" not in enriched[1] + + +def test_match_processes_vlm_order_once_each() -> None: + entries = [ + {"heading": "第一章", "level": 1, "page_number": 1}, + {"heading": "第二章", "level": 1, "page_number": 5}, + ] + links = [ + _link("第二章........5", 15), + _link("第一章........1", 11), + ] + + enriched, matched = match_toc_entries_to_links(entries, links) + + assert matched == 2 + assert [e["heading"] for e in enriched] == ["第一章", "第二章"] + assert enriched[0]["link"]["physical_page"] == 11 + assert enriched[1]["link"]["physical_page"] == 15 + + +def test_match_strips_stale_link_when_unmatched() -> None: + entries = [ + { + "heading": "无匹配", + "level": 1, + "page_number": 1, + "link": {"physical_page": 99}, + }, + ] + enriched, matched = match_toc_entries_to_links(entries, [_link("别的标题..1", 2)]) + + assert matched == 0 + assert "link" not in enriched[0] From 759930f0cd3592ef2e9f89a850578526a5ee2f65 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Sat, 15 Aug 2026 12:11:04 +0800 Subject: [PATCH 6/9] refactor: reorganize TOC anchoring imports and streamline code structure - Moved import statements for `anchor_hierarchy`, `finalize_calibration_result`, `pick_primary_offset`, `calibrate_offset`, and `graft_contained_toc` to the top of the file for better readability. - Removed redundant import statements to enhance clarity and maintainability of the `toc_anchoring.py` file. --- .../document_agent/structure/toc_anchoring.py | 29 +++++++++---------- 1 file changed, 14 insertions(+), 15 deletions(-) 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 ba9995261..353280913 100644 --- a/apps/worker/app/services/document_agent/structure/toc_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/toc_anchoring.py @@ -33,6 +33,10 @@ def run_toc_anchoring(ctx: ToolContext) -> None: """Anchor extracted TOC hierarchies onto the profile blackboard.""" + from app.services.document_agent.agents.calibration.orchestrator import ( + anchor_hierarchy, + ) + page_count = int(ctx.blackboard.page_count or 0) hierarchies = list(ctx.blackboard.toc_hierarchies or []) if page_count <= 0 or not hierarchies: @@ -251,21 +255,6 @@ def classify_toc_relationship( return "parallel" -from app.services.document_agent.agents.calibration.orchestrator import ( # noqa: E402 - anchor_hierarchy, -) -from app.services.document_agent.agents.calibration.procedure import ( # noqa: E402 - finalize_calibration_result, - pick_primary_offset, -) -from app.services.document_agent.agents.calibration.service import ( # noqa: E402 - calibrate_offset, -) -from app.services.document_agent.structure.toc_graft import ( # noqa: E402 - graft_contained_toc, -) - - def _anchor_pending_tocs( *, pending_tocs: list[dict[str, Any]], @@ -275,6 +264,14 @@ def _anchor_pending_tocs( body_pages: list[int], primary_ranges: list[ResolvedHierarchyRange], ) -> list[dict[str, Any]]: + from app.services.document_agent.agents.calibration.procedure import ( + finalize_calibration_result, + pick_primary_offset, + ) + from app.services.document_agent.agents.calibration.service import ( + calibrate_offset, + ) + records: list[dict[str, Any]] = [] for i, pending_toc in enumerate(pending_tocs): nodes = extract_toc_nodes([pending_toc]) @@ -351,6 +348,8 @@ def _graft_contained_pending( page_texts: dict[int, str], body_pages: list[int], ) -> tuple[list[TitleNode], SkeletonAnchor]: + from app.services.document_agent.structure.toc_graft import graft_contained_toc + nodes = resolve_nodes overrides = dict(skeleton_anchor.match_overrides) for record in pending_records: From e59c2fb1f4d215a6f8c4b3ca6d0e3bec1b2ddaaa Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Sat, 15 Aug 2026 12:13:51 +0800 Subject: [PATCH 7/9] fix: patch calibration symbols at import sites in TOC graft tests Lazy imports in toc_anchoring broke module-level mocks; point patches at orchestrator/service/procedure instead. Co-authored-by: Cursor --- .../tests/contract/test_toc_graft_contract.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/apps/worker/tests/contract/test_toc_graft_contract.py b/apps/worker/tests/contract/test_toc_graft_contract.py index f355bc265..0c68aa5af 100644 --- a/apps/worker/tests/contract/test_toc_graft_contract.py +++ b/apps/worker/tests/contract/test_toc_graft_contract.py @@ -261,15 +261,15 @@ def fake_anchor_hierarchy(**kwargs): with ( patch( - "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + "app.services.document_agent.agents.calibration.orchestrator.anchor_hierarchy", side_effect=fake_anchor_hierarchy, ), patch( - "app.services.document_agent.structure.toc_anchoring.calibrate_offset", + "app.services.document_agent.agents.calibration.service.calibrate_offset", return_value=object(), ), patch( - "app.services.document_agent.structure.toc_anchoring.pick_primary_offset", + "app.services.document_agent.agents.calibration.procedure.pick_primary_offset", return_value=0, ), patch( @@ -277,7 +277,7 @@ def fake_anchor_hierarchy(**kwargs): return_value="parallel", ), patch( - "app.services.document_agent.structure.toc_anchoring.finalize_calibration_result", + "app.services.document_agent.agents.calibration.procedure.finalize_calibration_result", return_value=([pending], _anchor({("App",): 22}), True), ), ): @@ -319,15 +319,15 @@ def test_profile_grafts_contained_and_keeps_original_pending() -> None: with ( patch( - "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + "app.services.document_agent.agents.calibration.orchestrator.anchor_hierarchy", return_value=([primary], _anchor({("第一章",): 10})), ), patch( - "app.services.document_agent.structure.toc_anchoring.calibrate_offset", + "app.services.document_agent.agents.calibration.service.calibrate_offset", return_value=object(), ), patch( - "app.services.document_agent.structure.toc_anchoring.pick_primary_offset", + "app.services.document_agent.agents.calibration.procedure.pick_primary_offset", return_value=0, ), patch( @@ -335,7 +335,7 @@ def test_profile_grafts_contained_and_keeps_original_pending() -> None: return_value="contained", ), patch( - "app.services.document_agent.structure.toc_anchoring.finalize_calibration_result", + "app.services.document_agent.agents.calibration.procedure.finalize_calibration_result", return_value=( [contained], _anchor({("第一章",): 10, ("第一章", "1.2"): 12}), From e035588f35bd2ee189863fc68afecaa361ad0448 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Sat, 15 Aug 2026 12:22:25 +0800 Subject: [PATCH 8/9] fix: patch TOC anchoring mocks at lazy-import sites Contract tests were still patching symbols on toc_anchoring that are only imported inside functions after the circular-import refactor. Co-authored-by: Cursor --- ...profile_skeleton_anchor_wiring_contract.py | 18 +++++++------- .../test_toc_link_attach_wiring_contract.py | 24 +++++++++++++++++++ 2 files changed, 33 insertions(+), 9 deletions(-) diff --git a/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py b/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py index 096494f52..0abc934c9 100644 --- a/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py +++ b/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py @@ -151,7 +151,7 @@ def fake_anchor_hierarchy(**_kwargs): return [_node()], _anchor() with patch( - "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + "app.services.document_agent.agents.calibration.orchestrator.anchor_hierarchy", side_effect=fake_anchor_hierarchy, ): run_toc_anchoring(ctx) @@ -294,7 +294,7 @@ def fake_finalize(**kwargs): with ( patch( - "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + "app.services.document_agent.agents.calibration.orchestrator.anchor_hierarchy", return_value=([_node()], _anchor()), ), patch( @@ -311,15 +311,15 @@ def fake_finalize(**kwargs): ], ), patch( - "app.services.document_agent.structure.toc_anchoring.calibrate_offset", + "app.services.document_agent.agents.calibration.service.calibrate_offset", return_value=object(), ), patch( - "app.services.document_agent.structure.toc_anchoring.pick_primary_offset", + "app.services.document_agent.agents.calibration.procedure.pick_primary_offset", return_value=0, ), patch( - "app.services.document_agent.structure.toc_anchoring.finalize_calibration_result", + "app.services.document_agent.agents.calibration.procedure.finalize_calibration_result", side_effect=fake_finalize, ), ): @@ -345,7 +345,7 @@ def _boom(*_args, **_kwargs): with ( patch( - "app.services.document_agent.structure.toc_anchoring.anchor_hierarchy", + "app.services.document_agent.agents.calibration.orchestrator.anchor_hierarchy", return_value=([_node()], _anchor()), ), patch( @@ -362,15 +362,15 @@ def _boom(*_args, **_kwargs): ], ), patch( - "app.services.document_agent.structure.toc_anchoring.calibrate_offset", + "app.services.document_agent.agents.calibration.service.calibrate_offset", return_value=object(), ), patch( - "app.services.document_agent.structure.toc_anchoring.pick_primary_offset", + "app.services.document_agent.agents.calibration.procedure.pick_primary_offset", return_value=0, ), patch( - "app.services.document_agent.structure.toc_anchoring.finalize_calibration_result", + "app.services.document_agent.agents.calibration.procedure.finalize_calibration_result", side_effect=_boom, ), ): diff --git a/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py b/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py index ed48f979f..49186dd80 100644 --- a/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py +++ b/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py @@ -71,10 +71,18 @@ def fake_anchor(ctx) -> None: "app.services.document_agent.coordinator.enrich_toc_hierarchies_with_links", side_effect=fake_enrich, ), + patch( + "app.services.document_agent.structure.toc_link_enrichment.enrich_toc_hierarchies_with_links", + side_effect=fake_enrich, + ), patch( "app.services.document_agent.coordinator.run_toc_anchoring", side_effect=fake_anchor, ), + patch( + "app.services.document_agent.structure.toc_anchoring.run_toc_anchoring", + side_effect=fake_anchor, + ), ): coordinator._run_toc_extraction_pipeline() @@ -106,10 +114,18 @@ def fake_anchor(ctx) -> None: "app.services.document_agent.coordinator.enrich_toc_hierarchies_with_links", side_effect=RuntimeError("pymupdf failed"), ), + patch( + "app.services.document_agent.structure.toc_link_enrichment.enrich_toc_hierarchies_with_links", + side_effect=RuntimeError("pymupdf failed"), + ), patch( "app.services.document_agent.coordinator.run_toc_anchoring", side_effect=fake_anchor, ), + patch( + "app.services.document_agent.structure.toc_anchoring.run_toc_anchoring", + side_effect=fake_anchor, + ), ): coordinator._run_toc_extraction_pipeline() @@ -142,10 +158,18 @@ def fake_anchor(_ctx) -> None: "app.services.document_agent.coordinator.enrich_toc_hierarchies_with_links", side_effect=fake_enrich, ), + patch( + "app.services.document_agent.structure.toc_link_enrichment.enrich_toc_hierarchies_with_links", + side_effect=fake_enrich, + ), patch( "app.services.document_agent.coordinator.run_toc_anchoring", side_effect=fake_anchor, ), + patch( + "app.services.document_agent.structure.toc_anchoring.run_toc_anchoring", + side_effect=fake_anchor, + ), ): coordinator._run_toc_extraction_pipeline() From 6895d0e00784be3e91b41fd8d87b3d04705e449c Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Sat, 15 Aug 2026 12:30:36 +0800 Subject: [PATCH 9/9] fix: isolate api/worker pytest and patch via function globals api and worker both use top-level `app`, so a combined pytest run can leave stale module bindings; patch the defining function globals and run the suites separately in CI. Co-authored-by: Cursor --- .github/workflows/pr-ci.yml | 8 ++- .../tests/contract/test_ocr_pages_contract.py | 8 +-- ...profile_skeleton_anchor_wiring_contract.py | 58 ++++++++-------- .../tests/contract/test_toc_graft_contract.py | 14 ++-- .../test_toc_link_attach_wiring_contract.py | 69 +++++++------------ 5 files changed, 73 insertions(+), 84 deletions(-) diff --git a/.github/workflows/pr-ci.yml b/.github/workflows/pr-ci.yml index aeffa9798..e72af0d92 100644 --- a/.github/workflows/pr-ci.yml +++ b/.github/workflows/pr-ci.yml @@ -109,5 +109,9 @@ jobs: - name: Install test system dependencies run: uv run python apps/api/scripts/ensure_test_environment.py --install - - name: Run tests - run: uv run pytest apps/api/tests apps/worker/tests/contract -q + # api/worker both expose top-level `app`; run separately to avoid module shadowing. + - name: Run API tests + run: uv run pytest apps/api/tests -q + + - name: Run worker contract tests + run: uv run pytest apps/worker/tests/contract -q diff --git a/apps/worker/tests/contract/test_ocr_pages_contract.py b/apps/worker/tests/contract/test_ocr_pages_contract.py index 55db1c7ce..d47497283 100644 --- a/apps/worker/tests/contract/test_ocr_pages_contract.py +++ b/apps/worker/tests/contract/test_ocr_pages_contract.py @@ -64,11 +64,11 @@ def __call__(self, _image_path: str): fake_mod = ModuleType("rapidocr_onnxruntime") fake_mod.RapidOCR = lambda: FakeEngine() # type: ignore[attr-defined] + def fake_render(*_args, **_kwargs): + return [{"page": 1, "png_path": "/tmp/ocr_page_1.png"}] + with ( - patch( - "app.services.document_agent.tools.ocr_pages.render_pages", - return_value=[{"page": 1, "png_path": "/tmp/ocr_page_1.png"}], - ), + patch.dict(ocr_pages.__globals__, {"render_pages": fake_render}), patch.dict(sys.modules, {"rapidocr_onnxruntime": fake_mod}), ): result = ocr_pages(ctx, {"pages": [1]}) diff --git a/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py b/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py index 0abc934c9..788836bdd 100644 --- a/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py +++ b/apps/worker/tests/contract/test_profile_skeleton_anchor_wiring_contract.py @@ -297,18 +297,20 @@ def fake_finalize(**kwargs): "app.services.document_agent.agents.calibration.orchestrator.anchor_hierarchy", return_value=([_node()], _anchor()), ), - patch( - "app.services.document_agent.structure.toc_anchoring.resolve_hierarchy_page_ranges", - return_value=[ - ResolvedHierarchyRange( - title="Ch1", - level=1, - start_page=2, - end_page=19, - path_titles=("Ch1",), - match=None, - ) - ], + patch.dict( + run_toc_anchoring.__globals__, + { + "resolve_hierarchy_page_ranges": lambda *_args, **_kwargs: [ + ResolvedHierarchyRange( + title="Ch1", + level=1, + start_page=2, + end_page=19, + path_titles=("Ch1",), + match=None, + ) + ] + }, ), patch( "app.services.document_agent.agents.calibration.service.calibrate_offset", @@ -348,18 +350,20 @@ def _boom(*_args, **_kwargs): "app.services.document_agent.agents.calibration.orchestrator.anchor_hierarchy", return_value=([_node()], _anchor()), ), - patch( - "app.services.document_agent.structure.toc_anchoring.resolve_hierarchy_page_ranges", - return_value=[ - ResolvedHierarchyRange( - title="Ch1", - level=1, - start_page=2, - end_page=19, - path_titles=("Ch1",), - match=None, - ) - ], + patch.dict( + run_toc_anchoring.__globals__, + { + "resolve_hierarchy_page_ranges": lambda *_args, **_kwargs: [ + ResolvedHierarchyRange( + title="Ch1", + level=1, + start_page=2, + end_page=19, + path_titles=("Ch1",), + match=None, + ) + ] + }, ), patch( "app.services.document_agent.agents.calibration.service.calibrate_offset", @@ -427,9 +431,9 @@ def test_c4_uses_persisted_pending_relationship_and_does_not_classify() -> None: def _boom(*_args, **_kwargs): raise AssertionError("C4 must not classify pending TOC") - with patch( - "app.services.document_agent.structure.toc_anchoring.classify_toc_relationship", - side_effect=_boom, + with patch.dict( + run_toc_anchoring.__globals__, + {"classify_toc_relationship": _boom}, ): skeletons = extract_section_skeletons( anatomy=anatomy, diff --git a/apps/worker/tests/contract/test_toc_graft_contract.py b/apps/worker/tests/contract/test_toc_graft_contract.py index 0c68aa5af..055a98e85 100644 --- a/apps/worker/tests/contract/test_toc_graft_contract.py +++ b/apps/worker/tests/contract/test_toc_graft_contract.py @@ -259,6 +259,7 @@ def fake_anchor_hierarchy(**kwargs): captured["body_pages"] = kwargs["body_pages"] return [primary], _anchor({("Ch1",): 2}) + anchoring_globals = run_toc_anchoring.__globals__ with ( patch( "app.services.document_agent.agents.calibration.orchestrator.anchor_hierarchy", @@ -272,9 +273,9 @@ def fake_anchor_hierarchy(**kwargs): "app.services.document_agent.agents.calibration.procedure.pick_primary_offset", return_value=0, ), - patch( - "app.services.document_agent.structure.toc_anchoring.classify_toc_relationship", - return_value="parallel", + patch.dict( + anchoring_globals, + {"classify_toc_relationship": lambda **_kwargs: "parallel"}, ), patch( "app.services.document_agent.agents.calibration.procedure.finalize_calibration_result", @@ -317,6 +318,7 @@ def test_profile_grafts_contained_and_keeps_original_pending() -> None: children=[TitleNode(title="1.2", level=2, printed_page=12)], ) + anchoring_globals = run_toc_anchoring.__globals__ with ( patch( "app.services.document_agent.agents.calibration.orchestrator.anchor_hierarchy", @@ -330,9 +332,9 @@ def test_profile_grafts_contained_and_keeps_original_pending() -> None: "app.services.document_agent.agents.calibration.procedure.pick_primary_offset", return_value=0, ), - patch( - "app.services.document_agent.structure.toc_anchoring.classify_toc_relationship", - return_value="contained", + patch.dict( + anchoring_globals, + {"classify_toc_relationship": lambda **_kwargs: "contained"}, ), patch( "app.services.document_agent.agents.calibration.procedure.finalize_calibration_result", diff --git a/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py b/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py index 49186dd80..6ed378141 100644 --- a/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py +++ b/apps/worker/tests/contract/test_toc_link_attach_wiring_contract.py @@ -37,6 +37,7 @@ def _coordinator() -> ProfileCoordinator: def test_profile_attaches_toc_links_before_anchoring() -> None: coordinator = _coordinator() seen: dict[str, object] = {} + globals_ = ProfileCoordinator._run_toc_extraction_pipeline.__globals__ def fake_enrich(*, pdf_path: str, toc_hierarchies: list[dict[str, object]]): assert pdf_path == "/tmp/doc.pdf" @@ -67,21 +68,12 @@ def fake_anchor(ctx) -> None: "_dispatch_profile_tool", return_value=ToolResult(status="ok"), ), - patch( - "app.services.document_agent.coordinator.enrich_toc_hierarchies_with_links", - side_effect=fake_enrich, - ), - patch( - "app.services.document_agent.structure.toc_link_enrichment.enrich_toc_hierarchies_with_links", - side_effect=fake_enrich, - ), - patch( - "app.services.document_agent.coordinator.run_toc_anchoring", - side_effect=fake_anchor, - ), - patch( - "app.services.document_agent.structure.toc_anchoring.run_toc_anchoring", - side_effect=fake_anchor, + patch.dict( + globals_, + { + "enrich_toc_hierarchies_with_links": fake_enrich, + "run_toc_anchoring": fake_anchor, + }, ), ): coordinator._run_toc_extraction_pipeline() @@ -98,33 +90,28 @@ def fake_anchor(ctx) -> None: def test_link_attach_failure_keeps_hierarchies_and_still_anchors() -> None: coordinator = _coordinator() seen = {"anchored": False} + globals_ = ProfileCoordinator._run_toc_extraction_pipeline.__globals__ def fake_anchor(ctx) -> None: seen["anchored"] = True entry = ctx.blackboard.toc_hierarchies[0]["toc_with_level"][0] assert entry.get("link") is None + def boom_enrich(**_kwargs): + raise RuntimeError("pymupdf failed") + with ( patch.object( ProfileCoordinator, "_dispatch_profile_tool", return_value=ToolResult(status="ok"), ), - patch( - "app.services.document_agent.coordinator.enrich_toc_hierarchies_with_links", - side_effect=RuntimeError("pymupdf failed"), - ), - patch( - "app.services.document_agent.structure.toc_link_enrichment.enrich_toc_hierarchies_with_links", - side_effect=RuntimeError("pymupdf failed"), - ), - patch( - "app.services.document_agent.coordinator.run_toc_anchoring", - side_effect=fake_anchor, - ), - patch( - "app.services.document_agent.structure.toc_anchoring.run_toc_anchoring", - side_effect=fake_anchor, + patch.dict( + globals_, + { + "enrich_toc_hierarchies_with_links": boom_enrich, + "run_toc_anchoring": fake_anchor, + }, ), ): coordinator._run_toc_extraction_pipeline() @@ -137,6 +124,7 @@ def test_skip_toc_anchoring_stops_after_link_attach() -> None: coordinator = _coordinator() coordinator.ctx.settings["skip_toc_anchoring"] = True seen = {"anchored": False} + globals_ = ProfileCoordinator._run_toc_extraction_pipeline.__globals__ def fake_enrich(*, pdf_path: str, toc_hierarchies: list[dict[str, object]]): return toc_hierarchies, SimpleNamespace( @@ -154,21 +142,12 @@ def fake_anchor(_ctx) -> None: "_dispatch_profile_tool", return_value=ToolResult(status="ok"), ), - patch( - "app.services.document_agent.coordinator.enrich_toc_hierarchies_with_links", - side_effect=fake_enrich, - ), - patch( - "app.services.document_agent.structure.toc_link_enrichment.enrich_toc_hierarchies_with_links", - side_effect=fake_enrich, - ), - patch( - "app.services.document_agent.coordinator.run_toc_anchoring", - side_effect=fake_anchor, - ), - patch( - "app.services.document_agent.structure.toc_anchoring.run_toc_anchoring", - side_effect=fake_anchor, + patch.dict( + globals_, + { + "enrich_toc_hierarchies_with_links": fake_enrich, + "run_toc_anchoring": fake_anchor, + }, ), ): coordinator._run_toc_extraction_pipeline()