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 01/11] 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 02/11] 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 03/11] 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 04/11] 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 05/11] 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 06/11] 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 07/11] 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 08/11] 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 09/11] 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() From 9937fc530cd58d5d92190ddf754cd08c0cc45e3a Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 15 Aug 2026 20:15:55 +0800 Subject: [PATCH 10/11] fix(ecs): preserve validated worker cpu capacity --- deploy/ecs/task-definition-worker.staging.json | 2 +- deploy/ecs/test_render_task_definitions.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/deploy/ecs/task-definition-worker.staging.json b/deploy/ecs/task-definition-worker.staging.json index 6caf86e16..f0c2a709d 100644 --- a/deploy/ecs/task-definition-worker.staging.json +++ b/deploy/ecs/task-definition-worker.staging.json @@ -4,7 +4,7 @@ "executionRoleArn": "${EXECUTION_ROLE_ARN}", "networkMode": "awsvpc", "requiresCompatibilities": ["FARGATE"], - "cpu": "1024", + "cpu": "2048", "memory": "4096", "ephemeralStorage": {"sizeInGiB": 21}, "runtimePlatform": { diff --git a/deploy/ecs/test_render_task_definitions.py b/deploy/ecs/test_render_task_definitions.py index bbabbf8e9..72b13457a 100644 --- a/deploy/ecs/test_render_task_definitions.py +++ b/deploy/ecs/test_render_task_definitions.py @@ -126,7 +126,9 @@ def test_staging_worker_preserves_evidence_selected_capacity() -> None: definition_path.read_text(encoding="utf-8") ) - assert definition["cpu"] == "1024" + # The measured production-envelope replay rejected 1-vCPU tasks and passed + # on two fixed 2-vCPU tasks, so CD must not restore the rejected task size. + assert definition["cpu"] == "2048" assert definition["memory"] == "4096" From 87d339c507fdaa8a57c46b71ca13703fa73918a4 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 15 Aug 2026 22:45:22 +0800 Subject: [PATCH 11/11] feat: add manual staging availability controls --- .github/workflows/manage-staging.yml | 257 +++++++++++++++++++++ deploy/ecs/README.md | 38 +++ deploy/ecs/test_manage_staging_workflow.py | 101 ++++++++ 3 files changed, 396 insertions(+) create mode 100644 .github/workflows/manage-staging.yml create mode 100644 deploy/ecs/test_manage_staging_workflow.py diff --git a/.github/workflows/manage-staging.yml b/.github/workflows/manage-staging.yml new file mode 100644 index 000000000..085ff8dd3 --- /dev/null +++ b/.github/workflows/manage-staging.yml @@ -0,0 +1,257 @@ +name: Manage staging ECS availability + +on: + workflow_dispatch: + inputs: + operation: + description: Read status, start staging, or stop staging after the drain window + required: true + default: status + type: choice + options: + - status + - start + - stop + approval_reference: + description: Required for start/stop; use the owner approval issue or comment URL + required: false + type: string + +permissions: + contents: read + +concurrency: + group: staging-ecs-management + cancel-in-progress: false + +jobs: + manage: + name: ${{ inputs.operation }} staging ECS + runs-on: ubuntu-latest + timeout-minutes: 60 + env: + AWS_REGION: us-east-1 + APPROVAL_REFERENCE: ${{ inputs.approval_reference }} + OPERATION: ${{ inputs.operation }} + RUNTIME_SECRET_ARN: ${{ secrets.AWS_ECS_STAGING_SECRETS_ARN }} + steps: + - name: Checkout operator source + uses: actions/checkout@v4 + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ env.AWS_REGION }} + + - name: Create the required named profile + shell: bash + run: | + set -euo pipefail + aws configure set aws_access_key_id "$AWS_ACCESS_KEY_ID" --profile knowhere + aws configure set aws_secret_access_key "$AWS_SECRET_ACCESS_KEY" --profile knowhere + aws configure set region "$AWS_REGION" --profile knowhere + if [ -n "${AWS_SESSION_TOKEN:-}" ]; then + aws configure set aws_session_token "$AWS_SESSION_TOKEN" --profile knowhere + fi + + - name: Ensure PostgreSQL client is available + shell: bash + run: | + set -euo pipefail + if ! command -v psql >/dev/null 2>&1; then + sudo apt-get update + sudo apt-get install --yes postgresql-client + fi + + - name: Run staging management operation + shell: bash + run: | + set -euo pipefail + report_path="$RUNNER_TEMP/staging-management-report.json" + startup_seconds="" + + # The named profile plus the account assertion prevent a manually + # dispatched staging operation from falling through to another AWS + # credential source or account. + unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN + account_id="$(aws --profile knowhere sts get-caller-identity --query Account --output text)" + if [ "$account_id" != "107424103509" ]; then + echo "Unexpected AWS account: $account_id" >&2 + exit 1 + fi + + if [ "$OPERATION" != "status" ]; then + approval_reference_pattern='^https://github\.com/Ontos-AI/[A-Za-z0-9._-]+/(issues/[0-9]+(#issuecomment-[0-9]+)?|pull/[0-9]+(#(issuecomment-[0-9]+|discussion_r[0-9]+))?)$' + if [[ ! "$APPROVAL_REFERENCE" =~ $approval_reference_pattern ]]; then + echo "start and stop require an Ontos-AI GitHub issue, pull request, or comment URL" >&2 + exit 1 + fi + fi + if [ -z "$RUNTIME_SECRET_ARN" ]; then + echo "AWS_ECS_STAGING_SECRETS_ARN is not configured" >&2 + exit 1 + fi + + update_service() { + local service_name="$1" + local desired_count="$2" + aws --profile knowhere ecs update-service \ + --cluster knowhere-fargate \ + --service "$service_name" \ + --desired-count "$desired_count" \ + --query 'service.serviceArn' \ + --output text >/dev/null + } + + wait_for_service() { + local service_name="$1" + aws --profile knowhere ecs wait services-stable \ + --cluster knowhere-fargate \ + --services "$service_name" + } + + read_services() { + aws --profile knowhere ecs describe-services \ + --cluster knowhere-fargate \ + --services knowhere-api-staging knowhere-worker-staging \ + --query 'services[].{name:serviceName,taskDefinition:taskDefinition,desired:desiredCount,running:runningCount,pending:pendingCount,status:status}' \ + --output json + } + + case "$OPERATION" in + status) + ;; + start) + start_started_epoch="$(date +%s)" + # Restore the accepted two-worker capacity before admitting API + # traffic; ECS stability alone does not prove container health. + update_service knowhere-worker-staging 2 + wait_for_service knowhere-worker-staging + mapfile -t worker_tasks < <( + aws --profile knowhere ecs list-tasks \ + --cluster knowhere-fargate \ + --service-name knowhere-worker-staging \ + --desired-status RUNNING \ + --query 'taskArns[]' \ + --output text | tr '\t' '\n' + ) + if [ "${#worker_tasks[@]}" -ne 2 ]; then + echo "Expected two running worker tasks, found ${#worker_tasks[@]}" >&2 + exit 1 + fi + healthy_workers="$(aws --profile knowhere ecs describe-tasks \ + --cluster knowhere-fargate \ + --tasks "${worker_tasks[@]}" \ + --query 'length(tasks[?lastStatus==`RUNNING` && healthStatus==`HEALTHY`])' \ + --output text)" + if [ "$healthy_workers" -ne 2 ]; then + echo "Expected two healthy workers, found $healthy_workers" >&2 + exit 1 + fi + + update_service knowhere-api-staging 1 + wait_for_service knowhere-api-staging + curl --fail --silent --show-error --max-time 30 \ + https://api-staging.knowhereto.ai/health >/dev/null + startup_seconds="$(($(date +%s) - start_started_epoch))" + ;; + stop) + # Close admission first, retain the accepted 30-minute drain, + # then stop workers. Any remainder uses the tested watchdog path. + update_service knowhere-api-staging 0 + wait_for_service knowhere-api-staging + sleep 1800 + update_service knowhere-worker-staging 0 + wait_for_service knowhere-worker-staging + ;; + *) + echo "Unsupported operation: $OPERATION" >&2 + exit 1 + ;; + esac + + services_json="$(read_services)" + + # The hosted runner cannot reach private ElastiCache. The jobs table + # is the permanent processing ledger, so it supplies the durable + # backlog and oldest-active-job signal without another VPC resource. + runtime_secret="$(aws --profile knowhere secretsmanager get-secret-value \ + --secret-id "$RUNTIME_SECRET_ARN" \ + --query SecretString \ + --output text)" + database_url="$(jq -r '.DATABASE_URL // empty' <<<"$runtime_secret")" + unset runtime_secret + if [ -z "$database_url" ]; then + echo "Staging runtime secret has no DATABASE_URL" >&2 + exit 1 + fi + psql_url="$(sed 's#^postgresql+asyncpg://#postgresql://#' <<<"$database_url")" + unset database_url + read -r -d '' ledger_query <<'SQL' || true + WITH backlog AS ( + SELECT status, count(*)::integer AS status_count, min(created_at) AS oldest + FROM jobs + WHERE status NOT IN ('done', 'failed') + GROUP BY status + ) + SELECT json_build_object( + 'total', coalesce(sum(status_count), 0), + 'queued', json_build_object( + 'count', coalesce( + sum(status_count) FILTER ( + WHERE status IN ('pending', 'waiting-file') + ), + 0 + ), + 'oldestCreatedAt', min(oldest) FILTER ( + WHERE status IN ('pending', 'waiting-file') + ) + ), + 'processing', json_build_object( + 'count', coalesce( + sum(status_count) FILTER ( + WHERE status IN ('running', 'converting') + ), + 0 + ), + 'oldestCreatedAt', min(oldest) FILTER ( + WHERE status IN ('running', 'converting') + ) + ), + 'byStatus', coalesce( + json_object_agg(status, status_count) FILTER (WHERE status IS NOT NULL), + '{}'::json + ) + )::text + FROM backlog; + SQL + ledger_json="$(psql "$psql_url" \ + --no-psqlrc \ + --set ON_ERROR_STOP=1 \ + --tuples-only \ + --no-align \ + --command "$ledger_query")" + unset psql_url ledger_query + + jq -n \ + --arg operation "$OPERATION" \ + --arg approvalReference "$APPROVAL_REFERENCE" \ + --arg startupSeconds "$startup_seconds" \ + --argjson services "$services_json" \ + --argjson jobLedger "$ledger_json" \ + '{ + operation: $operation, + approvalReference: (if $approvalReference == "" then null else $approvalReference end), + startupSeconds: (if $startupSeconds == "" then null else ($startupSeconds | tonumber) end), + services: $services, + jobLedger: $jobLedger + }' | tee "$report_path" + + { + echo '## Staging ECS management result' + echo '```json' + cat "$report_path" + echo '```' + } >> "$GITHUB_STEP_SUMMARY" diff --git a/deploy/ecs/README.md b/deploy/ecs/README.md index b02ebd3a2..564611fb3 100644 --- a/deploy/ecs/README.md +++ b/deploy/ecs/README.md @@ -33,6 +33,44 @@ The staging workflow in `.github/workflows/build-images.yml` expects these GitHu Before an ECS staging deployment, an operator must create and verify the ECS services, network configuration, API load-balancer target, CloudWatch log groups, and runtime secret. The workflow validates those resources and fails without registering or updating a service when any prerequisite is missing. It does not create or delete AWS resources. +## Manual staging availability + +`.github/workflows/manage-staging.yml` exposes three manually dispatched +operations. It does not run on a branch, release, cron, or other automatic +trigger, and it does not create or delete AWS resources. + +- `status` verifies AWS account `107424103509`, reports API and worker + desired/running/pending counts and task-definition revisions, and reads the + queued and processing backlog groups from the permanent `jobs` ledger. +- `start` requires an owner-approval issue/comment reference, starts the worker + service at desired count `2`, waits for two healthy workers, then starts the + API at desired count `1`, verifies the public health endpoint, and reports + the elapsed cold-start readiness time. +- `stop` requires an owner-approval issue/comment reference, stops the API, + waits the accepted 30-minute worker drain window, then stops the workers and + verifies both services at desired/running `0/0`. + +The workflow configures and uses only the named AWS profile `knowhere`; the +command verifies the account before any ECS call. It retrieves the runtime +database URL from the existing staging Secrets Manager secret only to query the +job ledger. Secret values and the database URL are never written to workflow +output or the GitHub step summary. The ledger is the durable business backlog; +it is intentionally reported instead of raw Redis queue length because the +hosted GitHub runner has no route to the private ElastiCache endpoint. + +For `start` or `stop`, paste the explicit owner-approval URL from an issue, pull +request, or comment in an `Ontos-AI` GitHub repository into +`approval_reference`. Free-form notes and URLs outside the organization are +rejected. `status` is read-only and does not require a reference. +Do not begin an upload, billing, webhook, deployment-validation, or long-running +worker test close to the scheduled cutoff. Once EventBridge schedules exist, +an extended staging session must disable both stop schedules with an owner and +expiry before starting such a test, then restore them afterward. + +The repository workflow is only the manual operator interface. The four +weekday start/stop schedules and their least-privilege execution role remain +AWS EventBridge Scheduler resources governed by infrastructure issue #30. + ## Local Docker smoke test LocalStack Community does not implement the ECS API used by this deployment, so it cannot validate Fargate orchestration. The runtime can still be checked locally with Docker: diff --git a/deploy/ecs/test_manage_staging_workflow.py b/deploy/ecs/test_manage_staging_workflow.py new file mode 100644 index 000000000..f51c2db98 --- /dev/null +++ b/deploy/ecs/test_manage_staging_workflow.py @@ -0,0 +1,101 @@ +"""Contracts for the manual staging ECS management workflow.""" + +from pathlib import Path + + +WORKFLOW_PATH: Path = ( + Path(__file__).parents[2] / ".github/workflows/manage-staging.yml" +) + + +def _read_workflow() -> str: + return WORKFLOW_PATH.read_text(encoding="utf-8") + + +def test_workflow_exposes_only_manual_status_start_and_stop() -> None: + """Availability changes require a deliberate workflow dispatch.""" + workflow: str = _read_workflow() + + assert "workflow_dispatch:" in workflow + assert " - status\n - start\n - stop" in workflow + assert "schedule:" not in workflow + assert "--approval-reference" not in workflow + assert 'if [ "$OPERATION" != "status" ]' in workflow + + +def test_mutations_require_an_ontos_ai_github_approval_url() -> None: + """A free-form audit note must not be accepted as mutation approval.""" + workflow: str = _read_workflow() + + assert "approval_reference_pattern=" in workflow + assert '[[ ! "$APPROVAL_REFERENCE" =~ $approval_reference_pattern ]]' in workflow + assert "Ontos-AI GitHub issue, pull request, or comment URL" in workflow + + +def test_workflow_uses_named_profile_and_approved_account() -> None: + """Every live AWS call is pinned to the Knowhere account and profile.""" + workflow: str = _read_workflow() + + assert "aws --profile knowhere sts get-caller-identity" in workflow + assert 'if [ "$account_id" != "107424103509" ]' in workflow + assert "aws ecs" not in workflow + assert "aws sts" not in workflow + assert "aws secretsmanager" not in workflow + + +def test_start_restores_healthy_workers_before_api() -> None: + """The public API cannot open before the two-worker floor is healthy.""" + workflow: str = _read_workflow() + start_block: str = workflow.split(" start)", maxsplit=1)[1].split( + " ;;", maxsplit=1 + )[0] + + assert start_block.index("update_service knowhere-worker-staging 2") < ( + start_block.index("wait_for_service knowhere-worker-staging") + ) + assert start_block.index("wait_for_service knowhere-worker-staging") < ( + start_block.index('if [ "$healthy_workers" -ne 2 ]') + ) + assert start_block.index('if [ "$healthy_workers" -ne 2 ]') < ( + start_block.index("update_service knowhere-api-staging 1") + ) + assert start_block.index("update_service knowhere-api-staging 1") < ( + start_block.index("https://api-staging.knowhereto.ai/health") + ) + + +def test_stop_closes_api_then_drains_before_workers() -> None: + """The stop operation preserves the accepted API-first 30-minute drain.""" + workflow: str = _read_workflow() + stop_block: str = workflow.split(" stop)", maxsplit=1)[1].split( + " ;;", maxsplit=1 + )[0] + + assert stop_block.index("update_service knowhere-api-staging 0") < ( + stop_block.index("sleep 1800") + ) + assert stop_block.index("sleep 1800") < stop_block.index( + "update_service knowhere-worker-staging 0" + ) + + +def test_status_uses_the_jobs_ledger_for_backlog() -> None: + """Status reports durable work without requiring private Redis access.""" + workflow: str = _read_workflow() + + assert "FROM jobs" in workflow + assert "WHERE status NOT IN ('done', 'failed')" in workflow + assert "'queued'" in workflow + assert "status IN ('pending', 'waiting-file')" in workflow + assert "'processing'" in workflow + assert "status IN ('running', 'converting')" in workflow + assert "get-secret-value" in workflow + + +def test_start_reports_readiness_time() -> None: + """Operators can compare cold-start readiness with the schedule lead time.""" + workflow: str = _read_workflow() + + assert 'start_started_epoch="$(date +%s)"' in workflow + assert 'startup_seconds="$(($(date +%s) - start_started_epoch))"' in workflow + assert "startupSeconds" in workflow