From 4711bd78b05147296cd5f144f95c496be5aaedf9 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 11 Aug 2026 17:26:14 +0800 Subject: [PATCH 01/12] feat: enhance document agent text probing and retrieval metrics - Added support for tracking invisible text length in document agent features. - Updated the `probe_page_features` and related functions to include metrics for both visible and invisible text. - Improved documentation for clarity on text rendering modes and their impact on text metrics. This change enhances the accuracy of text extraction and analysis in document processing. --- .../tests/contract/test_retrieval_contract.py | 1 + .../app/services/document_agent/manifest.py | 3 + .../tools/probe_page_features.py | 69 +++++++++++++- .../tests/unit/test_probe_visible_text.py | 95 +++++++++++++++++++ .../services/retrieval/nav/nav_control.py | 17 ++-- .../shared/services/retrieval/nav/nav_llm.py | 4 +- 6 files changed, 173 insertions(+), 16 deletions(-) create mode 100644 apps/worker/tests/unit/test_probe_visible_text.py diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 78dffed98..032e1a8a6 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -6,6 +6,7 @@ import pytest from httpx import AsyncClient from pytest import MonkeyPatch +from sqlalchemy.ext.asyncio import AsyncSession from tests.support.contract_database import ContractDatabase diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index d422a55b8..ad9367150 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -21,6 +21,7 @@ class PageFeature: page: int raw_text_length: int + """Visible extractable text length (excludes PDF ``3 Tr`` / invisible ink).""" text_density: float image_coverage: float image_count: int @@ -31,6 +32,8 @@ class PageFeature: height: float has_asset: bool is_blank_like: bool + invisible_text_length: int = 0 + """Text drawn with invisible rendering mode (``3 Tr``); not used by probe gates.""" def to_dict(self) -> dict[str, Any]: return asdict(self) diff --git a/apps/worker/app/services/document_agent/tools/probe_page_features.py b/apps/worker/app/services/document_agent/tools/probe_page_features.py index 2c3de3598..5c1e83553 100644 --- a/apps/worker/app/services/document_agent/tools/probe_page_features.py +++ b/apps/worker/app/services/document_agent/tools/probe_page_features.py @@ -1,4 +1,9 @@ -"""Full-page structural probing: text pass, then optional asset pass.""" +"""Full-page structural probing: text pass, then optional asset pass. + +Text metrics count **visible** extractable characters only. PDF text drawn with +rendering mode ``3 Tr`` (MuPDF texttrace ``type == 3``) is recorded separately as +``invisible_text_length`` and does not feed density / blank / extrema gates. +""" from __future__ import annotations @@ -24,6 +29,62 @@ # Near-full-page sparse stroke frames are treated as borders, not figures. _FIGURE_FULLPAGE_AREA_RATIO = 0.92 _FIGURE_FULLPAGE_MAX_PATHS = 25 +# MuPDF ``page.get_texttrace()`` type matching PDF text rendering mode ``3 Tr``. +_INVISIBLE_TEXT_TRACE_TYPE = 3 + + +def _trace_ucs_char(value: object) -> str: + if isinstance(value, str): + return value + if isinstance(value, int) and value > 0: + try: + return chr(value) + except ValueError: + return "" + return "" + + +def _text_lengths_from_trace(page: Any) -> tuple[int, int]: + """Return ``(visible_len, invisible_len)`` from MuPDF texttrace. + + Invisible spans use texttrace ``type == 3`` (PDF ``3 Tr``: neither fill nor + stroke). Probe metrics must use the visible length only so CAD/OCR ink that + is present in the content stream but not painted cannot dominate extrema. + """ + visible_parts: list[str] = [] + invisible_parts: list[str] = [] + try: + items = page.get_texttrace() or [] + except Exception: + items = [] + for item in items: + chars = item.get("chars") or [] + chunk = "".join( + _trace_ucs_char(ch[0]) + for ch in chars + if isinstance(ch, (list, tuple)) and ch + ) + if not chunk: + continue + if int(item.get("type") or 0) == _INVISIBLE_TEXT_TRACE_TYPE: + invisible_parts.append(chunk) + else: + visible_parts.append(chunk) + return len("".join(visible_parts).strip()), len("".join(invisible_parts).strip()) + + +def _probe_text_lengths(page: Any) -> tuple[int, int]: + """Visible/invisible text lengths for one page. + + Prefers texttrace so ``3 Tr`` ink can be excluded from ``raw_text_length``. + If the trace is empty, falls back to ``get_text()`` as visible-only so + environments without usable texttrace still probe native text. + """ + visible_len, invisible_len = _text_lengths_from_trace(page) + if visible_len == 0 and invisible_len == 0: + fallback = (page.get_text() or "").strip() + return len(fallback), 0 + return visible_len, invisible_len def _rect_area(rect: Any) -> float: @@ -284,12 +345,12 @@ def _probe_visual_assets( def _probe_text_one(page: Any, page_number: int) -> dict[str, Any]: rect = page.rect area = max(_rect_area(rect), 1.0) - text = page.get_text() or "" - raw_text_length = len(text.strip()) + raw_text_length, invisible_text_length = _probe_text_lengths(page) orientation = "landscape" if float(rect.width) > float(rect.height) else "portrait" return { "page": page_number, "raw_text_length": raw_text_length, + "invisible_text_length": invisible_text_length, "text_density": round(raw_text_length / area * 10000, 4), "orientation": orientation, "width": round(float(rect.width), 2), @@ -365,6 +426,7 @@ def probe_page_features(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: height=float(item.get("height") or 0.0), has_asset=False, is_blank_like=bool(item.get("is_blank_like")), + invisible_text_length=int(item.get("invisible_text_length") or 0), ) for item in (result.get("features") or []) ] @@ -427,6 +489,7 @@ def probe_page_assets(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: height=feature.height, has_asset=has_asset, is_blank_like=feature.raw_text_length < 50 and not has_asset, + invisible_text_length=feature.invisible_text_length, ) ) ctx.blackboard.page_features = sorted(updated, key=lambda f: f.page) diff --git a/apps/worker/tests/unit/test_probe_visible_text.py b/apps/worker/tests/unit/test_probe_visible_text.py new file mode 100644 index 000000000..3909fd396 --- /dev/null +++ b/apps/worker/tests/unit/test_probe_visible_text.py @@ -0,0 +1,95 @@ +"""Probe text metrics must ignore PDF invisible ink (``3 Tr``).""" + +from __future__ import annotations + +from pathlib import Path + +import fitz + +from app.services.document_agent.tools.probe_page_features import ( + _probe_text_lengths, + _probe_text_one, +) + +_VISIBLE_SAMPLE = "VisibleBodyTextXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" +_INVISIBLE_SAMPLE = "InvisibleOCRLayer" + + +def _write_mixed_visibility_pdf(path: Path) -> None: + doc = fitz.open() + page = doc.new_page(width=600, height=300) + + invisible = fitz.TextWriter(page.rect) + invisible.append((40, 80), _INVISIBLE_SAMPLE) + invisible.write_text(page, render_mode=3) + + visible = fitz.TextWriter(page.rect) + visible.append((40, 160), _VISIBLE_SAMPLE) + visible.write_text(page, render_mode=0) + + doc.save(path) + doc.close() + + +def _write_invisible_only_pdf(path: Path) -> None: + doc = fitz.open() + page = doc.new_page(width=600, height=300) + writer = fitz.TextWriter(page.rect) + writer.append((40, 80), "OnlyInvisibleInk") + writer.write_text(page, render_mode=3) + doc.save(path) + doc.close() + + +def test_probe_text_lengths_splits_visible_and_invisible(tmp_path: Path) -> None: + pdf_path = tmp_path / "mixed_tr.pdf" + _write_mixed_visibility_pdf(pdf_path) + + doc = fitz.open(pdf_path) + try: + page = doc[0] + extracted = page.get_text() or "" + assert _INVISIBLE_SAMPLE in extracted + assert "VisibleBodyText" in extracted + + visible_len, invisible_len = _probe_text_lengths(page) + finally: + doc.close() + + assert visible_len == len(_VISIBLE_SAMPLE) + assert invisible_len == len(_INVISIBLE_SAMPLE) + assert len(_VISIBLE_SAMPLE) >= 50 + + +def test_probe_text_one_uses_visible_length_only(tmp_path: Path) -> None: + pdf_path = tmp_path / "mixed_tr.pdf" + _write_mixed_visibility_pdf(pdf_path) + + doc = fitz.open(pdf_path) + try: + feature = _probe_text_one(doc[0], 1) + finally: + doc.close() + + assert len(_VISIBLE_SAMPLE) >= 50 + assert feature["raw_text_length"] == len(_VISIBLE_SAMPLE) + assert feature["invisible_text_length"] == len(_INVISIBLE_SAMPLE) + assert feature["is_blank_like"] is False + + +def test_invisible_only_page_counts_as_empty_for_probe(tmp_path: Path) -> None: + pdf_path = tmp_path / "invisible_only.pdf" + _write_invisible_only_pdf(pdf_path) + + doc = fitz.open(pdf_path) + try: + page = doc[0] + assert len((page.get_text() or "").strip()) == len("OnlyInvisibleInk") + feature = _probe_text_one(page, 1) + finally: + doc.close() + + assert feature["raw_text_length"] == 0 + assert feature["invisible_text_length"] == len("OnlyInvisibleInk") + assert feature["text_density"] == 0.0 + assert feature["is_blank_like"] is True diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_control.py b/packages/shared-python/shared/services/retrieval/nav/nav_control.py index 9402df3ad..698fa5874 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_control.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_control.py @@ -307,9 +307,7 @@ def plan_control( try: from .nav_llm import ( # type: ignore nav_chat, - planner_output_max_tokens, resolve_nav_model, - resolve_nav_thinking_mode, ) model = resolve_nav_model( @@ -317,18 +315,15 @@ def plan_control( model_env="NAV_PLANNER_MODEL", fallback_envs=("NAV_LLM_MODEL",), ) - # Control is short JSON (accept/widen/drop); thinking only on plan_query/replan. - max_tokens = planner_output_max_tokens( + # Control is short JSON (accept/widen/drop); never use planner thinking. + max_tokens = max( + 256, int(getattr(config, "planner_llm_max_tokens", 0) or 0) - or int(config.llm_max_tokens or 256) + or int(config.llm_max_tokens or 256), ) timeout_s = float(os.environ.get("NAV_PLANNER_TIMEOUT_SECONDS", "").strip() or "0") if timeout_s <= 0: - timeout_s = ( - 300.0 - if resolve_nav_thinking_mode(role="planner") == "enabled" - else 90.0 - ) + timeout_s = 90.0 cached = nav_chat( purpose=_CONTROL_PURPOSE, model=model, @@ -339,7 +334,7 @@ def plan_control( temperature=float(config.llm_temperature), max_tokens=max_tokens, response_format={"type": "json_object"}, - thinking_role="planner", + thinking_role="action", context="Nav Plan Control", api_key_env="NAV_PLANNER_API_KEY", base_url_env="NAV_PLANNER_BASE_URL", diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py b/packages/shared-python/shared/services/retrieval/nav/nav_llm.py index af75dcf9a..f8ba3187f 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_llm.py @@ -9,10 +9,10 @@ - ``action`` (navigate / harvest / refine / verify / score): always disabled — short JSON under ``llm_max_tokens`` (often 256). -- ``planner`` (plan_query / replan / plan_control): episode-bound +- ``planner`` (plan_query / replan only): episode-bound ``NavConfig.planner_thinking``, else ``NAV_PLANNER_THINKING`` for EXP scripts; unset → disabled. When enabled, callers should use - ``planner_output_max_tokens``. + ``planner_output_max_tokens``. ``plan_control`` uses ``action`` (thinking off). Migration to Knowhere: inject the production callable with ``set_nav_chat_backend`` (wrap ``llm_fn``); leave nav call sites unchanged. From edc924666f9e6ebf87ce3a96d0e0463893d274ce Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 11 Aug 2026 23:09:52 +0800 Subject: [PATCH 02/12] refactor: remove inspect.pages tool and update related logic - Deleted the `inspect.pages` tool from the document agent, streamlining the toolset. - Updated prompts and decision-making logic to reflect the removal of the inspect action. - Adjusted the planner and executor to handle legacy actions appropriately, ensuring compatibility with the new structure. - Enhanced the sampling logic in the planner to include optional random page selection when extrema are not available. This change simplifies the document profiling process and improves overall clarity in tool usage. --- .../document_agent/executor/prompts.py | 11 +- .../document_agent/executor/react_loop.py | 1 - .../document_agent/planner/planner.py | 62 +- .../document_agent/planner/prompts.py | 5 +- .../services/document_agent/tools/__init__.py | 1 - .../document_agent/tools/inspect_pages.py | 115 --- .../tools/propose_shard_plan.py | 55 -- .../app/services/document_agent/visual.py | 1 - .../test_profile_agent_protocol_contract.py | 7 +- .../tests/unit/test_planner_sample_pages.py | 35 + docs/pdf-profile-agent.html | 689 ++++++++++++++++++ 11 files changed, 782 insertions(+), 200 deletions(-) delete mode 100644 apps/worker/app/services/document_agent/tools/inspect_pages.py create mode 100644 apps/worker/tests/unit/test_planner_sample_pages.py create mode 100644 docs/pdf-profile-agent.html diff --git a/apps/worker/app/services/document_agent/executor/prompts.py b/apps/worker/app/services/document_agent/executor/prompts.py index 06fe0458a..a19cec544 100644 --- a/apps/worker/app/services/document_agent/executor/prompts.py +++ b/apps/worker/app/services/document_agent/executor/prompts.py @@ -4,12 +4,11 @@ "You are the executor of a document profiling agent. Decide the next tool " "call from the blackboard facts and available tools. Return strict JSON with " "keys: action (must be tool_call), rationale, tool_name, tool_args. " - "Use inspect.pages when more visual evidence is needed, grep.text when " - "native-PDF text evidence is needed, propose.shard_plan when evidence is " - "sufficient to shard, validate.anatomy_map after a shard plan exists, and " - "the verdict tool to finish: verdict(status=success) only after validation " - "succeeds, or verdict(status=abort, rationale=...) only when the document " - "cannot be profiled. Do not invent other finish actions." + "Use grep.text when native-PDF text evidence is needed, propose.shard_plan " + "when evidence is sufficient to shard, validate.anatomy_map after a shard " + "plan exists, and the verdict tool to finish: verdict(status=success) only " + "after validation succeeds, or verdict(status=abort, rationale=...) only " + "when the document cannot be profiled. Do not invent other finish actions." ) __all__ = ["REFLEXION_INSTRUCTIONS"] diff --git a/apps/worker/app/services/document_agent/executor/react_loop.py b/apps/worker/app/services/document_agent/executor/react_loop.py index 9917ac4fa..1bdfeff45 100644 --- a/apps/worker/app/services/document_agent/executor/react_loop.py +++ b/apps/worker/app/services/document_agent/executor/react_loop.py @@ -50,7 +50,6 @@ def _compact_blackboard(ctx: ToolContext) -> dict[str, Any]: "verdict": ctx.blackboard.verdict.to_dict() if ctx.blackboard.verdict else None, - "visual_inspections": ctx.blackboard.global_signals.get("visual_inspections", [])[-3:], "grep_history": ctx.blackboard.global_signals.get("grep_history", [])[-3:], "budget": ctx.budget.snapshot(), } diff --git a/apps/worker/app/services/document_agent/planner/planner.py b/apps/worker/app/services/document_agent/planner/planner.py index acf656003..daeff47fa 100644 --- a/apps/worker/app/services/document_agent/planner/planner.py +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -5,6 +5,7 @@ import base64 import json import os +import random import time from typing import Any, cast @@ -73,14 +74,20 @@ 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, ) -> list[int]: """Select representative pages for VLM profiling. Strategy (cap 10): - 1. Text extrema first (length/density min+max, ≤4 unique). + 1. Text extrema first (length/density min+max, ≤4 unique), unless the + caller passes an empty list (e.g. all visible text lengths are 0). 2. Fill remaining slots with 2/2/2 stratified samples from front/middle/back of the non-extrema pool. - 3. Hard truncate to ``_COARSE_SAMPLE_CAP``. + 3. Optionally append ``random_extra`` uniform pages from the leftover + pool (used when extrema were skipped because max text length is 0). + 4. Hard truncate to ``_COARSE_SAMPLE_CAP``. Args: page_count: Total number of pages. @@ -89,6 +96,9 @@ def _sample_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). + rng: Optional RNG for deterministic tests. """ if page_count <= 0: return [] @@ -107,10 +117,23 @@ def _sample_pages( + _segment_sample(middle or pool, middle_n) + _segment_sample(back or pool, back_n) ) - ordered = [] + ordered: list[int] = [] for page in extrema + sampled: if page not in ordered: ordered.append(page) + + extra_n = max(int(random_extra), 0) + if extra_n > 0: + leftover = [ + page + for page in range(1, page_count + 1) + if page not in ordered and page not in skip + ] + if leftover: + picker = rng if rng is not None else random.Random() + for page in picker.sample(leftover, k=min(extra_n, len(leftover))): + ordered.append(page) + return ordered[:_COARSE_SAMPLE_CAP] @@ -159,18 +182,12 @@ def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDec next_action = str(data.get("next_action") or "ready_to_shard").strip().lower() # Legacy models may still emit verdict_now; that is not a planner finish # signal — fall through to ready_to_shard so the executor owns success/abort. - if next_action == "verdict_now": + # Legacy inspect_more is ignored the same way (tool removed). + if next_action in {"verdict_now", "inspect_more"}: next_action = "ready_to_shard" tool_name: str | None = None tool_args: dict[str, Any] = {} - if next_action == "inspect_more": - pages = [int(page) for page in (data.get("inspect_pages") or [])] - tool_name = "inspect.pages" - tool_args = { - "pages": pages[:10], - "question": "Clarify the document structure and whether these pages change the profile or sharding strategy.", - } - elif next_action == "grep_text" and not profile.is_scanned: + if next_action == "grep_text" and not profile.is_scanned: query = str(data.get("grep_query") or "").strip() if query: tool_name = "grep.text" @@ -208,10 +225,28 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: if self.ctx.blackboard.toc_result else [] ) + text_max = float( + ( + ((self.ctx.blackboard.doc_stats or {}).get("raw_text_length") or {}).get( + "max" + ) + or {} + ).get("value") + or 0.0 + ) + # All-visible-zero docs: extrema collapse to a meaningless page-1 tie. + # Skip them and draw one uniform random page instead. + if text_max > 0: + extrema_pages = self.ctx.blackboard.extrema_pages + random_extra = 0 + else: + extrema_pages = [] + random_extra = 1 pages = _sample_pages( self.ctx.blackboard.page_count, - self.ctx.blackboard.extrema_pages, + extrema_pages, exclude_pages=toc_pages, + random_extra=random_extra, ) if not model: profile = DocumentProfile( @@ -268,7 +303,6 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: ) ], "available_actions": [ - "inspect.pages", "grep.text", "propose.shard_plan", "validate.anatomy_map", diff --git a/apps/worker/app/services/document_agent/planner/prompts.py b/apps/worker/app/services/document_agent/planner/prompts.py index 2c6c48145..ae4a8c334 100644 --- a/apps/worker/app/services/document_agent/planner/prompts.py +++ b/apps/worker/app/services/document_agent/planner/prompts.py @@ -5,7 +5,7 @@ "TOC/H1 evidence, 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, " - "inspect_pages, grep_query. " + "grep_query. " "category is a concise semantic document type in at most 5 English words. " "routing_category must be one of atlas, scanned, slides, generic. " "Set routing_category=atlas only when pages are primarily drawing/detail " @@ -17,8 +17,7 @@ "footer_y is the highest footer line you observe (smallest y) when any " "footer is present, otherwise null. When both are set, require " "header_y < footer_y. " - "next_action must be one of inspect_more, grep_text, ready_to_shard. " - "Use inspect_more only when extra page screenshots are needed. " + "next_action must be one of grep_text, ready_to_shard. " "Use grep_text only for native PDFs when a global text search would clarify " "structure. Use ready_to_shard when evidence is sufficient to propose shards. " "Do not finish or abort the profile run from next_action; the executor owns " diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index 177ff879f..43b91226a 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -5,7 +5,6 @@ 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 inspect_pages as inspect_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/inspect_pages.py b/apps/worker/app/services/document_agent/tools/inspect_pages.py deleted file mode 100644 index 2254dbe47..000000000 --- a/apps/worker/app/services/document_agent/tools/inspect_pages.py +++ /dev/null @@ -1,115 +0,0 @@ -"""Generic VLM inspection tool for selected PDF pages.""" - -from __future__ import annotations - -import base64 -import json -import os -import time -from typing import Any, cast - -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 -from shared.utils.token_estimate import estimate_tokens - - -@register_tool( - name="inspect.pages", - description="Render arbitrary PDF pages and ask the VLM a custom profiling question.", - parameters={ - "type": "object", - "properties": { - "pages": {"type": "array", "items": {"type": "integer"}}, - "question": {"type": "string"}, - }, - "required": ["pages", "question"], - }, - preconditions=(has_page_features,), -) -def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: - start = time.monotonic() - pages = sorted( - { - int(page) - for page in (args.get("pages") or []) - if 1 <= int(page) <= ctx.blackboard.page_count - } - )[:10] - if not pages: - return ToolResult( - status="error", - error="inspect.pages requires at least one valid page", - latency_ms=int((time.monotonic() - start) * 1000), - ) - question = str(args.get("question") or "Describe the document structure visible on these pages.") - pngs = render_pages(ctx, pages, folder_name="inspect_pages", prefix="inspect") - model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") - prompt = ( - "You are inspecting PDF page screenshots for a document profiling agent. " - "Answer strict JSON with keys: observations, implications, recommended_next_action. " - "observations must be an array of {page, summary, visual_kind}. " - f"Question: {question}" - ) - est = estimate_tokens(prompt) + len(pngs) * 800 - if not model: - payload = {"pages": pages, "pngs": pngs, "note": "No VLM model configured."} - ctx.blackboard.global_signals.setdefault("visual_inspections", []).append(payload) - return ToolResult( - status="ok", - payload=payload, - latency_ms=int((time.monotonic() - start) * 1000), - warnings=["No VLM model configured; returned rendered page paths only."], - ) - stage = "structural_react" - if not ctx.budget.try_reserve("visual", est, stage=stage): - return ToolResult( - status="error", - error="insufficient visual budget", - latency_ms=int((time.monotonic() - start) * 1000), - ) - - content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt}] - for item in pngs: - with open(str(item["png_path"]), "rb") as f: - img_b64 = base64.b64encode(f.read()).decode() - content_parts.append({"type": "text", "text": f"\n--- Page {item['page']} ---"}) - content_parts.append( - {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_b64}"}} - ) - try: - from shared.services.ai.llm_overrides import get_vision_client - - client, model = get_vision_client(requested_model=model) - raw, usage = client.chat_completion_with_usage( - messages=cast(Any, [{"role": "user", "content": content_parts}]), - model=model, - temperature=0.0, - max_tokens=1200, - response_format={"type": "json_object"}, - usage_task="document_agent.inspect_pages", - ) - ctx.budget.commit( - "visual", - actual=usage.get("total_tokens", est), - est=est, - stage=stage, - ) - try: - payload: dict[str, Any] = json.loads(raw) - except json.JSONDecodeError: - payload = {"raw": raw} - if isinstance(payload, dict): - payload.setdefault("pages", pages) - else: - payload = {"result": payload, "pages": pages} - ctx.blackboard.global_signals.setdefault("visual_inspections", []).append(payload) - return ToolResult( - status="ok", - payload=payload, - latency_ms=int((time.monotonic() - start) * 1000), - tokens_used=usage.get("total_tokens", 0), - ) - except Exception: - ctx.budget.refund("visual", est=est, stage=stage) - raise 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 1c2737db7..30962fcdb 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 @@ -367,61 +367,6 @@ def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> return shards -def _build_prompt( - *, - page_count: int, - min_pages: int, - max_pages: int, - doc_stats: dict[str, Any], - page_kind_counts: dict[str, int], - toc_pages: list[int], - leaf_pages: list[int], - profile: dict[str, Any] | None, - visual_evidence: list[dict[str, Any]], - grep_history: list[dict[str, Any]], -) -> str: - payload = { - "page_count": page_count, - "min_pages_per_shard": min_pages, - "max_pages_per_shard": max_pages, - "page_kind_counts": page_kind_counts, - "doc_stats": doc_stats, - "toc_pages": toc_pages, - "leaf_cut_pages": leaf_pages, - "document_profile": profile, - "visual_evidence": visual_evidence[-3:], - "grep_history": grep_history[-3:], - } - return ( - "You are a senior document parsing architect. Decide whether to split a PDF " - "and where to split it using document-scale features and TOC leaf-node evidence.\n" - "Rules:\n" - "- Return strict JSON only.\n" - "- Prefer TOC leaf-node pages as semantic boundaries, cutting at page-1 when possible.\n" - "- Do not blindly split on every leaf node. Consider total page_count, spacing, min/max " - "shard sizes, and over-fragmentation.\n" - "- Prefer fewer, semantically coherent shards over many tiny shards.\n" - "- Keep each cut rationale under 120 characters.\n" - "- Every resulting shard length must be between min_pages_per_shard and " - "max_pages_per_shard, except the final shard may be shorter only when no better " - "valid split exists. Check each segment length exactly before returning.\n" - "- If no split is useful, return enabled=false and cuts=[] even for a long document.\n" - "Output schema:\n" - "{\n" - ' "enabled": boolean,\n' - ' "cuts": [\n' - " {\"cut_after_page\": number, \"anchor_type\": \"h1_boundary\" | " - "\"blank_separator\" | \"forced_max_size\", " - "\"confidence\": number, \"rationale\": string}\n" - " ],\n" - ' "reason": "llm_boundary_decision" | "not_needed" | "too_large",\n' - ' "rationale": string\n' - "}\n" - "Payload:\n" - + json.dumps(payload, ensure_ascii=False) - ) - - def _build_chapter_prompt( *, page_count: int, diff --git a/apps/worker/app/services/document_agent/visual.py b/apps/worker/app/services/document_agent/visual.py index aebd108d3..f25bf70d6 100644 --- a/apps/worker/app/services/document_agent/visual.py +++ b/apps/worker/app/services/document_agent/visual.py @@ -19,7 +19,6 @@ "planner_pages", "page_locate_pages", "toc_pages", - "inspect_pages", "verify_pages", "agent_visuals", } diff --git a/apps/worker/tests/contract/test_profile_agent_protocol_contract.py b/apps/worker/tests/contract/test_profile_agent_protocol_contract.py index c0c0c88df..60a4c07aa 100644 --- a/apps/worker/tests/contract/test_profile_agent_protocol_contract.py +++ b/apps/worker/tests/contract/test_profile_agent_protocol_contract.py @@ -39,7 +39,6 @@ def test_planner_verdict_now_falls_through_to_ready_to_shard() -> None: "header_y": None, "footer_y": None, "next_action": "verdict_now", - "inspect_pages": [], "grep_query": "", } ) @@ -66,7 +65,7 @@ def test_planner_ready_to_shard_proposes_shard_plan() -> None: assert decision.tool_name == "propose.shard_plan" -def test_planner_inspect_more_maps_to_inspect_pages() -> None: +def test_planner_legacy_inspect_more_falls_through_to_ready_to_shard() -> None: raw = json.dumps( { "is_scanned": False, @@ -79,8 +78,8 @@ def test_planner_inspect_more_maps_to_inspect_pages() -> None: } ) _profile, decision = _parse_profile_and_decision(raw) - assert decision.tool_name == "inspect.pages" - assert decision.tool_args["pages"] == [3, 8] + assert decision.tool_name == "propose.shard_plan" + assert decision.tool_args == {} def test_executor_legacy_verdict_now_without_status_becomes_shard() -> None: diff --git a/apps/worker/tests/unit/test_planner_sample_pages.py b/apps/worker/tests/unit/test_planner_sample_pages.py new file mode 100644 index 000000000..31f9eadf9 --- /dev/null +++ b/apps/worker/tests/unit/test_planner_sample_pages.py @@ -0,0 +1,35 @@ +"""Coarse planner page sampling.""" + +from __future__ import annotations + +import random + +from app.services.document_agent.planner.planner import _sample_pages + + +def test_sample_pages_with_extrema_keeps_extrema_first() -> None: + pages = _sample_pages(273, [1, 270], rng=random.Random(0)) + assert pages[0] == 1 + assert 270 in pages + assert len(pages) <= 10 + + +def test_zero_text_skips_extrema_and_adds_one_random() -> None: + """When caller drops extrema (max visible text == 0), add one random page.""" + rng = random.Random(0) + pages = _sample_pages(273, [], random_extra=1, rng=rng) + + # Stratified front/mid/back only (no extrema): first/last of each third. + # pool=[1..273], third=91 → [1,91] + [92,182] + [183,273] + 1 random. + assert pages[:6] == [1, 91, 92, 182, 183, 273] + assert len(pages) == 7 + assert pages[6] not in {1, 91, 92, 182, 183, 273} + assert 1 <= pages[6] <= 273 + + +def test_zero_text_random_extra_is_deterministic_with_rng() -> None: + a = _sample_pages(273, [], random_extra=1, rng=random.Random(7)) + b = _sample_pages(273, [], random_extra=1, rng=random.Random(7)) + c = _sample_pages(273, [], random_extra=1, rng=random.Random(8)) + assert a == b + assert a != c diff --git a/docs/pdf-profile-agent.html b/docs/pdf-profile-agent.html new file mode 100644 index 000000000..fd6f2312b --- /dev/null +++ b/docs/pdf-profile-agent.html @@ -0,0 +1,689 @@ + + + + + + PDF PROFILE 阶段：流程 · 工具 · 门控 + + + +
+
+

Knowhere · Document Agent

+

PDF PROFILE 阶段手册

+

+ 给反复查阅用：PROFILE 主体是协调器固定流水线；真正可自由选用的侦察工具只剩「全文 GREP」； + 「提切分 / 校验 / 裁决」是确定性门控。page_memory 认真做目录，但不进真切分环。 + （已移除无效的 inspect.pages。） +

+
+ 范围：仅 PROFILE + 流程 + 侦察工具 + 门控 +
+
+ + + +
+

1. 一句话架构

+
+ 整体没有「模型随便选工具绕圈」。日常 PROFILE = 固定阶段。真正的多轮环只服务「超大 PDF + chunk 需要真切分」。 +
+ 
固定流水线（几乎每份 PDF）
+  ① 扫文字特征 → ② 统计 → ③（可选）找目录 + VLM 抄目录
+  → ④ VLM 粗分类（类别 / 图册? / 扫描件?）
+  → ⑤ 扫资产图表 → 再刷统计
+  → ⑥ 按「页数 / 是否图册」决定要不要解剖、要不要切分
+
+侦察（可选）
+  全文 GREP  —— 主要出现在超大切分环
+
+门控链（确定性触发）
+  提切分 → 校验 → 裁决 → 落盘
+
+page_memory：强制倾向做目录；跳过真切分，塞「整本一片」占位。
+
+ +
+

2. 三层分类：流程 / 侦察工具 / 门控

+

关键分界是「谁决定触发」，不是「里面有没有 LLM」。

+
+ 流程 · 协调器写死 + 侦察 · 可选选用 + 门控 · 触发确定性 +
+ +
+
+

流程 模型选不了

+

文字探测 · 页种类粗标 · 统计聚合（两次）· VLM 粗分类 · 资产初探 · 目录两步（固定顺序）· 是否进解剖 / 是否开环 · 单片占位 · 组装落盘 · ReAct 调度器本身

+
+
+

侦察工具 仅 GREP

+

全文 GREP（原生文字 PDF）。时机来自粗分类「下一步建议」或环里决策。已移除无效的「再看几页」。

+
+
+

门控 挂成工具壳

+

提切分 → 校验 → 裁决。成功路径被铁律卡住：不能跳过校验直接成功；结束必须走裁决。

+
+
+

目录两步的特殊身份

+

形态上是注册工具，但日常由协调器按死顺序调用（找锚点 → 抄目录），不是模型随手点的侦察工具。

+
+
+ + 
流程阶段
+  文字探测 → 统计 → 目录? → 粗分类 → 资产 → 再统计
+       │
+       ▼
+侦察工具（可选，仅超大切分环）
+  GREP
+       │
+       ▼
+门控链（确定性触发）
+  提切分 → 校验 → 裁决 → 落盘
+
+ +
+

3. 总流程图

+
+ 
+flowchart TD
+  A[PDF 进 PROFILE] --> B[① 全页文字探测]
+  B --> C[② 算统计 + 极值页]
+  C --> D{要做目录?}
+  D -->|page_memory 强制 / chunk 看开关| E[找目录锚点：关键词扫]
+  E --> F[VLM 确认 + 扩窗抄目录]
+  D -->|关| G[目录跳过]
+  F --> H[④ VLM 粗分类]
+  G --> H
+  H --> I[⑤ 全页资产初探]
+  I --> J[再刷统计]
+  J --> K{图册?}
+  K -->|是| Z[结束：跳过解剖]
+  K -->|否| L{页数超限? 约200}
+  L -->|否| M[轻量解剖]
+  L -->|是| N[完整结构解剖]
+  M --> M1{要真切分?}
+  M1 -->|page_memory 否| M2[单片占位 → 落盘]
+  M1 -->|chunk 是| M3[一次性提切分门控]
+  N --> N1{要真切分?}
+  N1 -->|page_memory 否| N2[单片占位 → 落盘]
+  N1 -->|chunk 是| N3[进入多轮环]
+  N3 --> R0[可选侦察: GREP]
+  R0 --> R1[门控: 提切分]
+  R1 --> R2[门控: 校验]
+  R2 --> R3[门控: 裁决]
+  R3 --> P[落盘解剖图]
+  M2 --> P
+  M3 --> P
+  N2 --> P
+
+
+
+ +
+

4. 能力清单与时机

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
能力分类做什么何时实现性质
文字探测流程每页字数、密度、横竖、空白感一进来必跑纯代码
页种类粗标流程横版 / 普通等规则标签文字探测后立刻纯代码
统计聚合流程均值分位、最短最长页；粗分类靠极值页抽样探测后一次；资产后再刷纯代码
找目录锚点流程壳扫「目录 / contents」，滤页眉假命中目录开时固定第一步文本扫描
抄目录流程壳确认真目录页，扩窗抽标题树紧接锚点VLM
粗分类流程语义类别、路由（图册/扫描/幻灯/通用）、页眉页脚带、「下一步建议」粗阶段必跑VLM（规划一步，不是工具）
资产 / 图表初探流程嵌入图、表、矢量块；用页眉页脚带去噪粗分类之后纯代码几何（不是 VLM 读图）
全文 GREP侦察子串 / 正则搜字超大切分环可选；扫描件不可用纯文本搜
提切分方案门控要不要切、刀切在哪见下一节触发规则混合（内部可有 LLM）
校验门控覆盖是否完整、每片是否超限成功收工前硬门纯代码
裁决门控成功 / 中止；唯一出口校验通过后或失败路径控制逻辑
+ +
+ 你常记得的三块都能对上：统计=流程必调；超长切割=页数超限时的门控链（chunk 才真切）；资产图表初探=粗分类后的流程。 +
+
+ +
+

5. 门控触发规则（按触发，不按是否用 LLM）

+ + + + + + + + + + + + + + + + + + + + + +
门控触发规则
提切分 + 轻量解剖：协调器必调一次（或 page_memory 直接塞单片占位，连门都不进）。 + 超大环：没方案时下一步默认就是它；非法动作也会被掰回它。 +
校验 + 硬门：想成功收工但还没校验通过 → 执行器强制改成先校验。 + 确定性模式：有方案后下一步固定是校验。 +
裁决 + 唯一出口。校验通过 → 才能成功收工；预算耗尽 / 失败 → 中止。 + 不是「探完再想想要不要结束」。 +
+ +

确定性模式写死的顺序

+ 
没方案 → 提切分 → 校验 →（过）成功裁决
+                    └（不过）回退单片再校验
+

有模型时环里表面上还能「选工具」，但成功路径仍被两条铁律卡住：不能跳过校验直接成功；结束必须走裁决。

+
+ +
+

6. 三条故事线

+ +

例 A · 40 页普通报告（page_memory / v2）

+ 
扫字 → 统计 → 关键词找目录 → VLM 抄目录
+→ VLM 粗分类（通用）→ 扫资产 → 再统计
+→ 页数没超 → 轻量解剖 → 不真切分（整本一片）
+→ 结束。后面建层级吃的是「目录树」，不吃切分计划。
+
+没有 agent 环。GREP 基本不出现。
+ +

例 B · 500 页法规/招股书，chunk 要喂 MinerU

+ 
前面粗阶段同上（目录看开关）
+→ 超限 + 非图册 → 完整结构
+→ 可选侦察：粗分类说「再搜 Chapter」→ GREP
+→ 提切分 → 校验 → 裁决
+→ 下游按刀口拆 PDF 分别 MinerU
+
+这里才有真正的多轮环；但环的骨架仍是门控链。
+ +

例 C · 图册 / 图纸为主

+ 
扫字 → 统计 →（可选目录）→ VLM 说「这是图册」
+→ 扫资产 → 直接结束，不做解剖 / 切分环
+→ chunk 下游改走图册解析；page_memory 仍走页管线，但不消费切分
+
+ +
+

7. page_memory vs chunk 分叉

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
维度chunk / MinerUpage_memory
目录可选（全局开关，常关）强制倾向开启；常在粗分类前做完
真切分超大时要真实切分计划跳过；单片占位即可
多轮环超大 + 非图册时进入即使超大也不进切分环
图册跳过解剖；下游走图册解析器PROFILE 可标图册并跳过解剖；解析侧不按图册绕道
下游主要吃什么路由类别、切分刀口、可选目录目录层级为主；切分字段基本是占位
页数门槛约 200 页为「超大」切分门另有页记忆硬上限；不消费切分计划
+
+ +
+

8. VLM / GREP / 纯代码对照

+ + + + + + + + + + + + + + + + +
步骤类型
全页文字特征、横竖版标签、统计极值、资产/表/矢量聚类纯代码
目录关键词扫全文、全文 GREP文本扫描 / GREP
粗分类、目录确认与抽取VLM
切分时印刷页→物理页校准（有目录时）常为 VLM
按章决定刀口文本 LLM（失败则纯代码回退）
无目录切分、校验、单片占位、阶段编排纯代码 / 固定协调
环里「下一步选哪个」文本 LLM 决策（仅超大切分路径；成功仍受门控约束）
+
+ +
+

9. 注册表备忘（实现对照）

+

下面只方便对照代码，日常心智模型请用上面的「流程 / 侦察 / 门控」。

+
+ 注册表里现有工具形态：真正自由侦察只剩 GREP；提切分 / 校验 / 裁决是门控；目录两步是流程步骤的工具壳。（inspect.pages 已删除） +
+ + + + + + + + + + + + + + + + +
注册名心智分类谁触发
find.toc_anchor_pages流程壳协调器固定
extract.toc_with_boundaries流程壳协调器固定
grep.text侦察粗分类建议 / 环可选
propose.shard_plan门控确定性：轻量必调 / 环默认
validate.anatomy_map门控成功收工前硬门
verdict门控唯一出口
+

未进注册表、但会跑的流程步骤

+ +

+ 相关代码位置（查阅用）： + apps/worker/app/services/document_parser/profiling/doc_profiler.py、 + apps/worker/app/services/document_agent/coordinator.py、 + apps/worker/app/services/document_agent/registry.py、 + apps/worker/app/services/document_agent/executor/react_loop.py、 + apps/worker/app/services/document_agent/tools/ +

+
+ + +
+ + + + From d8c2d6cec35045629c446b9fdc731ffe7ab8da8b Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 12 Aug 2026 13:12:45 +0800 Subject: [PATCH 03/12] refactor: update structure anchoring and TOC handling in document agent - Replaced the offset calibration method in `propose_shard_plan` with a new implementation from `structure_anchoring`. - Consolidated TOC node extraction and hierarchy handling by integrating new functions for anchoring and offset calibration. - Removed obsolete pruning logic for out-of-scope nodes, streamlining the skeleton extraction process. - Enhanced the handling of TOC ranges and improved the overall clarity of the document agent's structure. These changes improve the maintainability and efficiency of the document agent's processing logic. --- .../document_agent/agents/__init__.py | 1 + .../agents/calibration/SKILL.md | 84 ++ .../agents/calibration/__init__.py | 21 + .../document_agent/agents/calibration/loop.py | 456 ++++++++++ .../agents/calibration/procedure.py | 305 +++++++ .../agents/calibration/service.py | 68 ++ .../agents/calibration/tools.py | 140 +++ .../agents/calibration/types.py | 151 ++++ .../structure/structure_anchoring.py | 825 ++++++++++++++++++ .../structure/toc_link_enrichment.py | 310 +++++++ .../document_agent/tools/inspect_pages.py | 163 ++++ .../tools/propose_shard_plan.py | 6 +- .../page_memory/skeleton_extractor.py | 767 +--------------- .../test_structure_anchoring_contract.py | 200 +++++ 14 files changed, 2765 insertions(+), 732 deletions(-) create mode 100644 apps/worker/app/services/document_agent/agents/__init__.py create mode 100644 apps/worker/app/services/document_agent/agents/calibration/SKILL.md create mode 100644 apps/worker/app/services/document_agent/agents/calibration/__init__.py create mode 100644 apps/worker/app/services/document_agent/agents/calibration/loop.py create mode 100644 apps/worker/app/services/document_agent/agents/calibration/procedure.py create mode 100644 apps/worker/app/services/document_agent/agents/calibration/service.py create mode 100644 apps/worker/app/services/document_agent/agents/calibration/tools.py create mode 100644 apps/worker/app/services/document_agent/agents/calibration/types.py create mode 100644 apps/worker/app/services/document_agent/structure/structure_anchoring.py create mode 100644 apps/worker/app/services/document_agent/structure/toc_link_enrichment.py create mode 100644 apps/worker/app/services/document_agent/tools/inspect_pages.py create mode 100644 apps/worker/tests/contract/test_structure_anchoring_contract.py diff --git a/apps/worker/app/services/document_agent/agents/__init__.py b/apps/worker/app/services/document_agent/agents/__init__.py new file mode 100644 index 000000000..9cd87f956 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/__init__.py @@ -0,0 +1 @@ +"""Calibration agents package.""" diff --git a/apps/worker/app/services/document_agent/agents/calibration/SKILL.md b/apps/worker/app/services/document_agent/agents/calibration/SKILL.md new file mode 100644 index 000000000..e1629b004 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/SKILL.md @@ -0,0 +1,84 @@ +# Calibration SubAgent Skill + +## Goal + +For the **current TOC region**, discover page-numbering **regimes** and an +**initial offset** for each regime that has usable entries. Submit candidate +offsets via `calibration.submit`. After submit, a deterministic completion pass +runs the production tail-verify → binary-search → small-step recalibrate loop +(using the same visual page confirmer as production). Only **complete segments** +are usable for coarse structure; unrecognized pages are treated as **no TOC**. + +## Do not use + +- Do not scan a fixed window after the TOC (the old “TOC end + N pages” probe). +- Do not invent physical pages you did not inspect or obtain from `link`. + +## Mandatory first step — partition regimes + +Inspect every `page_number` label on the current TOC entries and partition them +into **page-numbering regimes** (distinct numbering systems / label shapes: +decimal digits, roman numerals, prefixed folio labels, etc.). + +- Do not mix samples across regimes when computing an offset. +- Include `entry_indices` (0-based indices into `toc_region.entries`) for each + regime you submit. +- Run the same initial-calibration procedure independently for each regime that + has usable entries. + +## Phase 1 — Initial offset (your job via tools) + +For each regime: + +1. Select a small set of entries (prefer spread: early / middle / late when + enough entries exist). +2. Candidate physical page: + - If the entry has `link.physical_page`, use it as the primary candidate. + - Otherwise derive a coarse physical candidate from the printed label and + `page_count`, then confirm with vision. +3. Call `inspect.pages` to confirm the heading starts on that page and to read + the folio/printed label when useful. +4. If wrong, inspect nearby physical pages and revise. +5. Compute `offset = physical - printed` using this regime’s interpretation of + the printed label. +6. Submit **candidate** offsets. Do not treat Phase 1 alone as a finished + coarse-structure calibration. + +## Phase 2 — Completion (deterministic after submit; production path) + +For each TOC region with a candidate primary offset (prefer decimal): + +1. Build TitleNodes via production `extract_toc_nodes` (integer `page_number` + only; roman / prefixed labels become `printed_page=None`). +2. Run production `anchor_hierarchy_from_offset`: + prune → tail verify → binary-search breakpoint → small-step recalibrate → + null-page parent locate. +3. Emit production `SkeletonAnchor` (`offset`, `offset_status`, + `match_overrides`, `null_page_report`, `bulk_count`, `pruned_count`, + `locate_agent`). +4. On recalibrate/budget failure: keep the complete **prefix**; mark only the + unresolved **suffix** as no TOC. Never fall back to a fixed post-TOC window. + +## Usability bar + +- Coarse structure may use the result when `SkeletonAnchor.offset_status=ok` + and `bulk_count > 0` (at least one complete production segment). +- Otherwise downstream treats the document as no-TOC / Root fallback. + +## Tools + +- `inspect.pages`: primary tool for Phase 1. Open physical pages, render, answer + your question. Prefer batching related pages when the same question applies. +- `calibration.submit`: finish Phase 1. Pass the full result under + `tool_args.result` (or result fields directly in `tool_args`). + +## Output rules + +- Submit `status`, `regimes`, top-level `offset` / `offset_status` for the + primary decimal-digit regime when identifiable, `tool_calls`, `notes`. +- Each regime must include `kind`, candidate `offset`, `offset_status`, + `entry_indices`, `samples` (with `title`, `printed_label`, `physical` when + known), and `posterior` if you already inspected a late check. +- Keep `kind` values consistent within one run (`decimal`, `roman`, `prefixed`, + or `other`). +- Stay within the tool/round budget announced in the payload. diff --git a/apps/worker/app/services/document_agent/agents/calibration/__init__.py b/apps/worker/app/services/document_agent/agents/calibration/__init__.py new file mode 100644 index 000000000..5c7f30b9f --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/__init__.py @@ -0,0 +1,21 @@ +"""Calibration SubAgent package.""" + +from app.services.document_agent.agents.calibration.loop import ( + run_calibration_agent, + run_calibration_for_all_regions, +) +from app.services.document_agent.agents.calibration.procedure import ( + build_calibration_payload, + finalize_calibration_result, +) +from app.services.document_agent.agents.calibration.service import calibrate_offset +from app.services.document_agent.agents.calibration.types import CalibrationResult + +__all__ = [ + "CalibrationResult", + "build_calibration_payload", + "calibrate_offset", + "finalize_calibration_result", + "run_calibration_agent", + "run_calibration_for_all_regions", +] diff --git a/apps/worker/app/services/document_agent/agents/calibration/loop.py b/apps/worker/app/services/document_agent/agents/calibration/loop.py new file mode 100644 index 000000000..5b3a99af9 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/loop.py @@ -0,0 +1,456 @@ +"""ReAct loop for the calibration SubAgent.""" + +from __future__ import annotations + +import json +import os +import time +from pathlib import Path +from typing import Any + +from loguru import logger + +from app.services.document_agent.agents.calibration.tools import ( + build_calibration_registry, + strip_toc_links, +) +from app.services.document_agent.agents.calibration.procedure import ( + build_calibration_payload, + finalize_calibration_result, +) +from app.services.document_agent.agents.calibration.types import ( + CalibrationResult, + calibration_result_from_dict, +) +from app.services.document_agent.budget import BudgetTracker +from app.services.document_agent.manifest import ToolContext, ToolResult +from app.services.document_agent.state import AgentBlackboard +from app.services.document_agent.structure.structure_anchoring import ( + deserialize_skeleton_anchor, + serialize_skeleton_anchor, +) +from shared.utils.token_estimate import estimate_tokens + +_SKILL_PATH = Path(__file__).resolve().parent / "SKILL.md" + +_DECISION_INSTRUCTIONS = """ +You are the calibration SubAgent. Follow the Skill strictly. +Each turn return a JSON object with keys: + action: "tool_call" + rationale: string + tool_name: one of the available tools + tool_args: object +Your job is Phase 1 only: partition regimes and find candidate offsets via +inspect.pages, then call calibration.submit. +Phase 2 (tail verify, binary search, small-step recalibrate) runs automatically +after submit. Do not use a fixed post-TOC page window. +Include the word json in your response. +""".strip() + + +def _load_skill() -> str: + return _SKILL_PATH.read_text(encoding="utf-8") + + +def _parse_decision(raw: str) -> dict[str, Any]: + data = json.loads(raw) + if not isinstance(data, dict): + return {"tool_name": None, "tool_args": {}, "rationale": "invalid decision"} + tool_name = data.get("tool_name") or data.get("name") or data.get("tool") + tool_args = data.get("tool_args") or data.get("arguments") or data.get("args") or {} + if not isinstance(tool_args, dict): + tool_args = {} + return { + "tool_name": tool_name, + "tool_args": dict(tool_args), + "rationale": str(data.get("rationale") or ""), + } + + +def _attach_history(result: CalibrationResult, history: list[dict[str, Any]]) -> CalibrationResult: + result.history_tail = history[-12:] + return result + + +def _toc_region_payload( + hierarchies: list[dict[str, Any]], + region_index: int, +) -> dict[str, Any]: + if region_index < 0 or region_index >= len(hierarchies): + raise IndexError(f"region_index out of range: {region_index}") + region = hierarchies[region_index] + entries = region.get("toc_with_level") if isinstance(region, dict) else None + return { + "region_index": region_index, + "toc_range": region.get("toc_range") if isinstance(region, dict) else None, + "entries": entries if isinstance(entries, list) else [], + } + + +def run_calibration_phase1( + *, + ctx: ToolContext, + toc_hierarchies: list[dict[str, Any]], + region_index: int = 0, + page_count: int | None = None, + no_links: bool = False, + max_rounds: int = 16, + inspect_page_cap: int = 5, + inspect_page_budget: int = 24, +) -> CalibrationResult: + """Agent Phase-1 only: partition regimes + candidate offsets, then submit. + + Reuses the caller's ``ToolContext`` (budget / pdf / settings). Does **not** + run production Phase-2 bulk anchoring. + """ + hierarchies = list(toc_hierarchies or []) + if no_links: + hierarchies = strip_toc_links(hierarchies) + if not hierarchies: + return CalibrationResult(status="failed", notes="toc_hierarchies empty") + region_payload = _toc_region_payload(hierarchies, region_index) + resolved_page_count = int( + page_count or ctx.blackboard.page_count or 0 + ) + if resolved_page_count: + ctx.blackboard.page_count = resolved_page_count + + ctx.settings.setdefault("inspect_page_cap", inspect_page_cap) + ctx.settings.setdefault("inspect_page_budget", inspect_page_budget) + + blackboard = ctx.blackboard + blackboard.global_signals["calibration_region_index"] = region_index + blackboard.global_signals["calibration_tool_calls"] = 0 + blackboard.global_signals["calibration_inspect_pages_used"] = int( + blackboard.global_signals.get("calibration_inspect_pages_used") or 0 + ) + blackboard.global_signals["calibration_done"] = False + blackboard.global_signals.pop("calibration_result", None) + + registry = build_calibration_registry() + skill = _load_skill() + history: list[dict[str, Any]] = [] + + for round_index in range(max_rounds): + available = registry.openai_specs(blackboard) + payload = { + "skill": skill, + "page_count": resolved_page_count, + "no_links": no_links, + "budgets": { + "max_rounds": max_rounds, + "round_index": round_index, + "rounds_remaining": max_rounds - round_index, + "inspect_page_cap_per_call": int( + ctx.settings.get("inspect_page_cap") or inspect_page_cap + ), + "inspect_page_budget_total": int( + ctx.settings.get("inspect_page_budget") or inspect_page_budget + ), + "inspect_pages_used": blackboard.global_signals.get( + "calibration_inspect_pages_used" + ), + }, + "toc_region": region_payload, + "history_tail": history[-8:], + "available_tools": available, + } + prompt = _DECISION_INSTRUCTIONS + "\nPayload:\n" + json.dumps( + payload, ensure_ascii=False + ) + model = ctx.settings.get("model") or ctx.settings.get("vlm_model") + if not model: + return _attach_history( + CalibrationResult( + status="failed", + notes="planner model missing", + region_index=region_index, + ), + history, + ) + + est = estimate_tokens(prompt) + if not ctx.budget.try_reserve("plan", est): + return _attach_history( + CalibrationResult( + status="failed", + notes="planner budget exhausted", + region_index=region_index, + tool_calls=int( + blackboard.global_signals.get("calibration_tool_calls") or 0 + ), + ), + history, + ) + + try: + from shared.services.ai.llm_overrides import get_text_client + + client, model = get_text_client(requested_model=str(model)) + raw, usage = client.chat_completion_with_usage( + messages=[{"role": "user", "content": prompt}], + model=model, + temperature=0.0, + max_tokens=2500, + response_format={"type": "json_object"}, + usage_task="calibration.react_loop", + ) + ctx.budget.commit("plan", actual=usage.get("total_tokens", est), est=est) + decision = _parse_decision(raw) + except Exception as exc: + ctx.budget.refund("plan", est=est) + logger.warning("[calibration] decision failed round={}: {}", round_index, exc) + return _attach_history( + CalibrationResult( + status="failed", + notes=f"decision failed: {exc}", + region_index=region_index, + tool_calls=int( + blackboard.global_signals.get("calibration_tool_calls") or 0 + ), + ), + history, + ) + + tool_name = str(decision.get("tool_name") or "").strip() + tool_args = dict(decision.get("tool_args") or {}) + if not tool_name: + history.append( + { + "round": round_index, + "error": "missing tool_name", + "decision": decision, + } + ) + continue + + tool_result: ToolResult = registry.dispatch(tool_name, ctx, tool_args) + blackboard.global_signals["calibration_tool_calls"] = ( + int(blackboard.global_signals.get("calibration_tool_calls") or 0) + 1 + ) + history.append( + { + "round": round_index, + "rationale": decision.get("rationale"), + "tool_name": tool_name, + "tool_args": tool_args, + "tool_status": tool_result.status, + "tool_payload": tool_result.output_summary + if tool_result.status == "ok" + else tool_result.payload, + "tool_error": tool_result.error, + } + ) + logger.info( + "[calibration] region={} round={} tool={} status={}", + region_index, + round_index, + tool_name, + tool_result.status, + ) + + if blackboard.global_signals.get("calibration_done"): + raw_result = blackboard.global_signals.get("calibration_result") or {} + if isinstance(raw_result, dict): + parsed = calibration_result_from_dict(raw_result) + parsed.region_index = region_index + parsed.tool_calls = int( + blackboard.global_signals.get("calibration_tool_calls") or 0 + ) + return _attach_history(parsed, history) + + return _attach_history( + CalibrationResult( + status="failed", + notes="max rounds reached without calibration.submit", + region_index=region_index, + tool_calls=int(blackboard.global_signals.get("calibration_tool_calls") or 0), + ), + history, + ) + + +def run_calibration_agent( + *, + pdf_path: str, + page_count: int, + toc_hierarchies: list[dict[str, Any]], + region_index: int = 0, + output_dir: str, + vlm_model: str | None = None, + planner_model: str | None = None, + no_links: bool = False, + max_rounds: int = 16, + inspect_page_cap: int = 5, + inspect_page_budget: int = 24, + budget: BudgetTracker | None = None, + page_texts: dict[int, str] | None = None, + body_pages: list[int] | None = None, +) -> tuple[CalibrationResult, dict[str, Any]]: + """Debug/full path: Phase-1 agent + production Phase-2 finalize.""" + hierarchies = list(toc_hierarchies or []) + if no_links: + hierarchies = strip_toc_links(hierarchies) + region_hierarchies = [hierarchies[region_index]] if hierarchies else [] + region_payload = _toc_region_payload(hierarchies, region_index) if hierarchies else { + "entries": [] + } + + blackboard = AgentBlackboard() + blackboard.page_count = page_count + if page_texts: + blackboard.page_full_text_cache = dict(page_texts) + + ctx = ToolContext( + pdf_path=pdf_path, + job_id=f"calibration-region-{region_index}", + blackboard=blackboard, + budget=budget + or BudgetTracker( + plan_budget=int(os.environ.get("PARSE_AGENT_PLAN_BUDGET", "50000")), + visual_budget=int(os.environ.get("PARSE_AGENT_VISUAL_BUDGET", "80000")), + ), + trace=None, + output_dir=output_dir, + settings={ + "vlm_model": vlm_model or "", + "model": planner_model or vlm_model or "", + "inspect_page_cap": inspect_page_cap, + "inspect_page_budget": inspect_page_budget, + }, + ) + + phase1 = run_calibration_phase1( + ctx=ctx, + toc_hierarchies=hierarchies, + region_index=region_index, + page_count=page_count, + no_links=False, # already stripped above when requested + max_rounds=max_rounds, + inspect_page_cap=inspect_page_cap, + inspect_page_budget=inspect_page_budget, + ) + if phase1.status == "failed" and not phase1.regimes: + return phase1, {} + + anchor, finalized = finalize_calibration_result( + result=phase1, + entries=list(region_payload.get("entries") or []), + toc_hierarchies=region_hierarchies, + ctx=ctx, + page_count=page_count, + page_texts=page_texts, + body_pages=body_pages, + ) + finalized.history_tail = list(phase1.history_tail) + return finalized, serialize_skeleton_anchor(anchor) + + +def run_calibration_for_all_regions( + *, + pdf_path: str, + page_count: int, + toc_hierarchies: list[dict[str, Any]], + output_dir: str, + vlm_model: str | None = None, + planner_model: str | None = None, + no_links: bool = False, + max_rounds: int = 16, + budget: BudgetTracker | None = None, + page_texts: dict[int, str] | None = None, + body_pages: list[int] | None = None, +) -> dict[str, Any]: + """Calibrate each TOC region; return production SkeletonAnchor-shaped payload.""" + hierarchies = list(toc_hierarchies or []) + if not hierarchies: + return { + "offset": None, + "offset_status": "failed", + "match_overrides": {}, + "null_page_report": [], + "bulk_count": 0, + "pruned_count": 0, + "locate_agent": "offset_only", + "status": "failed", + "regimes": [], + "regions": [], + "notes": "toc_hierarchies empty", + "tool_calls": 0, + "no_links": no_links, + } + + region_results: list[dict[str, Any]] = [] + all_regimes: list[dict[str, Any]] = [] + tool_calls = 0 + primary_anchor: dict[str, Any] | None = None + primary_result: CalibrationResult | None = None + + for idx in range(len(hierarchies)): + t0 = time.time() + result, anchor_dict = run_calibration_agent( + pdf_path=pdf_path, + page_count=page_count, + toc_hierarchies=hierarchies, + region_index=idx, + output_dir=output_dir, + vlm_model=vlm_model, + planner_model=planner_model, + no_links=no_links, + max_rounds=max_rounds, + budget=budget, + page_texts=page_texts, + body_pages=body_pages, + ) + payload = result.to_dict() + payload["elapsed_s"] = round(time.time() - t0, 2) + payload["skeleton_anchor"] = anchor_dict + region_results.append(payload) + tool_calls += int(result.tool_calls or 0) + for regime in payload.get("regimes") or []: + if isinstance(regime, dict): + tagged = dict(regime) + tagged["region_index"] = idx + all_regimes.append(tagged) + if primary_anchor is None and anchor_dict.get("offset") is not None: + primary_anchor = anchor_dict + primary_result = result + + if primary_anchor is None: + primary_anchor = { + "offset": None, + "offset_status": "failed", + "match_overrides": {}, + "null_page_report": [], + "bulk_count": 0, + "pruned_count": 0, + "locate_agent": "offset_only", + } + if primary_result is None: + primary_result = CalibrationResult( + status="failed", notes="no region produced offset" + ) + + anchor = deserialize_skeleton_anchor(primary_anchor) + status = ( + "ok" + if anchor.offset is not None and int(anchor.bulk_count or 0) > 0 + else "failed" + ) + merged = build_calibration_payload( + anchor=anchor, + result=CalibrationResult( + status=status, + regimes=[], + offset=anchor.offset, + offset_status=anchor.offset_status, + tool_calls=tool_calls, + notes=primary_result.notes, + ), + no_links=no_links, + region_payloads=region_results, + tool_calls=tool_calls, + ) + merged["status"] = status + merged["regimes"] = all_regimes + merged["regions"] = region_results + return merged diff --git a/apps/worker/app/services/document_agent/agents/calibration/procedure.py b/apps/worker/app/services/document_agent/agents/calibration/procedure.py new file mode 100644 index 000000000..bcd8fb096 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/procedure.py @@ -0,0 +1,305 @@ +"""Phase-2 completion aligned with production anchoring. + +After the agent submits candidate regime offsets, this module: +1. Picks the primary (usually decimal) candidate offset +2. Builds TitleNodes the same way production does (``extract_toc_nodes``) +3. Runs ``anchor_hierarchy_from_offset`` (prune → bulk/bisect → null-page) + +The returned ``SkeletonAnchor`` is the production schema swap point. +Regime metadata is retained only as experiment diagnostics. +""" + +from __future__ import annotations + +import re +from typing import Any + +from loguru import logger + +from app.services.document_agent.agents.calibration.types import ( + CalibrationRegime, + CalibrationResult, + CalibrationSegment, +) +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.structure.hierarchy_locator import ( + TitleMatch, + extract_toc_nodes, + iter_leaf_title_nodes, +) +from app.services.document_agent.structure.structure_anchoring import ( + SkeletonAnchor, + anchor_hierarchy_from_offset, + serialize_skeleton_anchor, +) + + +_ROMAN_MAP = {"i": 1, "v": 5, "x": 10, "l": 50, "c": 100, "d": 500, "m": 1000} + + +def classify_page_number_kind(label: Any) -> str: + text = str(label or "").strip() + if not text: + return "other" + if re.fullmatch(r"\d+", text): + return "decimal" + if re.fullmatch(r"[ivxlcdm]+", text, flags=re.IGNORECASE): + return "roman" + if re.fullmatch(r"[A-Za-z]+-\d+", text): + return "prefixed" + return "other" + + +def parse_printed_page(label: Any, *, kind: str) -> int | None: + text = str(label or "").strip() + if not text: + return None + kind_l = (kind or "").lower() + if kind_l in {"decimal", "arabic", "arabic_digits"}: + return int(text) if text.isdigit() else None + if kind_l == "roman": + return _roman_to_int(text) + if kind_l in {"prefixed", "folio"}: + match = re.fullmatch(r"[A-Za-z]+-(\d+)", text) + return int(match.group(1)) if match else None + if text.isdigit(): + return int(text) + if re.fullmatch(r"[ivxlcdm]+", text, flags=re.IGNORECASE): + return _roman_to_int(text) + match = re.fullmatch(r"[A-Za-z]+-(\d+)", text) + return int(match.group(1)) if match else None + + +def _roman_to_int(text: str) -> int | None: + raw = text.strip().lower() + if not raw or not re.fullmatch(r"[ivxlcdm]+", raw): + return None + total = 0 + prev = 0 + for ch in reversed(raw): + value = _ROMAN_MAP.get(ch) + if value is None: + return None + if value < prev: + total -= value + else: + total += value + prev = value + return total if total > 0 else None + + +def normalize_kind(kind: str) -> str: + text = (kind or "other").strip().lower() + if text in {"arabic", "arabic_digits", "decimal"}: + return "decimal" + if text == "roman": + return "roman" + if text in {"prefixed", "folio"}: + return "prefixed" + return text or "other" + + +def pick_primary_offset(result: CalibrationResult) -> int | None: + """Prefer decimal-regime candidate offset; else first regime with an offset.""" + for regime in result.regimes: + if normalize_kind(regime.kind) == "decimal" and regime.offset is not None: + return int(regime.offset) + for regime in result.regimes: + if regime.offset is not None: + return int(regime.offset) + if result.offset is not None: + return int(result.offset) + return None + + +def _seed_overrides_from_samples( + *, + result: CalibrationResult, + nodes: list[Any], +) -> dict[tuple[str, ...], TitleMatch]: + """Map Phase-1 confirmed samples onto leaf paths when titles match.""" + title_to_path: dict[str, tuple[str, ...]] = {} + for path, node in iter_leaf_title_nodes(nodes): + title_to_path[node.title] = path + + overrides: dict[tuple[str, ...], TitleMatch] = {} + for regime in result.regimes: + for sample in regime.samples: + if sample.physical is None or not sample.title: + continue + path = title_to_path.get(sample.title.strip()) + if path is None: + # Soft match: normalized equality + needle = sample.title.strip().lower() + for title, candidate in title_to_path.items(): + if title.lower() == needle: + path = candidate + break + if path is None: + continue + overrides[path] = TitleMatch( + page=int(sample.physical), + confidence=0.85, + source="agent_vlm", + matched_line="", + score=0.85, + candidates=[int(sample.physical)], + evidence={ + "calibration": True, + "printed_label": sample.printed_label, + "method": sample.method or "agent_phase1", + "regime_kind": regime.kind, + }, + ) + return overrides + + +def _annotate_regimes_from_anchor( + *, + result: CalibrationResult, + anchor: SkeletonAnchor, + entries: list[dict[str, Any]], +) -> list[CalibrationRegime]: + """Attach production segment view onto agent regimes for diagnostics.""" + out: list[CalibrationRegime] = [] + for regime in result.regimes: + kind = normalize_kind(regime.kind) + indices = list(regime.entry_indices or []) + if not indices: + indices = [ + idx + for idx, entry in enumerate(entries) + if isinstance(entry, dict) + and classify_page_number_kind(entry.get("page_number")) == kind + ] + + # Production trees only integer-print leaves enter bulk; decimal regime + # maps directly onto SkeletonAnchor bulk when Phase-2 succeeded. + if ( + kind == "decimal" + and anchor.offset is not None + and int(anchor.bulk_count or 0) > 0 + ): + ok_indices = indices + no_toc: list[int] = [] + segments = [ + CalibrationSegment( + offset=int(anchor.offset), + leaf_start=0, + leaf_end=max(0, len(ok_indices) - 1), + entry_indices=ok_indices, + status="ok", + ) + ] + else: + ok_indices = [] + no_toc = list(indices) + segments = [] + + out.append( + CalibrationRegime( + kind=kind, + offset=anchor.offset if kind == "decimal" else regime.offset, + offset_status="ok" if segments else "failed", + entry_indices=indices, + samples=list(regime.samples), + posterior=list(regime.posterior), + segments=segments, + no_toc_entry_indices=no_toc, + notes=( + f"production_bulk={anchor.bulk_count}; " + f"locate_agent={anchor.locate_agent}" + ), + ) + ) + return out + + +def finalize_calibration_result( + *, + result: CalibrationResult, + entries: list[dict[str, Any]], + toc_hierarchies: list[dict[str, Any]], + ctx: ToolContext, + page_count: int, + page_texts: dict[int, str] | None = None, + body_pages: list[int] | None = None, +) -> tuple[SkeletonAnchor, CalibrationResult]: + """Run production Phase-2 from an agent candidate offset.""" + offset_hint = pick_primary_offset(result) + texts = dict(page_texts or {}) + bodies = list(body_pages or sorted(texts.keys()) or list(range(1, page_count + 1))) + # Same TitleNode prep as C4 / extract_section_skeletons before anchoring. + from app.services.page_memory.skeleton_extractor import ( + _collapse_intermediate_single_child_chains, + ) + + nodes = _collapse_intermediate_single_child_chains( + extract_toc_nodes(toc_hierarchies) + ) + seed = _seed_overrides_from_samples(result=result, nodes=nodes) + + working, anchor = anchor_hierarchy_from_offset( + nodes=nodes, + offset_hint=offset_hint, + calibration_overrides=seed, + page_texts=texts, + body_pages=bodies, + page_count=page_count, + ctx=ctx, + ) + logger.info( + "[calibration.completion] offset={} status={} bulk={} pruned={} nodes={}", + anchor.offset, + anchor.offset_status, + anchor.bulk_count, + anchor.pruned_count, + len(working), + ) + + regimes = _annotate_regimes_from_anchor( + result=result, anchor=anchor, entries=entries + ) + complete = sum(len(r.segments) for r in regimes) + notes_parts = [result.notes] if result.notes else [] + notes_parts.append( + f"phase2 production locate_agent={anchor.locate_agent} " + f"bulk={anchor.bulk_count} complete_regime_segments={complete}" + ) + finalized = CalibrationResult( + status="ok" if anchor.offset_status == "ok" and anchor.bulk_count > 0 else "failed", + regimes=regimes, + offset=anchor.offset, + offset_status=anchor.offset_status, + tool_calls=result.tool_calls, + notes="; ".join(p for p in notes_parts if p), + region_index=result.region_index, + history_tail=list(result.history_tail), + ) + return anchor, finalized + + +def build_calibration_payload( + *, + anchor: SkeletonAnchor, + result: CalibrationResult, + no_links: bool, + region_payloads: list[dict[str, Any]] | None = None, + tool_calls: int | None = None, +) -> dict[str, Any]: + """Production SkeletonAnchor fields + experiment diagnostics.""" + payload = serialize_skeleton_anchor(anchor) + payload.update( + { + "status": result.status, + "regimes": [regime.to_dict() if hasattr(regime, "to_dict") else regime for regime in ( + # dataclasses asdict via CalibrationResult + result.to_dict().get("regimes") or [] + )], + "regions": list(region_payloads or []), + "tool_calls": int(tool_calls if tool_calls is not None else result.tool_calls), + "notes": result.notes, + "no_links": no_links, + } + ) + return payload diff --git a/apps/worker/app/services/document_agent/agents/calibration/service.py b/apps/worker/app/services/document_agent/agents/calibration/service.py new file mode 100644 index 000000000..c6e688022 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/service.py @@ -0,0 +1,68 @@ +"""Production calibration entry: Agent Phase-1 offset discovery. + +Same return shape as the former ``calibrate_offset_via_vlm`` so callers can +swap without changing prune / bulk / null-page. +""" + +from __future__ import annotations + +from typing import Any + +from loguru import logger + +from app.services.document_agent.agents.calibration.loop import run_calibration_phase1 +from app.services.document_agent.agents.calibration.procedure import ( + pick_primary_offset, + _seed_overrides_from_samples, +) +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.structure.hierarchy_locator import TitleMatch, TitleNode + + +def calibrate_offset( + *, + nodes: list[TitleNode], + toc_hierarchies: list[dict[str, Any]] | None, + ctx: ToolContext | None, + page_texts: dict[int, str], + page_count: int, +) -> tuple[int | None, dict[tuple[str, ...], TitleMatch]]: + """Discover printed→physical offset via the calibration SubAgent (Phase 1). + + Returns ``(offset, seed_overrides)``. Phase 2 (tail / bisect / null-page) + stays in ``anchor_hierarchy_from_offset`` using the caller's node tree. + """ + if ctx is None: + return None, {} + hierarchies = list(toc_hierarchies or []) + if not hierarchies: + return None, {} + + if page_count and not ctx.blackboard.page_count: + ctx.blackboard.page_count = int(page_count) + if page_texts and not ctx.blackboard.page_full_text_cache: + ctx.blackboard.page_full_text_cache = dict(page_texts) + + try: + phase1 = run_calibration_phase1( + ctx=ctx, + toc_hierarchies=hierarchies, + region_index=0, + page_count=int(page_count or ctx.blackboard.page_count or 0), + ) + except Exception as exc: + logger.warning("[calibration] Phase-1 failed: {}", exc) + return None, {} + + offset = pick_primary_offset(phase1) + if offset is None: + logger.info("[calibration] Phase-1 produced no primary offset") + return None, {} + + seed = _seed_overrides_from_samples(result=phase1, nodes=nodes) + logger.info( + "[calibration] Phase-1 offset={} seed_overrides={}", + offset, + len(seed), + ) + return offset, seed diff --git a/apps/worker/app/services/document_agent/agents/calibration/tools.py b/apps/worker/app/services/document_agent/agents/calibration/tools.py new file mode 100644 index 000000000..a2d7c44b3 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/tools.py @@ -0,0 +1,140 @@ +"""Tools for the calibration SubAgent (local registry, not PROFILE gates).""" + +from __future__ import annotations + +import time +from typing import Any + +from app.services.document_agent.agents.calibration.types import ( + calibration_result_from_dict, +) +from app.services.document_agent.manifest import ToolContext, ToolResult +from app.services.document_agent.registry import ToolRegistry, ToolSpec +from app.services.document_agent.tools.inspect_pages import inspect_pages + + +def build_calibration_registry() -> ToolRegistry: + registry = ToolRegistry() + registry.register( + ToolSpec( + name="inspect.pages", + description=( + "Open one or more physical PDF pages, render them, and answer " + "the given question about those pages." + ), + parameters={ + "type": "object", + "properties": { + "pages": { + "type": "array", + "items": {"type": "integer"}, + "description": "1-based physical page numbers", + }, + "question": { + "type": "string", + "description": "Question to answer from the rendered pages", + }, + }, + "required": ["pages", "question"], + }, + preconditions=(), + handler=_calibration_inspect_pages, + ) + ) + registry.register( + ToolSpec( + name="calibration.submit", + description="Submit the final CalibrationResult and finish.", + parameters={ + "type": "object", + "properties": { + "result": { + "type": "object", + "description": ( + "Phase-1 CalibrationResult with candidate regime " + "offsets and entry_indices; Phase-2 completion runs after submit" + ), + }, + }, + "required": ["result"], + }, + preconditions=(), + handler=calibration_submit, + ) + ) + return registry + + +def _calibration_inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + merged = dict(args) + merged.setdefault("folder_name", "calibration_inspect") + merged.setdefault("prefix", "calib") + merged.setdefault("usage_task", "calibration.inspect_pages") + return inspect_pages(ctx, merged) + + +def calibration_submit(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + raw = args.get("result") + if not isinstance(raw, dict): + if any(key in args for key in ("status", "regimes", "offset")): + raw = { + key: args.get(key) + for key in ( + "status", + "regimes", + "offset", + "offset_status", + "tool_calls", + "notes", + "region_index", + ) + if key in args + } + else: + return ToolResult( + status="error", + error="calibration.submit requires result object", + latency_ms=int((time.monotonic() - start) * 1000), + ) + result = calibration_result_from_dict(raw) + if not result.regimes and result.status == "ok": + return ToolResult( + status="error", + error="ok result must include regimes", + latency_ms=int((time.monotonic() - start) * 1000), + ) + tool_calls = int(ctx.blackboard.global_signals.get("calibration_tool_calls") or 0) + result.tool_calls = tool_calls + region_index = ctx.blackboard.global_signals.get("calibration_region_index") + if region_index is not None: + result.region_index = int(region_index) + ctx.blackboard.global_signals["calibration_result"] = result.to_dict() + ctx.blackboard.global_signals["calibration_done"] = True + return ToolResult( + status="ok", + payload=result.to_dict(), + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={"status": result.status, "regimes": len(result.regimes)}, + ) + + +def strip_toc_links(hierarchies: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Return hierarchies with entry ``link`` keys removed.""" + out: list[dict[str, Any]] = [] + for hierarchy in hierarchies: + if not isinstance(hierarchy, dict): + continue + cloned = dict(hierarchy) + entries = hierarchy.get("toc_with_level") + if isinstance(entries, list): + new_entries: list[dict[str, Any]] = [] + for entry in entries: + if not isinstance(entry, dict): + continue + item = dict(entry) + item.pop("link", None) + new_entries.append(item) + cloned["toc_with_level"] = new_entries + out.append(cloned) + return out diff --git a/apps/worker/app/services/document_agent/agents/calibration/types.py b/apps/worker/app/services/document_agent/agents/calibration/types.py new file mode 100644 index 000000000..3b22cd8b5 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/types.py @@ -0,0 +1,151 @@ +"""Calibration SubAgent result types.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any + + +@dataclass +class CalibrationSample: + title: str + printed_label: str | int | None = None + physical: int | None = None + method: str | None = None + + +@dataclass +class CalibrationPosterior: + title: str + expected_physical: int | None = None + confirmed: bool | None = None + method: str | None = None + + +@dataclass +class CalibrationSegment: + """A contiguous leaf range that fully completed Phase-2 for one offset.""" + + offset: int + leaf_start: int + leaf_end: int + entry_indices: list[int] = field(default_factory=list) + status: str = "ok" + + +@dataclass +class CalibrationRegime: + kind: str + offset: int | None = None + offset_status: str = "failed" + entry_indices: list[int] = field(default_factory=list) + samples: list[CalibrationSample] = field(default_factory=list) + posterior: list[CalibrationPosterior] = field(default_factory=list) + segments: list[CalibrationSegment] = field(default_factory=list) + no_toc_entry_indices: list[int] = field(default_factory=list) + notes: str = "" + + +@dataclass +class CalibrationResult: + status: str + regimes: list[CalibrationRegime] = field(default_factory=list) + offset: int | None = None + offset_status: str = "failed" + tool_calls: int = 0 + notes: str = "" + region_index: int | None = None + # Debug-only trail from the ReAct loop (not part of submit schema). + history_tail: list[dict[str, Any]] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def _as_optional_int(value: Any) -> int | None: + if value is None or isinstance(value, bool): + return None + if isinstance(value, int): + return value + if isinstance(value, float): + return int(value) + if isinstance(value, str): + text = value.strip() + if not text: + return None + try: + return int(text) + except ValueError: + return None + return None + + +def _as_int_list(value: Any) -> list[int]: + if not isinstance(value, list): + return [] + out: list[int] = [] + for item in value: + parsed = _as_optional_int(item) + if parsed is not None: + out.append(parsed) + return out + + +def calibration_result_from_dict(data: dict[str, Any]) -> CalibrationResult: + regimes: list[CalibrationRegime] = [] + for raw in data.get("regimes") or []: + if not isinstance(raw, dict): + continue + samples = [ + CalibrationSample( + title=str(s.get("title") or ""), + printed_label=s.get("printed_label"), + physical=_as_optional_int(s.get("physical")), + method=s.get("method") if isinstance(s.get("method"), str) else None, + ) + for s in (raw.get("samples") or []) + if isinstance(s, dict) + ] + posterior = [ + CalibrationPosterior( + title=str(p.get("title") or ""), + expected_physical=_as_optional_int(p.get("expected_physical")), + confirmed=p.get("confirmed") if isinstance(p.get("confirmed"), bool) else None, + method=p.get("method") if isinstance(p.get("method"), str) else None, + ) + for p in (raw.get("posterior") or []) + if isinstance(p, dict) + ] + segments = [ + CalibrationSegment( + offset=_as_optional_int(seg.get("offset")) or 0, + leaf_start=_as_optional_int(seg.get("leaf_start")) or 0, + leaf_end=_as_optional_int(seg.get("leaf_end")) or 0, + entry_indices=_as_int_list(seg.get("entry_indices")), + status=str(seg.get("status") or "ok"), + ) + for seg in (raw.get("segments") or []) + if isinstance(seg, dict) and _as_optional_int(seg.get("offset")) is not None + ] + regimes.append( + CalibrationRegime( + kind=str(raw.get("kind") or "other"), + offset=_as_optional_int(raw.get("offset")), + offset_status=str(raw.get("offset_status") or "failed"), + entry_indices=_as_int_list(raw.get("entry_indices")), + samples=samples, + posterior=posterior, + segments=segments, + no_toc_entry_indices=_as_int_list(raw.get("no_toc_entry_indices")), + notes=str(raw.get("notes") or ""), + ) + ) + return CalibrationResult( + status=str(data.get("status") or "failed"), + regimes=regimes, + offset=_as_optional_int(data.get("offset")), + offset_status=str(data.get("offset_status") or "failed"), + tool_calls=_as_optional_int(data.get("tool_calls")) or 0, + notes=str(data.get("notes") or ""), + region_index=_as_optional_int(data.get("region_index")), + ) diff --git a/apps/worker/app/services/document_agent/structure/structure_anchoring.py b/apps/worker/app/services/document_agent/structure/structure_anchoring.py new file mode 100644 index 000000000..d6232fc1a --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/structure_anchoring.py @@ -0,0 +1,825 @@ +"""Shared hierarchy anchoring: offset calibrate, null-page locate, bulk apply. + +Extracted from page_memory.skeleton_extractor so profile-time skeleton phase +and page-memory C4 share one implementation. No page_memory imports. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any + +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.structure.hierarchy_locator import ( + TitleMatch, + TitleNode, + first_leaf_start_under, + iter_leaf_title_nodes, + last_leaf_start_under, + locate_title_compact_strict, +) +from app.services.document_agent.structure.page_locate_agent import ( + verify_section_page_choice, +) +from loguru import logger + +def prune_out_of_scope_nodes( + nodes: list[TitleNode], + *, + offset: int, + page_count: int, +) -> tuple[list[TitleNode], int]: + """Remove leaf nodes whose printed_page + offset exceeds page_count. + + Bottom-up: prune out-of-scope leaves, then remove intermediate nodes + that become childless after pruning. Returns (pruned_tree, removed_count). + """ + removed = 0 + + def _prune(node: TitleNode) -> TitleNode | None: + nonlocal removed + if not node.children: + if node.printed_page is not None: + expected = node.printed_page + offset + if expected > page_count or expected < 1: + removed += 1 + return None + return node + pruned_children = [] + for child in node.children: + result = _prune(child) + if result is not None: + pruned_children.append(result) + if not pruned_children: + removed += 1 + return None + return replace(node, children=pruned_children) + + pruned = [] + for node in nodes: + result = _prune(node) + if result is not None: + pruned.append(result) + + if removed: + logger.info( + "[structure_anchoring] pruned {} out-of-scope TOC nodes " + "(printed_page + offset={} exceeds page_count={})", + removed, + offset, + page_count, + ) + + return pruned, removed + + +def toc_range_start(hierarchy: dict[str, Any]) -> int | None: + toc_range = hierarchy.get("toc_range") + if not isinstance(toc_range, (list, tuple)) or not toc_range: + return None + try: + return int(toc_range[0]) + except (TypeError, ValueError): + return None + + +def toc_range_end(hierarchy: dict[str, Any]) -> int | None: + toc_range = hierarchy.get("toc_range") + if not isinstance(toc_range, (list, tuple)) or not toc_range: + return None + try: + return int(toc_range[-1]) + except (TypeError, ValueError): + return None + + +# ── Null-page parent locate (compact-strict + RTL visual) ─────────────────── + +_NULL_PARENT_VISUAL_CONFIDENCE = 0.6 + + +def locate_null_page_parent_overrides( + *, + nodes: list[TitleNode], + match_overrides: dict[tuple[str, ...], TitleMatch], + page_texts: dict[int, str], + body_pages: list[int], + ctx: ToolContext | None, +) -> tuple[dict[tuple[str, ...], TitleMatch], list[dict[str, Any]]]: + """Locate TOC parents with ``printed_page=None`` into ``match_overrides``. + + Window for parent P: ``[last leaf start under previous same-level sibling, + first leaf start under P]``. Text path is compact→strict unique page; on + miss/ambiguity, scan right→left with ``verify_section_page_choice``. + + Returns ``(overrides, report)`` where *report* lists every null-page parent + attempt (for debug / LLM-call accounting). + """ + if not nodes or not body_pages: + return dict(match_overrides), [] + + out = dict(match_overrides) + body_set = set(body_pages) + parent_scope_start = body_pages[0] + report: list[dict[str, Any]] = [] + + def walk( + sibling_nodes: list[TitleNode], + parent_titles: tuple[str, ...], + scope_start: int, + ) -> None: + for index, node in enumerate(sibling_nodes): + path_titles = (*parent_titles, node.title) + if ( + node.children + and node.printed_page is None + and path_titles not in out + ): + if index > 0: + left = last_leaf_start_under( + sibling_nodes[index - 1], parent_titles, out + ) + if left is None: + left = scope_start + else: + left = scope_start + right = first_leaf_start_under(node, parent_titles, out) + entry: dict[str, Any] = { + "path_titles": list(path_titles), + "title": node.title, + "printed_page": None, + "window": None, + "result": "skipped_no_right", + "page": None, + "accept": None, + "visual_verify_calls": 0, + } + if right is None or right < left: + report.append(entry) + logger.info( + "[structure_anchoring] null-page parent skipped: " + "title={!r} reason=no_located_first_child left={}", + node.title, + left, + ) + else: + entry["window"] = [left, right] + scope_pages = [ + page for page in body_pages if left <= page <= right + ] + match = locate_title_compact_strict( + node.title, + scope_pages=scope_pages, + page_texts=page_texts, + ) + visual_calls = 0 + if match is None and ctx is not None: + match, visual_calls = _visual_rtl_locate_parent( + title=node.title, + left=left, + right=right, + body_set=body_set, + ctx=ctx, + ) + entry["visual_verify_calls"] = visual_calls + if match is not None and match.page in body_set: + out[path_titles] = match + entry["result"] = str(match.evidence.get("accept") or match.source) + entry["page"] = match.page + entry["accept"] = match.evidence.get("accept") + logger.info( + "[structure_anchoring] null-page parent located: " + "title={!r} page={} window={} accept={} visual_calls={}", + node.title, + match.page, + [left, right], + match.evidence.get("accept"), + visual_calls, + ) + else: + entry["result"] = "unresolved" + logger.info( + "[structure_anchoring] null-page parent unresolved: " + "title={!r} window={} visual_calls={}", + node.title, + [left, right], + visual_calls, + ) + report.append(entry) + if node.children: + child_scope_start = ( + out[path_titles].page if path_titles in out else scope_start + ) + walk(node.children, path_titles, child_scope_start) + + walk(nodes, (), parent_scope_start) + logger.info( + "[structure_anchoring] null-page parent locate summary: " + "attempted={} located={} unresolved={} visual_verify_calls={}", + len(report), + sum(1 for row in report if row.get("page") is not None), + sum(1 for row in report if row.get("result") == "unresolved"), + sum(int(row.get("visual_verify_calls") or 0) for row in report), + ) + return out, report + + +def _visual_rtl_locate_parent( + *, + title: str, + left: int, + right: int, + body_set: set[int], + ctx: ToolContext, +) -> tuple[TitleMatch | None, int]: + """Confirm parent title from right boundary toward left via VLM verify.""" + visual_calls = 0 + for page in range(right, left - 1, -1): + if page not in body_set: + continue + candidate = TitleMatch( + page=page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[page], + evidence={"null_page_parent_probe": True}, + ) + visual_calls += 1 + result = verify_section_page_choice( + ctx=ctx, + title=title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + selected = result.get("selected_page") + confidence = float(result.get("confidence") or 0.0) + if selected != page or confidence < _NULL_PARENT_VISUAL_CONFIDENCE: + continue + if result.get("source") == "agent_vlm": + return ( + TitleMatch( + page=page, + confidence=confidence, + source="agent_vlm", + matched_line="", + score=confidence, + candidates=[page], + evidence={ + "accept": "visual_rtl", + "reason": result.get("reason", ""), + "visual_verify_calls": visual_calls, + }, + ), + visual_calls, + ) + return ( + TitleMatch( + page=page, + confidence=confidence, + source="agent_heuristic", + matched_line="", + score=confidence, + candidates=[page], + evidence={ + "accept": "visual_rtl", + "reason": result.get("reason", ""), + "visual_verify_calls": visual_calls, + }, + ), + visual_calls, + ) + return None, visual_calls + + +# ── Offset calibration (Agent Phase-1) ─────────────────────────────────────── + + +def calibrate_offset( + *, + nodes: list[TitleNode], + toc_hierarchies: list[dict[str, Any]] | None, + ctx: ToolContext | None, + page_texts: dict[int, str], + page_count: int, +) -> tuple[int | None, dict[tuple[str, ...], TitleMatch]]: + """Discover printed→physical offset (calibration SubAgent Phase-1). + + Production and debug share this entry. Phase-2 bulk/bisect stays in + ``anchor_hierarchy_from_offset``. + """ + from app.services.document_agent.agents.calibration.service import ( + calibrate_offset as _agent_calibrate_offset, + ) + + return _agent_calibrate_offset( + nodes=nodes, + toc_hierarchies=toc_hierarchies, + ctx=ctx, + page_texts=page_texts, + page_count=page_count, + ) + + +def calibrate_offset_via_vlm( + *, + nodes: list[TitleNode], + toc_hierarchies: list[dict[str, Any]] | None, + ctx: ToolContext | None, + page_texts: dict[int, str], + page_count: int, +) -> tuple[int | None, dict[tuple[str, ...], TitleMatch]]: + """Deprecated alias — use ``calibrate_offset`` (Agent Phase-1).""" + return calibrate_offset( + nodes=nodes, + toc_hierarchies=toc_hierarchies, + ctx=ctx, + page_texts=page_texts, + page_count=page_count, + ) + + +def toc_cluster_end_page(toc_hierarchies: list[dict[str, Any]] | None) -> int | None: + """Get the last physical page of the primary TOC cluster.""" + if not toc_hierarchies: + return None + end_pages: list[int] = [] + for hierarchy in toc_hierarchies: + end = toc_range_end(hierarchy) + if end is not None: + end_pages.append(end) + return max(end_pages) if end_pages else None + + +# ── Offset-guided bulk anchoring with recursive recalibrate (Phase A3) ─────── + +_TAIL_VERIFY_CONFIDENCE_THRESHOLD = 0.6 +_MAX_RECALIBRATE_DEPTH = 5 +_MAX_RECALIBRATE_DELTA = 5 + + +def _verify_offset_tail( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, +) -> bool: + """VLM-verify that the offset holds for the last leaf entry (Theorem 1). + + If head offset == tail offset, monotonicity guarantees all intermediate + entries share the same offset. + + Prefers a tail leaf whose expected page is strictly less than page_count + (boundary pages are unreliable for VLM verification). + """ + tail_leaves = [ + (path, node) for path, node in reversed(leaves) if node.printed_page is not None + ] + if not tail_leaves: + return True + + # Prefer non-boundary: printed_page + offset < page_count + selected = None + for path, node in tail_leaves: + pp = node.printed_page + if pp is None: + continue + expected = pp + offset + if 1 <= expected < page_count: + selected = (path, node) + break + if selected is None: + # All leaves are at the boundary; fall back to the last one + selected = tail_leaves[0] + + path, node = selected + printed_page = node.printed_page + if printed_page is None: + return True + expected_page = printed_page + offset + if expected_page < 1 or expected_page > page_count: + return False + + candidate = TitleMatch( + page=expected_page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[expected_page], + evidence={"tail_verify_probe": True}, + ) + result = verify_section_page_choice( + ctx=ctx, + title=node.title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + confirmed = ( + result.get("selected_page") == expected_page + and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD + ) + logger.info( + "[structure_anchoring] tail verify: title={!r} expected_page={} confirmed={} confidence={}", + node.title, + expected_page, + confirmed, + result.get("confidence", 0), + ) + return confirmed + + +def _vlm_confirm_single_page( + *, + ctx: ToolContext, + title: str, + expected_page: int, + page_count: int, +) -> bool: + """Single-page VLM confirmation for binary search steps.""" + if expected_page < 1 or expected_page > page_count: + return False + candidate = TitleMatch( + page=expected_page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[expected_page], + evidence={"bisect_probe": True}, + ) + result = verify_section_page_choice( + ctx=ctx, + title=title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + return ( + result.get("selected_page") == expected_page + and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD + ) + + +def _bisect_offset_breakpoint( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, +) -> int: + """Binary search for the last leaf index where offset is valid. O(log n) VLM calls.""" + lo, hi = 0, len(leaves) - 1 + while lo < hi: + mid = (lo + hi + 1) // 2 + _, node = leaves[mid] + if node.printed_page is None: + hi = mid - 1 + continue + expected = node.printed_page + offset + if _vlm_confirm_single_page( + ctx=ctx, title=node.title, expected_page=expected, page_count=page_count + ): + lo = mid + else: + hi = mid - 1 + logger.info( + "[structure_anchoring] bisect breakpoint: last_valid_index={} / total={}", + lo, + len(leaves), + ) + return lo + + +def bulk_offset_matches( + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, +) -> dict[tuple[str, ...], TitleMatch]: + """Generate TitleMatch overrides for all leaves using offset. No VLM calls.""" + matches: dict[tuple[str, ...], TitleMatch] = {} + for path_titles, node in leaves: + if node.printed_page is None: + continue + page = node.printed_page + offset + matches[path_titles] = TitleMatch( + page=page, + confidence=0.88, + source="agent_vlm", + matched_line="", + score=0.88, + candidates=[page], + evidence={ + "bulk_offset": True, + "offset": offset, + "printed_page": node.printed_page, + }, + ) + return matches + + +def _recalibrate_after_breakpoint( + *, + entry_node: TitleNode, + old_offset: int, + ctx: ToolContext, + page_count: int, +) -> int | None: + """Probe offsets old_offset+1, +2, ... to find new offset after breakpoint. + + Monotonicity guarantees new offset > old offset, so search space is tiny. + """ + entry_printed_page = entry_node.printed_page + if entry_printed_page is None: + return None + for delta in range(1, _MAX_RECALIBRATE_DELTA + 1): + new_offset = old_offset + delta + if _vlm_confirm_single_page( + ctx=ctx, + title=entry_node.title, + expected_page=entry_printed_page + new_offset, + page_count=page_count, + ): + logger.info( + "[structure_anchoring] recalibrate: title={!r} new_offset={} (delta=+{})", + entry_node.title, + new_offset, + delta, + ) + return new_offset + return None + + +def offset_guided_anchoring( + *, + nodes: list[TitleNode], + offset: int, + ctx: ToolContext, + page_count: int, + calibration_overrides: dict[tuple[str, ...], TitleMatch], +) -> dict[tuple[str, ...], TitleMatch] | None: + """Offset-guided bulk anchoring with recursive recalibrate on breakpoints. + + Strategy: + 1. Tail verify last leaf with current offset + 2. If pass → bulk apply all leaves (Theorem 1) + 3. If fail → binary search for breakpoint + 4. Bulk apply leaves before breakpoint + 5. Recalibrate: probe remaining[0] with offset+1, +2, ... (monotonicity) + 6. Recurse on remaining segment with new offset + 7. If recalibrate fails → return partial (caller falls back for remainder) + + Returns match_overrides for all anchored leaves, or None for full fallback. + """ + leaves = [ + (path, node) + for path, node in iter_leaf_title_nodes(nodes) + if node.printed_page is not None + ] + if len(leaves) < 2: + return None + + all_matches: dict[tuple[str, ...], TitleMatch] = {} + all_matches.update(calibration_overrides) + + _anchor_segment_recursive( + leaves=leaves, + offset=offset, + ctx=ctx, + page_count=page_count, + matches=all_matches, + depth=0, + ) + + if not all_matches: + return None + + logger.info( + "[structure_anchoring] offset bulk anchoring: {} / {} leaves anchored", + len(all_matches), + len(leaves), + ) + return all_matches + + +def _anchor_segment_recursive( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, + matches: dict[tuple[str, ...], TitleMatch], + depth: int, +) -> None: + """Recursively anchor a segment of leaves, handling multiple breakpoints.""" + if not leaves or depth >= _MAX_RECALIBRATE_DEPTH: + return + + if _verify_offset_tail(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count): + bulk = bulk_offset_matches(leaves, offset) + matches.update(bulk) + return + + bp = _bisect_offset_breakpoint(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count) + confirmed_leaves = leaves[: bp + 1] + if confirmed_leaves: + bulk = bulk_offset_matches(confirmed_leaves, offset) + matches.update(bulk) + + remaining = leaves[bp + 1:] + if not remaining: + return + + _, first_remaining_node = remaining[0] + new_offset = _recalibrate_after_breakpoint( + entry_node=first_remaining_node, + old_offset=offset, + ctx=ctx, + page_count=page_count, + ) + if new_offset is None: + return + + _anchor_segment_recursive( + leaves=remaining, + offset=new_offset, + ctx=ctx, + page_count=page_count, + matches=matches, + depth=depth + 1, + ) + + +@dataclass +class SkeletonAnchor: + offset: int | None + offset_status: str + match_overrides: dict[tuple[str, ...], TitleMatch] + null_page_report: list[dict[str, Any]] + bulk_count: int + pruned_count: int = 0 + locate_agent: str = "offset_only" + + +def serialize_title_match(match: TitleMatch) -> dict[str, Any]: + return { + "page": match.page, + "confidence": match.confidence, + "source": match.source, + "matched_line": match.matched_line, + "score": match.score, + "candidates": list(match.candidates), + "evidence": dict(match.evidence or {}), + } + + +def serialize_skeleton_anchor(anchor: SkeletonAnchor) -> dict[str, Any]: + """JSON-friendly SkeletonAnchor (path tuples joined by ' / ').""" + overrides: dict[str, Any] = {} + for path, match in (anchor.match_overrides or {}).items(): + key = " / ".join(str(part) for part in path) + overrides[key] = serialize_title_match(match) + return { + "offset": anchor.offset, + "offset_status": anchor.offset_status, + "match_overrides": overrides, + "null_page_report": list(anchor.null_page_report or []), + "bulk_count": int(anchor.bulk_count or 0), + "pruned_count": int(anchor.pruned_count or 0), + "locate_agent": anchor.locate_agent, + } + + +def deserialize_title_match(data: dict[str, Any]) -> TitleMatch: + return TitleMatch( + page=int(data["page"]), + confidence=float(data.get("confidence") or 0.0), + source=data.get("source") or "agent_vlm", # type: ignore[arg-type] + matched_line=str(data.get("matched_line") or ""), + score=float(data.get("score") or 0.0), + candidates=[int(p) for p in (data.get("candidates") or [])], + evidence=dict(data.get("evidence") or {}), + ) + + +def deserialize_skeleton_anchor(data: dict[str, Any]) -> SkeletonAnchor: + raw_overrides = data.get("match_overrides") or {} + overrides: dict[tuple[str, ...], TitleMatch] = {} + if isinstance(raw_overrides, dict): + for key, value in raw_overrides.items(): + if not isinstance(value, dict): + continue + if isinstance(key, str): + path = tuple(part.strip() for part in key.split(" / ") if part.strip()) + elif isinstance(key, (list, tuple)): + path = tuple(str(part) for part in key) + else: + continue + if path: + overrides[path] = deserialize_title_match(value) + return SkeletonAnchor( + offset=data.get("offset") if data.get("offset") is None else int(data["offset"]), + offset_status=str(data.get("offset_status") or "failed"), + match_overrides=overrides, + null_page_report=list(data.get("null_page_report") or []), + bulk_count=int(data.get("bulk_count") or 0), + pruned_count=int(data.get("pruned_count") or 0), + locate_agent=str(data.get("locate_agent") or "offset_only"), + ) + + +def anchor_hierarchy_from_offset( + *, + nodes: list[TitleNode], + offset_hint: int | None, + calibration_overrides: dict[tuple[str, ...], TitleMatch] | None = None, + page_texts: dict[int, str], + body_pages: list[int], + page_count: int, + ctx: ToolContext | None, +) -> tuple[list[TitleNode], SkeletonAnchor]: + """Production prune → bulk → null-page given a precomputed offset. + + Swap point for agent / VLM calibration: both feed ``offset_hint`` here. + """ + seed_overrides = dict(calibration_overrides or {}) + pruned_count = 0 + working = nodes + if offset_hint is not None: + working, pruned_count = prune_out_of_scope_nodes( + working, offset=offset_hint, page_count=page_count + ) + + offset_matches: dict[tuple[str, ...], TitleMatch] | None = None + if offset_hint is not None and ctx is not None and working: + offset_matches = offset_guided_anchoring( + nodes=working, + offset=offset_hint, + ctx=ctx, + page_count=page_count, + calibration_overrides=seed_overrides, + ) + + if offset_matches is not None: + match_overrides = offset_matches + locate_agent = "offset_guided_bulk" + bulk_count = len(offset_matches) + else: + match_overrides = seed_overrides + locate_agent = "offset_only" + bulk_count = 0 + + match_overrides, null_page_report = locate_null_page_parent_overrides( + nodes=working, + match_overrides=match_overrides, + page_texts=page_texts, + body_pages=body_pages, + ctx=ctx, + ) + + if offset_hint is None: + offset_status = "failed" if ctx is not None else "skipped" + else: + offset_status = "ok" + + return working, SkeletonAnchor( + offset=offset_hint, + offset_status=offset_status, + match_overrides=match_overrides, + null_page_report=null_page_report, + bulk_count=bulk_count, + pruned_count=pruned_count, + locate_agent=locate_agent, + ) + + +def anchor_hierarchy( + *, + nodes: list[TitleNode], + toc_hierarchies: list[dict[str, Any]] | None, + page_texts: dict[int, str], + body_pages: list[int], + page_count: int, + ctx: ToolContext | None, +) -> tuple[list[TitleNode], SkeletonAnchor]: + """Run offset → prune → bulk → null-page in production order. + + Returns possibly-pruned nodes and the anchor payload. Caller owns + resolve_hierarchy_page_ranges / skeleton assembly. + """ + offset_hint, calibration_overrides = calibrate_offset( + nodes=nodes, + toc_hierarchies=toc_hierarchies, + ctx=ctx, + page_texts=page_texts, + page_count=page_count, + ) + return anchor_hierarchy_from_offset( + nodes=nodes, + offset_hint=offset_hint, + calibration_overrides=calibration_overrides, + page_texts=page_texts, + body_pages=body_pages, + page_count=page_count, + ctx=ctx, + ) 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 new file mode 100644 index 000000000..e332f1eb0 --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/toc_link_enrichment.py @@ -0,0 +1,310 @@ +"""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). +""" + +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: + toc_page: int # 1-based + dest_physical_page: int # 1-based + anchor_text: str + kind: int | None = None + + +@dataclass(frozen=True) +class TocLinkEnrichStats: + toc_pages_scanned: list[int] + links_raw: int + links_internal: int + entries_total: int + entries_matched: int + skipped_no_links: bool = False + + +def _toc_pages_from_hierarchy(hierarchy: dict[str, Any]) -> list[int]: + """Physical pages that are actual TOC content (not VLM scan expansion).""" + pages: set[int] = set() + toc_range = hierarchy.get("toc_range") + if isinstance(toc_range, (list, tuple)) and len(toc_range) >= 2: + start, end = int(toc_range[0]), int(toc_range[1]) + if start > 0 and end >= start: + pages.update(range(start, end + 1)) + # Do NOT include scan_range: that window often covers non-TOC body pages + # used only for VLM boundary detection. + return sorted(pages) + + +def _anchor_text_for_rect(page: Any, rect: Any) -> str: + import fitz + + words = page.get_text("words") or [] + hit: list[tuple[float, float, str]] = [] + target = fitz.Rect(rect) + # Slightly expand so thin link boxes still catch title glyphs. + target = target + (-2, -2, 2, 2) + for word in words: + x0, y0, x1, y1, text = word[:5] + if not str(text).strip(): + continue + if fitz.Rect(x0, y0, x1, y1).intersects(target): + hit.append((float(y0), float(x0), str(text))) + hit.sort() + return " ".join(part for _, _, part in hit).strip() + + +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.""" + import fitz + + if not toc_pages: + return [] + + out: list[TocPageLink] = [] + doc = fitz.open(pdf_path) + try: + for toc_page in toc_pages: + if toc_page < 1 or toc_page > doc.page_count: + 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: + continue + try: + dest_physical = int(dest_idx) + 1 + 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 + out.append( + TocPageLink( + toc_page=toc_page, + dest_physical_page=dest_physical, + anchor_text=anchor, + kind=int(kind) if kind is not None else None, + ) + ) + finally: + doc.close() + 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"} + ] + + matched = 0 + enriched: list[dict[str, Any]] = [] + used_dest_for_heading: set[str] = set() + + for entry in entries: + if not isinstance(entry, dict): + continue + new_entry = { + "heading": entry.get("heading"), + "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), + ) + 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: + enriched.append(new_entry) + continue + used_dest_for_heading.add(heading_key) + + new_entry["link"] = { + "physical_page": chosen.dest_physical_page, + } + matched += 1 + enriched.append(new_entry) + + return enriched, matched + + +def enrich_toc_hierarchies_with_links( + *, + pdf_path: str, + toc_hierarchies: list[dict[str, Any]] | None, +) -> tuple[list[dict[str, Any]], TocLinkEnrichStats]: + """Attach optional ``link`` fields onto matching TOC entries. + + If TOC pages have no internal links, hierarchies are returned unchanged. + """ + hierarchies = [dict(h) for h in (toc_hierarchies or []) if isinstance(h, dict)] + if not hierarchies: + return [], TocLinkEnrichStats( + toc_pages_scanned=[], + links_raw=0, + links_internal=0, + entries_total=0, + entries_matched=0, + skipped_no_links=True, + ) + + toc_pages: list[int] = [] + seen: set[int] = set() + for hierarchy in hierarchies: + for page in _toc_pages_from_hierarchy(hierarchy): + if page not in seen: + seen.add(page) + toc_pages.append(page) + + links = collect_toc_page_links(pdf_path, toc_pages) + if not links: + logger.info( + "[toc_link_enrich] no internal links on TOC pages {}; skip", + toc_pages, + ) + return hierarchies, TocLinkEnrichStats( + toc_pages_scanned=toc_pages, + links_raw=0, + links_internal=0, + entries_total=sum( + len(h.get("toc_with_level") or []) + for h in hierarchies + if isinstance(h.get("toc_with_level"), list) + ), + entries_matched=0, + skipped_no_links=True, + ) + + total_entries = 0 + total_matched = 0 + out: list[dict[str, Any]] = [] + for hierarchy in hierarchies: + entries = hierarchy.get("toc_with_level") + if not isinstance(entries, list): + out.append(hierarchy) + continue + enriched_entries, matched = match_toc_entries_to_links(entries, links) + total_entries += len(enriched_entries) + total_matched += matched + new_hierarchy = dict(hierarchy) + new_hierarchy["toc_with_level"] = enriched_entries + out.append(new_hierarchy) + + stats = TocLinkEnrichStats( + toc_pages_scanned=toc_pages, + links_raw=len(links), + links_internal=len(links), + entries_total=total_entries, + entries_matched=total_matched, + skipped_no_links=False, + ) + logger.info( + "[toc_link_enrich] toc_pages={} links={} entries={}/{} matched", + toc_pages, + len(links), + total_matched, + total_entries, + ) + return out, stats diff --git a/apps/worker/app/services/document_agent/tools/inspect_pages.py b/apps/worker/app/services/document_agent/tools/inspect_pages.py new file mode 100644 index 000000000..3946f71ea --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/inspect_pages.py @@ -0,0 +1,163 @@ +"""Generic inspect.pages tool: open physical pages, render, answer a question.""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from typing import Any, cast + +from loguru import logger + +from app.services.document_agent.manifest import ToolContext, ToolResult + + +_DEFAULT_PAGE_CAP = 5 + + +def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + """Open one or more physical PDF pages, render them, and answer ``question``.""" + start = time.monotonic() + raw_pages = args.get("pages") or [] + question = str(args.get("question") or "").strip() + if not isinstance(raw_pages, list) or not raw_pages: + return ToolResult( + status="error", + error="inspect.pages requires pages[]", + latency_ms=int((time.monotonic() - start) * 1000), + ) + if not question: + return ToolResult( + status="error", + error="inspect.pages requires question", + latency_ms=int((time.monotonic() - start) * 1000), + ) + + page_count = int(ctx.blackboard.page_count or 0) + pages: list[int] = [] + for item in raw_pages: + try: + page = int(item) + except (TypeError, ValueError): + continue + if 1 <= page <= page_count and page not in pages: + pages.append(page) + page_cap = int(ctx.settings.get("inspect_page_cap") or _DEFAULT_PAGE_CAP) + pages = pages[: max(page_cap, 1)] + if not pages: + return ToolResult( + status="error", + error="no valid physical pages in range", + latency_ms=int((time.monotonic() - start) * 1000), + ) + + used = int(ctx.blackboard.global_signals.get("inspect_pages_used") or 0) + # Backward-compatible calibration counter. + used = max( + used, + int(ctx.blackboard.global_signals.get("calibration_inspect_pages_used") or 0), + ) + page_budget = int(ctx.settings.get("inspect_page_budget") or 0) + if page_budget > 0 and used >= page_budget: + return ToolResult( + status="error", + error="inspect page budget exhausted", + latency_ms=int((time.monotonic() - start) * 1000), + ) + if page_budget > 0: + remain = max(page_budget - used, 0) + pages = pages[:remain] + if not pages: + return ToolResult( + status="error", + error="inspect page budget exhausted", + latency_ms=int((time.monotonic() - start) * 1000), + ) + + from app.services.document_agent.visual import render_pages + + folder_name = str(args.get("folder_name") or "inspect_pages") + prefix = str(args.get("prefix") or "inspect") + rendered = render_pages( + ctx, + pages, + folder_name=folder_name, + prefix=prefix, + timeout=120, + ) + if not rendered: + return ToolResult( + status="error", + error="render failed", + latency_ms=int((time.monotonic() - start) * 1000), + ) + + model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") + if not model: + return ToolResult( + status="error", + error="vlm_model missing", + latency_ms=int((time.monotonic() - start) * 1000), + ) + + prompt = ( + "Answer the question about the provided PDF page image(s). " + "Return strict JSON object with keys: " + '{"answer": string, "page_notes": [{"page": number, "note": string}], ' + '"confidence": number}. ' + "Include the word json in your reasoning.\n\n" + f"Pages: {pages}\nQuestion: {question}\n" + ) + content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + for item in rendered: + with open(str(item["png_path"]), "rb") as image_file: + img_b64 = base64.b64encode(image_file.read()).decode() + content_parts.append({"type": "text", "text": f"\n--- Page {item['page']} ---"}) + content_parts.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + } + ) + + usage_task = str(args.get("usage_task") or "document_agent.inspect_pages") + try: + from shared.services.ai.llm_overrides import get_vision_client + + client, model = get_vision_client(requested_model=str(model)) + raw, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=model, + temperature=0.0, + max_tokens=800, + response_format={"type": "json_object"}, + usage_task=usage_task, + ) + payload = json.loads(raw) if raw else {} + except Exception as exc: + logger.warning("[inspect.pages] VLM failed: {}", exc) + return ToolResult( + status="error", + error=f"vlm failed: {exc}", + latency_ms=int((time.monotonic() - start) * 1000), + ) + + next_used = used + len(pages) + ctx.blackboard.global_signals["inspect_pages_used"] = next_used + ctx.blackboard.global_signals["calibration_inspect_pages_used"] = next_used + return ToolResult( + status="ok", + payload={ + "pages": pages, + "question": question, + "answer": payload.get("answer"), + "page_notes": payload.get("page_notes") or [], + "confidence": payload.get("confidence"), + "raw": payload, + }, + latency_ms=int((time.monotonic() - start) * 1000), + tokens_used=int((usage or {}).get("total_tokens") or 0), + output_summary={"pages": pages, "answer": payload.get("answer")}, + ) + 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 30962fcdb..7aa3a57cf 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 @@ -645,10 +645,12 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: offset_hint: int | None = None if ctx.blackboard.toc_hierarchies: from app.services.document_agent.structure.hierarchy_locator import extract_toc_nodes - from app.services.page_memory.skeleton_extractor import _calibrate_offset_via_vlm + from app.services.document_agent.structure.structure_anchoring import ( + calibrate_offset, + ) nodes = extract_toc_nodes(ctx.blackboard.toc_hierarchies) - offset_hint, _ = _calibrate_offset_via_vlm( + offset_hint, _ = calibrate_offset( nodes=nodes, toc_hierarchies=ctx.blackboard.toc_hierarchies, ctx=ctx, diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index ce835659d..abce3c08b 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -14,20 +14,21 @@ PageAnatomyMap, ToolContext, ) -from app.services.document_agent.structure.page_locate_agent import ( - verify_section_page_choice, -) from app.services.document_agent.structure.hierarchy_locator import ( ResolvedHierarchyRange, - TitleMatch, TitleNode, extract_toc_nodes, - first_leaf_start_under, iter_leaf_title_nodes, - last_leaf_start_under, - locate_title_compact_strict, resolve_hierarchy_page_ranges, ) +from app.services.document_agent.structure.structure_anchoring import ( + anchor_hierarchy, + calibrate_offset, + locate_null_page_parent_overrides, + offset_guided_anchoring, + toc_range_end, + toc_range_start, +) from loguru import logger from shared.services.chunks.path_segments import ( append_document_path, @@ -37,56 +38,6 @@ _FRONT_TOC_REGION_GAP_PAGES = 5 -def _prune_out_of_scope_nodes( - nodes: list[TitleNode], - *, - offset: int, - page_count: int, -) -> tuple[list[TitleNode], int]: - """Remove leaf nodes whose printed_page + offset exceeds page_count. - - Bottom-up: prune out-of-scope leaves, then remove intermediate nodes - that become childless after pruning. Returns (pruned_tree, removed_count). - """ - from dataclasses import replace as _replace - - removed = 0 - - def _prune(node: TitleNode) -> TitleNode | None: - nonlocal removed - if not node.children: - if node.printed_page is not None: - expected = node.printed_page + offset - if expected > page_count or expected < 1: - removed += 1 - return None - return node - pruned_children = [] - for child in node.children: - result = _prune(child) - if result is not None: - pruned_children.append(result) - if not pruned_children: - removed += 1 - return None - return _replace(node, children=pruned_children) - - pruned = [] - for node in nodes: - result = _prune(node) - if result is not None: - pruned.append(result) - - if removed: - logger.info( - "[page_memory.skeleton] pruned {} out-of-scope TOC nodes " - "(printed_page + offset={} exceeds page_count={})", - removed, - offset, - page_count, - ) - - return pruned, removed @dataclass(frozen=True) @@ -161,73 +112,47 @@ def extract_section_skeletons( if pending_tocs: pending_starts: list[int] = [] for t in pending_tocs: - start = _toc_range_start(t) + start = toc_range_start(t) if start is not None: pending_starts.append(start) if pending_starts: primary_page_count = min(pending_starts) - 1 primary_body_pages = [p for p in body_pages if p <= primary_page_count] - offset_hint, calibration_overrides = _calibrate_offset_via_vlm( + resolve_nodes, skeleton_anchor = anchor_hierarchy( nodes=nodes, toc_hierarchies=toc_hierarchies if not hierarchy_nodes else None, - ctx=ctx, page_texts=page_texts, + body_pages=primary_body_pages, page_count=page_count, + ctx=ctx, ) + if skeleton_anchor.pruned_count and not resolve_nodes: + return [ + _root_skeleton( + root_path=root_path, + filename=filename, + page_count=page_count, + reason="all_toc_nodes_out_of_scope", + ) + ] - # Prune TOC nodes whose printed_page + offset exceeds the physical PDF. - pruned_count = 0 - if offset_hint is not None: - nodes, pruned_count = _prune_out_of_scope_nodes( - nodes, offset=offset_hint, page_count=page_count, - ) - if not nodes: - return [ - _root_skeleton( - root_path=root_path, - filename=filename, - page_count=page_count, - reason="all_toc_nodes_out_of_scope", - ) - ] - - # Phase A3: offset-guided bulk anchoring for printed-page leaves. - offset_matches: dict[tuple[str, ...], TitleMatch] | None = None - if offset_hint is not None and ctx is not None: - offset_matches = _offset_guided_anchoring( - nodes=nodes, - offset=offset_hint, - ctx=ctx, - page_count=page_count, - calibration_overrides=calibration_overrides, - ) - - if offset_matches is not None: - match_overrides = offset_matches + match_overrides = skeleton_anchor.match_overrides + null_page_report = skeleton_anchor.null_page_report + if skeleton_anchor.locate_agent == "offset_guided_bulk": locate_summary: dict[str, Any] = { "agent": "offset_guided_bulk", - "offset": offset_hint, - "bulk_count": len(offset_matches), - "pruned_out_of_scope": pruned_count, + "offset": skeleton_anchor.offset, + "bulk_count": skeleton_anchor.bulk_count, + "pruned_out_of_scope": skeleton_anchor.pruned_count, } else: - match_overrides = calibration_overrides locate_summary = { "agent": "offset_only", - "offset": offset_hint, + "offset": skeleton_anchor.offset, "reason": "offset_guided_anchoring_skipped_or_empty", - "pruned_out_of_scope": pruned_count, + "pruned_out_of_scope": skeleton_anchor.pruned_count, } - resolve_nodes = nodes - - match_overrides, null_page_report = locate_null_page_parent_overrides( - nodes=resolve_nodes, - match_overrides=match_overrides, - page_texts=page_texts, - body_pages=primary_body_pages, - ctx=ctx, - ) 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), @@ -423,19 +348,19 @@ def _select_global_toc_hierarchies( page_based = [ hierarchy for hierarchy in hierarchies - if hierarchy.get("toc_range_unit") == "page" and _toc_range_start(hierarchy) is not None + 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) + 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) + start = toc_range_start(hierarchy) + end = toc_range_end(hierarchy) if start is None or end is None: selected_indices.add(original_index) continue @@ -472,24 +397,6 @@ def _select_global_toc_hierarchies( return (selected or None), pending, summary -def _toc_range_start(hierarchy: dict[str, Any]) -> int | None: - toc_range = hierarchy.get("toc_range") - if not isinstance(toc_range, (list, tuple)) or not toc_range: - return None - try: - return int(toc_range[0]) - except (TypeError, ValueError): - return None - - -def _toc_range_end(hierarchy: dict[str, Any]) -> int | None: - toc_range = hierarchy.get("toc_range") - if not isinstance(toc_range, (list, tuple)) or not toc_range: - return None - try: - return int(toc_range[-1]) - except (TypeError, ValueError): - return None def _body_pages(*, anatomy: Any | None, page_count: int) -> list[int]: @@ -499,606 +406,6 @@ def _body_pages(*, anatomy: Any | None, page_count: int) -> list[int]: return [page for page in range(1, page_count + 1) if page not in excluded] -# ── Null-page parent locate (compact-strict + RTL visual) ─────────────────── - -_NULL_PARENT_VISUAL_CONFIDENCE = 0.6 - - -def locate_null_page_parent_overrides( - *, - nodes: list[TitleNode], - match_overrides: dict[tuple[str, ...], TitleMatch], - page_texts: dict[int, str], - body_pages: list[int], - ctx: ToolContext | None, -) -> tuple[dict[tuple[str, ...], TitleMatch], list[dict[str, Any]]]: - """Locate TOC parents with ``printed_page=None`` into ``match_overrides``. - - Window for parent P: ``[last leaf start under previous same-level sibling, - first leaf start under P]``. Text path is compact→strict unique page; on - miss/ambiguity, scan right→left with ``verify_section_page_choice``. - - Returns ``(overrides, report)`` where *report* lists every null-page parent - attempt (for debug / LLM-call accounting). - """ - if not nodes or not body_pages: - return dict(match_overrides), [] - - out = dict(match_overrides) - body_set = set(body_pages) - parent_scope_start = body_pages[0] - report: list[dict[str, Any]] = [] - - def walk( - sibling_nodes: list[TitleNode], - parent_titles: tuple[str, ...], - scope_start: int, - ) -> None: - for index, node in enumerate(sibling_nodes): - path_titles = (*parent_titles, node.title) - if ( - node.children - and node.printed_page is None - and path_titles not in out - ): - if index > 0: - left = last_leaf_start_under( - sibling_nodes[index - 1], parent_titles, out - ) - if left is None: - left = scope_start - else: - left = scope_start - right = first_leaf_start_under(node, parent_titles, out) - entry: dict[str, Any] = { - "path_titles": list(path_titles), - "title": node.title, - "printed_page": None, - "window": None, - "result": "skipped_no_right", - "page": None, - "accept": None, - "visual_verify_calls": 0, - } - if right is None or right < left: - report.append(entry) - logger.info( - "[page_memory.skeleton] null-page parent skipped: " - "title={!r} reason=no_located_first_child left={}", - node.title, - left, - ) - else: - entry["window"] = [left, right] - scope_pages = [ - page for page in body_pages if left <= page <= right - ] - match = locate_title_compact_strict( - node.title, - scope_pages=scope_pages, - page_texts=page_texts, - ) - visual_calls = 0 - if match is None and ctx is not None: - match, visual_calls = _visual_rtl_locate_parent( - title=node.title, - left=left, - right=right, - body_set=body_set, - ctx=ctx, - ) - entry["visual_verify_calls"] = visual_calls - if match is not None and match.page in body_set: - out[path_titles] = match - entry["result"] = str(match.evidence.get("accept") or match.source) - entry["page"] = match.page - entry["accept"] = match.evidence.get("accept") - logger.info( - "[page_memory.skeleton] null-page parent located: " - "title={!r} page={} window={} accept={} visual_calls={}", - node.title, - match.page, - [left, right], - match.evidence.get("accept"), - visual_calls, - ) - else: - entry["result"] = "unresolved" - logger.info( - "[page_memory.skeleton] null-page parent unresolved: " - "title={!r} window={} visual_calls={}", - node.title, - [left, right], - visual_calls, - ) - report.append(entry) - if node.children: - child_scope_start = ( - out[path_titles].page if path_titles in out else scope_start - ) - walk(node.children, path_titles, child_scope_start) - - walk(nodes, (), parent_scope_start) - logger.info( - "[page_memory.skeleton] null-page parent locate summary: " - "attempted={} located={} unresolved={} visual_verify_calls={}", - len(report), - sum(1 for row in report if row.get("page") is not None), - sum(1 for row in report if row.get("result") == "unresolved"), - sum(int(row.get("visual_verify_calls") or 0) for row in report), - ) - return out, report - - -def _visual_rtl_locate_parent( - *, - title: str, - left: int, - right: int, - body_set: set[int], - ctx: ToolContext, -) -> tuple[TitleMatch | None, int]: - """Confirm parent title from right boundary toward left via VLM verify.""" - visual_calls = 0 - for page in range(right, left - 1, -1): - if page not in body_set: - continue - candidate = TitleMatch( - page=page, - confidence=0.4, - source="agent_heuristic", - matched_line="", - score=0.4, - candidates=[page], - evidence={"null_page_parent_probe": True}, - ) - visual_calls += 1 - result = verify_section_page_choice( - ctx=ctx, - title=title, - candidate_matches=[candidate], - candidate_page_cap=1, - ) - selected = result.get("selected_page") - confidence = float(result.get("confidence") or 0.0) - if selected != page or confidence < _NULL_PARENT_VISUAL_CONFIDENCE: - continue - if result.get("source") == "agent_vlm": - return ( - TitleMatch( - page=page, - confidence=confidence, - source="agent_vlm", - matched_line="", - score=confidence, - candidates=[page], - evidence={ - "accept": "visual_rtl", - "reason": result.get("reason", ""), - "visual_verify_calls": visual_calls, - }, - ), - visual_calls, - ) - return ( - TitleMatch( - page=page, - confidence=confidence, - source="agent_heuristic", - matched_line="", - score=confidence, - candidates=[page], - evidence={ - "accept": "visual_rtl", - "reason": result.get("reason", ""), - "visual_verify_calls": visual_calls, - }, - ), - visual_calls, - ) - return None, visual_calls - - -# ── VLM offset calibration (Phase A1) ─────────────────────────────────────── -_CALIBRATION_WINDOW_PAGES = 10 -_CALIBRATION_LEAF_PROBE_COUNT = 3 - - -def _calibrate_offset_via_vlm( - *, - nodes: list[TitleNode], - toc_hierarchies: list[dict[str, Any]] | None, - ctx: ToolContext | None, - page_texts: dict[int, str], - page_count: int, -) -> tuple[int | None, dict[tuple[str, ...], TitleMatch]]: - """Scan pages after the TOC to find the first leaf entry via VLM. - - Computes offset = confirmed_physical_page - printed_page. - Returns (offset, match_overrides) where match_overrides contains the - confirmed entry so downstream locate doesn't re-process it. - """ - if ctx is None: - return None, {} - - toc_physical_end = _toc_cluster_end_page(toc_hierarchies) - if toc_physical_end is None: - return None, {} - - scan_start = toc_physical_end + 1 - scan_end = min(scan_start + _CALIBRATION_WINDOW_PAGES - 1, page_count) - if scan_start > page_count: - return None, {} - - leaves = list(iter_leaf_title_nodes(nodes)) - probe_leaves = [ - (path_titles, node) - for path_titles, node in leaves - if node.printed_page is not None - ][:_CALIBRATION_LEAF_PROBE_COUNT] - - if not probe_leaves: - return None, {} - - scan_pages = list(range(scan_start, scan_end + 1)) - candidates = [ - TitleMatch( - page=page, - confidence=0.4, - source="agent_heuristic", - matched_line="", - score=0.4, - candidates=scan_pages, - evidence={"calibration_probe": True}, - ) - for page in scan_pages - ] - - for path_titles, node in probe_leaves: - result = verify_section_page_choice( - ctx=ctx, - title=node.title, - candidate_matches=candidates, - candidate_page_cap=len(scan_pages), - ) - selected = result.get("selected_page") - if selected is not None and result.get("confidence", 0) >= 0.6: - offset = selected - node.printed_page - match = TitleMatch( - page=selected, - confidence=result.get("confidence", 0.75), - source="agent_vlm", - matched_line="", - score=result.get("confidence", 0.75), - candidates=[selected], - evidence={ - "calibration": True, - "printed_page": node.printed_page, - "reason": result.get("reason", ""), - }, - ) - logger.info( - "[page_memory.skeleton] calibration confirmed: title={!r} " - "printed_page={} physical_page={} offset={}", - node.title, - node.printed_page, - selected, - offset, - ) - return offset, {path_titles: match} - - logger.info("[page_memory.skeleton] calibration: no leaf confirmed in scan window") - return None, {} - - -def _toc_cluster_end_page(toc_hierarchies: list[dict[str, Any]] | None) -> int | None: - """Get the last physical page of the primary TOC cluster.""" - if not toc_hierarchies: - return None - end_pages: list[int] = [] - for hierarchy in toc_hierarchies: - end = _toc_range_end(hierarchy) - if end is not None: - end_pages.append(end) - return max(end_pages) if end_pages else None - - -# ── Offset-guided bulk anchoring with recursive recalibrate (Phase A3) ─────── - -_TAIL_VERIFY_CONFIDENCE_THRESHOLD = 0.6 -_MAX_RECALIBRATE_DEPTH = 5 -_MAX_RECALIBRATE_DELTA = 5 - - -def _verify_offset_tail( - *, - leaves: list[tuple[tuple[str, ...], TitleNode]], - offset: int, - ctx: ToolContext, - page_count: int, -) -> bool: - """VLM-verify that the offset holds for the last leaf entry (Theorem 1). - - If head offset == tail offset, monotonicity guarantees all intermediate - entries share the same offset. - - Prefers a tail leaf whose expected page is strictly less than page_count - (boundary pages are unreliable for VLM verification). - """ - tail_leaves = [ - (path, node) for path, node in reversed(leaves) if node.printed_page is not None - ] - if not tail_leaves: - return True - - # Prefer non-boundary: printed_page + offset < page_count - selected = None - for path, node in tail_leaves: - pp = node.printed_page - if pp is None: - continue - expected = pp + offset - if 1 <= expected < page_count: - selected = (path, node) - break - if selected is None: - # All leaves are at the boundary; fall back to the last one - selected = tail_leaves[0] - - path, node = selected - printed_page = node.printed_page - if printed_page is None: - return True - expected_page = printed_page + offset - if expected_page < 1 or expected_page > page_count: - return False - - candidate = TitleMatch( - page=expected_page, - confidence=0.4, - source="agent_heuristic", - matched_line="", - score=0.4, - candidates=[expected_page], - evidence={"tail_verify_probe": True}, - ) - result = verify_section_page_choice( - ctx=ctx, - title=node.title, - candidate_matches=[candidate], - candidate_page_cap=1, - ) - confirmed = ( - result.get("selected_page") == expected_page - and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD - ) - logger.info( - "[page_memory.skeleton] tail verify: title={!r} expected_page={} confirmed={} confidence={}", - node.title, - expected_page, - confirmed, - result.get("confidence", 0), - ) - return confirmed - - -def _vlm_confirm_single_page( - *, - ctx: ToolContext, - title: str, - expected_page: int, - page_count: int, -) -> bool: - """Single-page VLM confirmation for binary search steps.""" - if expected_page < 1 or expected_page > page_count: - return False - candidate = TitleMatch( - page=expected_page, - confidence=0.4, - source="agent_heuristic", - matched_line="", - score=0.4, - candidates=[expected_page], - evidence={"bisect_probe": True}, - ) - result = verify_section_page_choice( - ctx=ctx, - title=title, - candidate_matches=[candidate], - candidate_page_cap=1, - ) - return ( - result.get("selected_page") == expected_page - and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD - ) - - -def _bisect_offset_breakpoint( - *, - leaves: list[tuple[tuple[str, ...], TitleNode]], - offset: int, - ctx: ToolContext, - page_count: int, -) -> int: - """Binary search for the last leaf index where offset is valid. O(log n) VLM calls.""" - lo, hi = 0, len(leaves) - 1 - while lo < hi: - mid = (lo + hi + 1) // 2 - _, node = leaves[mid] - if node.printed_page is None: - hi = mid - 1 - continue - expected = node.printed_page + offset - if _vlm_confirm_single_page( - ctx=ctx, title=node.title, expected_page=expected, page_count=page_count - ): - lo = mid - else: - hi = mid - 1 - logger.info( - "[page_memory.skeleton] bisect breakpoint: last_valid_index={} / total={}", - lo, - len(leaves), - ) - return lo - - -def _bulk_offset_matches( - leaves: list[tuple[tuple[str, ...], TitleNode]], - offset: int, -) -> dict[tuple[str, ...], TitleMatch]: - """Generate TitleMatch overrides for all leaves using offset. No VLM calls.""" - matches: dict[tuple[str, ...], TitleMatch] = {} - for path_titles, node in leaves: - if node.printed_page is None: - continue - page = node.printed_page + offset - matches[path_titles] = TitleMatch( - page=page, - confidence=0.88, - source="agent_vlm", - matched_line="", - score=0.88, - candidates=[page], - evidence={ - "bulk_offset": True, - "offset": offset, - "printed_page": node.printed_page, - }, - ) - return matches - - -def _recalibrate_after_breakpoint( - *, - entry_node: TitleNode, - old_offset: int, - ctx: ToolContext, - page_count: int, -) -> int | None: - """Probe offsets old_offset+1, +2, ... to find new offset after breakpoint. - - Monotonicity guarantees new offset > old offset, so search space is tiny. - """ - entry_printed_page = entry_node.printed_page - if entry_printed_page is None: - return None - for delta in range(1, _MAX_RECALIBRATE_DELTA + 1): - new_offset = old_offset + delta - if _vlm_confirm_single_page( - ctx=ctx, - title=entry_node.title, - expected_page=entry_printed_page + new_offset, - page_count=page_count, - ): - logger.info( - "[page_memory.skeleton] recalibrate: title={!r} new_offset={} (delta=+{})", - entry_node.title, - new_offset, - delta, - ) - return new_offset - return None - - -def _offset_guided_anchoring( - *, - nodes: list[TitleNode], - offset: int, - ctx: ToolContext, - page_count: int, - calibration_overrides: dict[tuple[str, ...], TitleMatch], -) -> dict[tuple[str, ...], TitleMatch] | None: - """Offset-guided bulk anchoring with recursive recalibrate on breakpoints. - - Strategy: - 1. Tail verify last leaf with current offset - 2. If pass → bulk apply all leaves (Theorem 1) - 3. If fail → binary search for breakpoint - 4. Bulk apply leaves before breakpoint - 5. Recalibrate: probe remaining[0] with offset+1, +2, ... (monotonicity) - 6. Recurse on remaining segment with new offset - 7. If recalibrate fails → return partial (caller falls back for remainder) - - Returns match_overrides for all anchored leaves, or None for full fallback. - """ - leaves = [ - (path, node) - for path, node in iter_leaf_title_nodes(nodes) - if node.printed_page is not None - ] - if len(leaves) < 2: - return None - - all_matches: dict[tuple[str, ...], TitleMatch] = {} - all_matches.update(calibration_overrides) - - _anchor_segment_recursive( - leaves=leaves, - offset=offset, - ctx=ctx, - page_count=page_count, - matches=all_matches, - depth=0, - ) - - if not all_matches: - return None - - logger.info( - "[page_memory.skeleton] offset bulk anchoring: {} / {} leaves anchored", - len(all_matches), - len(leaves), - ) - return all_matches - - -def _anchor_segment_recursive( - *, - leaves: list[tuple[tuple[str, ...], TitleNode]], - offset: int, - ctx: ToolContext, - page_count: int, - matches: dict[tuple[str, ...], TitleMatch], - depth: int, -) -> None: - """Recursively anchor a segment of leaves, handling multiple breakpoints.""" - if not leaves or depth >= _MAX_RECALIBRATE_DEPTH: - return - - if _verify_offset_tail(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count): - bulk = _bulk_offset_matches(leaves, offset) - matches.update(bulk) - return - - bp = _bisect_offset_breakpoint(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count) - confirmed_leaves = leaves[: bp + 1] - if confirmed_leaves: - bulk = _bulk_offset_matches(confirmed_leaves, offset) - matches.update(bulk) - - remaining = leaves[bp + 1:] - if not remaining: - return - - _, first_remaining_node = remaining[0] - new_offset = _recalibrate_after_breakpoint( - entry_node=first_remaining_node, - old_offset=offset, - ctx=ctx, - page_count=page_count, - ) - if new_offset is None: - return - - _anchor_segment_recursive( - leaves=remaining, - offset=new_offset, - ctx=ctx, - page_count=page_count, - matches=matches, - depth=depth + 1, - ) # ── Multi-TOC grafting (Track B) ───────────────────────────────────────────── @@ -1133,11 +440,11 @@ def _resolve_pending_tocs( 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_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]) + 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 @@ -1146,7 +453,7 @@ def _resolve_pending_tocs( if p <= toc_scope_end and (toc_scope_start is None or p >= toc_scope_start) ] - offset, cal_overrides = _calibrate_offset_via_vlm( + offset, cal_overrides = calibrate_offset( nodes=nodes, toc_hierarchies=[pending_toc], ctx=ctx, @@ -1174,7 +481,7 @@ def _resolve_pending_tocs( ) continue - offset_matches = _offset_guided_anchoring( + offset_matches = offset_guided_anchoring( nodes=nodes, offset=offset, ctx=ctx, diff --git a/apps/worker/tests/contract/test_structure_anchoring_contract.py b/apps/worker/tests/contract/test_structure_anchoring_contract.py new file mode 100644 index 000000000..84261f4fc --- /dev/null +++ b/apps/worker/tests/contract/test_structure_anchoring_contract.py @@ -0,0 +1,200 @@ +"""Contract tests for structure_anchoring (moved from skeleton_extractor).""" + +from __future__ import annotations + +import os +from typing import Any +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 ToolContext +from app.services.document_agent.state import AgentBlackboard +from app.services.document_agent.structure.hierarchy_locator import TitleMatch, TitleNode +from app.services.document_agent.structure import structure_anchoring as anchoring + + +def _ctx() -> ToolContext: + return ToolContext( + pdf_path="/tmp/doc.pdf", + job_id="job-anchor", + blackboard=AgentBlackboard(), + budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), + trace=None, + settings={}, + ) + + +def _leaf(title: str, page: int) -> TitleNode: + return TitleNode(title=title, level=1, printed_page=page, children=[]) + + +def test_prune_out_of_scope_nodes_removes_overflow_leaves() -> None: + nodes = [ + _leaf("A", 1), + _leaf("B", 50), + ] + pruned, removed = anchoring.prune_out_of_scope_nodes( + nodes, offset=0, page_count=10 + ) + assert removed == 1 + assert [n.title for n in pruned] == ["A"] + + +def test_null_page_parent_skipped_without_right_anchor() -> None: + parent = TitleNode( + title="Chapter", + level=1, + printed_page=None, + children=[TitleNode(title="Orphan", level=2, printed_page=None, children=[])], + ) + overrides, report = anchoring.locate_null_page_parent_overrides( + nodes=[parent], + match_overrides={}, + page_texts={1: "Chapter\nHello"}, + body_pages=[1, 2, 3], + ctx=None, + ) + assert overrides == {} + assert len(report) == 1 + assert report[0]["result"] == "skipped_no_right" + + +def test_null_page_parent_located_via_compact_text() -> None: + child = TitleNode(title="1.1 Detail", level=2, printed_page=5, children=[]) + parent = TitleNode( + title="1 Overview", + level=1, + printed_page=None, + children=[child], + ) + leaf_match = anchoring.bulk_offset_matches( + [(("1 Overview", "1.1 Detail"), child)], + offset=0, + ) + page_texts = { + 4: "noise", + 5: "1 Overview\n1.1 Detail\nbody", + 6: "more", + } + overrides, report = anchoring.locate_null_page_parent_overrides( + nodes=[parent], + match_overrides=leaf_match, + page_texts=page_texts, + body_pages=[4, 5, 6], + ctx=None, + ) + assert ("1 Overview",) in overrides + assert overrides[("1 Overview",)].page == 5 + assert report[0]["result"] != "unresolved" + assert report[0]["page"] == 5 + + +def test_calibrate_and_bulk_via_mocked_offset() -> None: + leaves = [ + _leaf("Intro", 3), + _leaf("Body", 10), + _leaf("End", 20), + ] + ctx = _ctx() + seed = { + ("Intro",): TitleMatch( + page=5, + confidence=0.9, + source="agent_vlm", + matched_line="", + score=0.9, + candidates=[5], + evidence={"calibration": True, "printed_page": 3}, + ) + } + + def fake_verify(**kwargs: Any) -> dict[str, Any]: + expected = kwargs["candidate_matches"][0].page + return {"selected_page": expected, "confidence": 0.9, "reason": "ok"} + + with ( + patch.object( + anchoring, + "calibrate_offset", + return_value=(2, seed), + ), + patch.object( + anchoring, + "verify_section_page_choice", + side_effect=fake_verify, + ), + ): + offset, seed_overrides = anchoring.calibrate_offset( + nodes=leaves, + toc_hierarchies=[{"toc_range": [1, 2], "toc_tree": {}}], + ctx=ctx, + page_texts={}, + page_count=30, + ) + assert offset == 2 + assert seed_overrides + matches = anchoring.offset_guided_anchoring( + nodes=leaves, + offset=offset, + ctx=ctx, + page_count=30, + calibration_overrides=seed_overrides, + ) + assert matches is not None + assert len(matches) >= 3 + assert matches[("Intro",)].page == 5 + assert matches[("Body",)].page == 12 + assert matches[("End",)].page == 22 + + +def test_anchor_hierarchy_returns_skeleton_anchor_fields() -> None: + leaves = [_leaf("Only", 2)] + toc_hierarchies = [{"toc_range": [1, 1], "toc_tree": {}}] + ctx = _ctx() + seed = { + ("Only",): TitleMatch( + page=4, + confidence=0.95, + source="agent_vlm", + matched_line="", + score=0.95, + candidates=[4], + evidence={"calibration": True}, + ) + } + + def fake_verify(**kwargs: Any) -> dict[str, Any]: + expected = kwargs["candidate_matches"][0].page + return {"selected_page": expected, "confidence": 0.95, "reason": "ok"} + + with ( + patch.object(anchoring, "calibrate_offset", return_value=(2, seed)), + patch.object( + anchoring, + "verify_section_page_choice", + side_effect=fake_verify, + ), + ): + nodes, anchor = anchoring.anchor_hierarchy( + nodes=leaves, + toc_hierarchies=toc_hierarchies, + page_texts={4: "Only\ntext"}, + body_pages=[2, 3, 4, 5], + page_count=10, + ctx=ctx, + ) + assert isinstance(anchor, anchoring.SkeletonAnchor) + assert anchor.offset == 2 + assert anchor.offset_status == "ok" + assert isinstance(anchor.match_overrides, dict) + assert isinstance(anchor.null_page_report, list) + assert isinstance(anchor.bulk_count, int) + assert isinstance(anchor.pruned_count, int) + assert nodes From 6a716423922d38930e8e5be2a3946bba9527bfd4 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 12 Aug 2026 17:27:15 +0800 Subject: [PATCH 04/12] feat: enhance calibration process in document agent - Updated budget parameters in `.env.example` files for both API and worker to support increased visual and calibration budgets. - Introduced a new calibration stage in the `BudgetStage` and updated the `BudgetTracker` to accommodate the new calibration budget logic. - Enhanced the `calibration` phase in various modules, including `loop.py`, `procedure.py`, and `service.py`, to improve the handling of offsets and visual stages. - Refactored the `inspect_pages` tool to clarify budget limits and ensure proper handling of visual stages during inspections. These changes improve the efficiency and accuracy of the calibration process within the document agent, allowing for better handling of visual budgets and offsets. --- apps/api/.env.example | 4 +- apps/worker/.env.example | 4 +- .../agents/calibration/SKILL.md | 55 +- .../document_agent/agents/calibration/loop.py | 68 ++- .../agents/calibration/procedure.py | 487 ++++++++++++++---- .../agents/calibration/service.py | 42 +- .../agents/calibration/tools.py | 1 + .../app/services/document_agent/budget.py | 2 + .../services/document_agent/coordinator.py | 8 +- .../structure/hierarchy_locator.py | 90 +++- .../structure/page_locate_agent.py | 4 +- .../structure/structure_anchoring.py | 188 ++++--- .../document_agent/tools/inspect_pages.py | 60 ++- .../tools/propose_shard_plan.py | 10 +- .../page_memory/skeleton_extractor.py | 59 +-- .../test_structure_anchoring_contract.py | 272 ++++++++-- 16 files changed, 995 insertions(+), 359 deletions(-) diff --git a/apps/api/.env.example b/apps/api/.env.example index 400d74747..8ab1c0aac 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -113,13 +113,15 @@ OVERSIZED_PDF_SOFT_LIMIT=1500 PDF_PROFILE_TOC_ENABLED=false MINERU_SHARD_CONCURRENCY=3 PARSE_AGENT_PLAN_BUDGET=50000 -PARSE_AGENT_VISUAL_BUDGET=80000 +PARSE_AGENT_VISUAL_BUDGET=120000 PARSE_AGENT_TOC_CONFIRM_MIN_BUDGET=8000 PARSE_AGENT_TOC_CONFIRM_CAP=24000 PARSE_AGENT_COARSE_PLANNER_MIN_BUDGET=12000 PARSE_AGENT_COARSE_PLANNER_CAP=36000 PARSE_AGENT_STRUCTURAL_REACT_MIN_BUDGET=24000 PARSE_AGENT_STRUCTURAL_REACT_CAP=64000 +PARSE_AGENT_CALIBRATION_MIN_BUDGET=12000 +PARSE_AGENT_CALIBRATION_CAP=40000 PARSE_AGENT_PAGE_TAGGING_MIN_BUDGET=0 PARSE_AGENT_PAGE_TAGGING_CAP=0 diff --git a/apps/worker/.env.example b/apps/worker/.env.example index 3742d1e30..7c140aa80 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -130,13 +130,15 @@ PDF_PROFILE_TOC_ENABLED=false # JAVA_HOME=/opt/homebrew/opt/java MINERU_SHARD_CONCURRENCY=3 PARSE_AGENT_PLAN_BUDGET=50000 -PARSE_AGENT_VISUAL_BUDGET=80000 +PARSE_AGENT_VISUAL_BUDGET=120000 PARSE_AGENT_TOC_CONFIRM_MIN_BUDGET=8000 PARSE_AGENT_TOC_CONFIRM_CAP=24000 PARSE_AGENT_COARSE_PLANNER_MIN_BUDGET=12000 PARSE_AGENT_COARSE_PLANNER_CAP=36000 PARSE_AGENT_STRUCTURAL_REACT_MIN_BUDGET=24000 PARSE_AGENT_STRUCTURAL_REACT_CAP=64000 +PARSE_AGENT_CALIBRATION_MIN_BUDGET=12000 +PARSE_AGENT_CALIBRATION_CAP=40000 PARSE_AGENT_PAGE_TAGGING_MIN_BUDGET=0 PARSE_AGENT_PAGE_TAGGING_CAP=0 PARSE_AGENT_PAGE_LOCATE_MIN_BUDGET=0 diff --git a/apps/worker/app/services/document_agent/agents/calibration/SKILL.md b/apps/worker/app/services/document_agent/agents/calibration/SKILL.md index e1629b004..9910ac271 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/SKILL.md +++ b/apps/worker/app/services/document_agent/agents/calibration/SKILL.md @@ -13,6 +13,8 @@ are usable for coarse structure; unrecognized pages are treated as **no TOC**. - Do not scan a fixed window after the TOC (the old “TOC end + N pages” probe). - Do not invent physical pages you did not inspect or obtain from `link`. +- Do not track a total page-count budget. Limits are **token budgets** and + **max_rounds** in the payload. A per-call page cap is only a batch-size limit. ## Mandatory first step — partition regimes @@ -36,28 +38,41 @@ For each regime: - If the entry has `link.physical_page`, use it as the primary candidate. - Otherwise derive a coarse physical candidate from the printed label and `page_count`, then confirm with vision. -3. Call `inspect.pages` to confirm the heading starts on that page and to read - the folio/printed label when useful. -4. If wrong, inspect nearby physical pages and revise. -5. Compute `offset = physical - printed` using this regime’s interpretation of +3. Progressive `inspect.pages` for that title (start small, expand only if needed): + - **1st call**: inspect **1** candidate page only. + - **2nd call** (if miss): inspect up to **3** nearby pages. + - **3rd call** (if still miss): inspect up to **5** nearby pages. + Never open with a full 5-page batch when a single page has not been tried. +4. Compute `offset = physical - printed` using this regime’s interpretation of the printed label. -6. Submit **candidate** offsets. Do not treat Phase 1 alone as a finished +5. Submit **candidate** offsets. Do not treat Phase 1 alone as a finished coarse-structure calibration. -## Phase 2 — Completion (deterministic after submit; production path) +If `inspect.pages` returns budget exhausted, or rounds run out before a reliable +offset: treat that sample / regime as **not found**, submit whatever regimes you +already confirmed (or `status=failed`), and let production fallback handle the +rest. Do not guess pages. -For each TOC region with a candidate primary offset (prefer decimal): +## Phase 2 — Completion (deterministic after submit; production path) -1. Build TitleNodes via production `extract_toc_nodes` (integer `page_number` - only; roman / prefixed labels become `printed_page=None`). -2. Run production `anchor_hierarchy_from_offset`: - prune → tail verify → binary-search breakpoint → small-step recalibrate → - null-page parent locate. -3. Emit production `SkeletonAnchor` (`offset`, `offset_status`, - `match_overrides`, `null_page_report`, `bulk_count`, `pruned_count`, - `locate_agent`). -4. On recalibrate/budget failure: keep the complete **prefix**; mark only the - unresolved **suffix** as no TOC. Never fall back to a fixed post-TOC window. +For each TOC region, every regime with a candidate offset is completed +independently, then merged by **physical page**: + +1. Build TitleNodes via production `extract_toc_nodes` (regime-aware parse: + decimal / roman / prefixed labels → `printed_page` + `page_kind`). +2. For **each** regime with an offset: + - Project leaves belonging to that regime + - Run production Phase-2: prune → tail verify → binary-search → + small-step recalibrate (single-leaf regimes apply offset directly) +3. Merge all regime `match_overrides` (physical pages), then null-page parent + locate once on the combined tree. +4. Emit production `SkeletonAnchor` (`offset` = primary decimal summary, + `match_overrides` = union of all regimes, `null_page_report`, `bulk_count`, + `pruned_count`, `locate_agent`). +5. On recalibrate/budget failure inside one regime: keep that regime's complete + **prefix**; **drop** unresolved **suffix** leaves from the TOC tree (no TOC), + then run null-page parent locate on what remains. Never fall back to a fixed + post-TOC window. ## Usability bar @@ -68,7 +83,9 @@ For each TOC region with a candidate primary offset (prefer decimal): ## Tools - `inspect.pages`: primary tool for Phase 1. Open physical pages, render, answer - your question. Prefer batching related pages when the same question applies. + your question. Prefer the progressive 1→3→5 schedule above. Per-call page + count is capped; overall spend is limited by the calibration visual token + budget and `max_rounds`. - `calibration.submit`: finish Phase 1. Pass the full result under `tool_args.result` (or result fields directly in `tool_args`). @@ -81,4 +98,4 @@ For each TOC region with a candidate primary offset (prefer decimal): known), and `posterior` if you already inspected a late check. - Keep `kind` values consistent within one run (`decimal`, `roman`, `prefixed`, or `other`). -- Stay within the tool/round budget announced in the payload. +- Stay within the token / round budgets announced in the payload. diff --git a/apps/worker/app/services/document_agent/agents/calibration/loop.py b/apps/worker/app/services/document_agent/agents/calibration/loop.py index 5b3a99af9..852c9991f 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/loop.py +++ b/apps/worker/app/services/document_agent/agents/calibration/loop.py @@ -22,7 +22,7 @@ CalibrationResult, calibration_result_from_dict, ) -from app.services.document_agent.budget import BudgetTracker +from app.services.document_agent.budget import BudgetTracker, StageEnvelope from app.services.document_agent.manifest import ToolContext, ToolResult from app.services.document_agent.state import AgentBlackboard from app.services.document_agent.structure.structure_anchoring import ( @@ -44,10 +44,26 @@ inspect.pages, then call calibration.submit. Phase 2 (tail verify, binary search, small-step recalibrate) runs automatically after submit. Do not use a fixed post-TOC page window. +Hard limits are token budgets and max_rounds — not a total page-count ledger. Include the word json in your response. """.strip() +def _default_calibration_budget() -> BudgetTracker: + return BudgetTracker( + plan_budget=int(os.environ.get("PARSE_AGENT_PLAN_BUDGET", "50000")), + visual_budget=int(os.environ.get("PARSE_AGENT_VISUAL_BUDGET", "120000")), + visual_stage_envelopes={ + "calibration": StageEnvelope( + min_guarantee=int( + os.environ.get("PARSE_AGENT_CALIBRATION_MIN_BUDGET", "12000") + ), + cap=int(os.environ.get("PARSE_AGENT_CALIBRATION_CAP", "40000")), + ), + }, + ) + + def _load_skill() -> str: return _SKILL_PATH.read_text(encoding="utf-8") @@ -96,12 +112,14 @@ def run_calibration_phase1( no_links: bool = False, max_rounds: int = 16, inspect_page_cap: int = 5, - inspect_page_budget: int = 24, ) -> CalibrationResult: """Agent Phase-1 only: partition regimes + candidate offsets, then submit. Reuses the caller's ``ToolContext`` (budget / pdf / settings). Does **not** run production Phase-2 bulk anchoring. + + Hard limits: planner/visual token budgets + ``max_rounds``. Per-call + ``inspect_page_cap`` is only a batch-size cap (not a total page ledger). """ hierarchies = list(toc_hierarchies or []) if no_links: @@ -116,7 +134,7 @@ def run_calibration_phase1( ctx.blackboard.page_count = resolved_page_count ctx.settings.setdefault("inspect_page_cap", inspect_page_cap) - ctx.settings.setdefault("inspect_page_budget", inspect_page_budget) + ctx.settings.setdefault("inspect_visual_stage", "calibration") blackboard = ctx.blackboard blackboard.global_signals["calibration_region_index"] = region_index @@ -133,6 +151,15 @@ def run_calibration_phase1( for round_index in range(max_rounds): available = registry.openai_specs(blackboard) + snap = ctx.budget.snapshot() if ctx.budget is not None else {} + visual_stages = ( + snap.get("visual_stages") if isinstance(snap, dict) else {} + ) or {} + calib_stage = ( + visual_stages.get("calibration") + if isinstance(visual_stages, dict) + else None + ) payload = { "skill": skill, "page_count": resolved_page_count, @@ -144,10 +171,10 @@ def run_calibration_phase1( "inspect_page_cap_per_call": int( ctx.settings.get("inspect_page_cap") or inspect_page_cap ), - "inspect_page_budget_total": int( - ctx.settings.get("inspect_page_budget") or inspect_page_budget - ), - "inspect_pages_used": blackboard.global_signals.get( + "calibration_visual": calib_stage, + "plan": snap.get("plan") if isinstance(snap, dict) else None, + "visual": snap.get("visual") if isinstance(snap, dict) else None, + "inspect_pages_used_diagnostic": blackboard.global_signals.get( "calibration_inspect_pages_used" ), }, @@ -249,6 +276,21 @@ def run_calibration_phase1( tool_result.status, ) + if tool_result.status == "error" and "budget exhausted" in str( + tool_result.error or "" + ).lower(): + return _attach_history( + CalibrationResult( + status="failed", + notes=f"budget exhausted: {tool_result.error}", + region_index=region_index, + tool_calls=int( + blackboard.global_signals.get("calibration_tool_calls") or 0 + ), + ), + history, + ) + if blackboard.global_signals.get("calibration_done"): raw_result = blackboard.global_signals.get("calibration_result") or {} if isinstance(raw_result, dict): @@ -282,7 +324,6 @@ def run_calibration_agent( no_links: bool = False, max_rounds: int = 16, inspect_page_cap: int = 5, - inspect_page_budget: int = 24, budget: BudgetTracker | None = None, page_texts: dict[int, str] | None = None, body_pages: list[int] | None = None, @@ -305,18 +346,14 @@ def run_calibration_agent( pdf_path=pdf_path, job_id=f"calibration-region-{region_index}", blackboard=blackboard, - budget=budget - or BudgetTracker( - plan_budget=int(os.environ.get("PARSE_AGENT_PLAN_BUDGET", "50000")), - visual_budget=int(os.environ.get("PARSE_AGENT_VISUAL_BUDGET", "80000")), - ), + budget=budget or _default_calibration_budget(), trace=None, output_dir=output_dir, settings={ "vlm_model": vlm_model or "", "model": planner_model or vlm_model or "", "inspect_page_cap": inspect_page_cap, - "inspect_page_budget": inspect_page_budget, + "inspect_visual_stage": "calibration", }, ) @@ -328,12 +365,11 @@ def run_calibration_agent( no_links=False, # already stripped above when requested max_rounds=max_rounds, inspect_page_cap=inspect_page_cap, - inspect_page_budget=inspect_page_budget, ) if phase1.status == "failed" and not phase1.regimes: return phase1, {} - anchor, finalized = finalize_calibration_result( + _working, anchor, finalized = finalize_calibration_result( result=phase1, entries=list(region_payload.get("entries") or []), toc_hierarchies=region_hierarchies, diff --git a/apps/worker/app/services/document_agent/agents/calibration/procedure.py b/apps/worker/app/services/document_agent/agents/calibration/procedure.py index bcd8fb096..a27126f63 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/procedure.py +++ b/apps/worker/app/services/document_agent/agents/calibration/procedure.py @@ -1,17 +1,17 @@ """Phase-2 completion aligned with production anchoring. After the agent submits candidate regime offsets, this module: -1. Picks the primary (usually decimal) candidate offset -2. Builds TitleNodes the same way production does (``extract_toc_nodes``) -3. Runs ``anchor_hierarchy_from_offset`` (prune → bulk/bisect → null-page) +1. Builds TitleNodes the same way production does +2. Runs Phase-2 **per regime** (prune → bulk/bisect → recalibrate) +3. Merges physical-page ``match_overrides`` across regimes +4. Runs null-page parent locate once on the combined tree -The returned ``SkeletonAnchor`` is the production schema swap point. -Regime metadata is retained only as experiment diagnostics. +Returns production ``SkeletonAnchor`` plus regime diagnostics for debug payloads. """ from __future__ import annotations -import re +from dataclasses import replace from typing import Any from loguru import logger @@ -24,79 +24,23 @@ from app.services.document_agent.manifest import ToolContext from app.services.document_agent.structure.hierarchy_locator import ( TitleMatch, + TitleNode, + classify_page_number_kind, extract_toc_nodes, iter_leaf_title_nodes, + normalize_page_kind, + parse_printed_page, ) from app.services.document_agent.structure.structure_anchoring import ( SkeletonAnchor, - anchor_hierarchy_from_offset, + locate_null_page_parent_overrides, + offset_guided_anchoring, + prune_unanchored_toc_leaves, serialize_skeleton_anchor, ) - -_ROMAN_MAP = {"i": 1, "v": 5, "x": 10, "l": 50, "c": 100, "d": 500, "m": 1000} - - -def classify_page_number_kind(label: Any) -> str: - text = str(label or "").strip() - if not text: - return "other" - if re.fullmatch(r"\d+", text): - return "decimal" - if re.fullmatch(r"[ivxlcdm]+", text, flags=re.IGNORECASE): - return "roman" - if re.fullmatch(r"[A-Za-z]+-\d+", text): - return "prefixed" - return "other" - - -def parse_printed_page(label: Any, *, kind: str) -> int | None: - text = str(label or "").strip() - if not text: - return None - kind_l = (kind or "").lower() - if kind_l in {"decimal", "arabic", "arabic_digits"}: - return int(text) if text.isdigit() else None - if kind_l == "roman": - return _roman_to_int(text) - if kind_l in {"prefixed", "folio"}: - match = re.fullmatch(r"[A-Za-z]+-(\d+)", text) - return int(match.group(1)) if match else None - if text.isdigit(): - return int(text) - if re.fullmatch(r"[ivxlcdm]+", text, flags=re.IGNORECASE): - return _roman_to_int(text) - match = re.fullmatch(r"[A-Za-z]+-(\d+)", text) - return int(match.group(1)) if match else None - - -def _roman_to_int(text: str) -> int | None: - raw = text.strip().lower() - if not raw or not re.fullmatch(r"[ivxlcdm]+", raw): - return None - total = 0 - prev = 0 - for ch in reversed(raw): - value = _ROMAN_MAP.get(ch) - if value is None: - return None - if value < prev: - total -= value - else: - total += value - prev = value - return total if total > 0 else None - - -def normalize_kind(kind: str) -> str: - text = (kind or "other").strip().lower() - if text in {"arabic", "arabic_digits", "decimal"}: - return "decimal" - if text == "roman": - return "roman" - if text in {"prefixed", "folio"}: - return "prefixed" - return text or "other" +# Re-export under prior names so existing imports keep working. +normalize_kind = normalize_page_kind def pick_primary_offset(result: CalibrationResult) -> int | None: @@ -112,7 +56,7 @@ def pick_primary_offset(result: CalibrationResult) -> int | None: return None -def _seed_overrides_from_samples( +def seed_overrides_from_samples( *, result: CalibrationResult, nodes: list[Any], @@ -129,7 +73,6 @@ def _seed_overrides_from_samples( continue path = title_to_path.get(sample.title.strip()) if path is None: - # Soft match: normalized equality needle = sample.title.strip().lower() for title, candidate in title_to_path.items(): if title.lower() == needle: @@ -154,16 +97,320 @@ def _seed_overrides_from_samples( return overrides +def _iter_all_title_nodes( + nodes: list[TitleNode], + *, + parent_titles: tuple[str, ...] = (), +) -> list[tuple[tuple[str, ...], TitleNode]]: + rows: list[tuple[tuple[str, ...], TitleNode]] = [] + for node in nodes: + path = (*parent_titles, node.title) + rows.append((path, node)) + if node.children: + rows.extend( + _iter_all_title_nodes(node.children, parent_titles=path) + ) + return rows + + +def flat_toc_entries(toc_hierarchies: list[dict[str, Any]] | None) -> list[dict[str, Any]]: + entries: list[dict[str, Any]] = [] + for hierarchy in toc_hierarchies or []: + raw = hierarchy.get("toc_with_level") if isinstance(hierarchy, dict) else None + if isinstance(raw, list): + entries.extend(e for e in raw if isinstance(e, dict)) + return entries + + +def _entry_titles_for_regime( + *, + regime: CalibrationRegime, + entries: list[dict[str, Any]], +) -> set[str] | None: + """Titles belonging to this regime; None means fall back to page_kind match.""" + from app.services.document_parser.structure.body_boundary import ( + normalize_heading_text, + ) + + kind = normalize_kind(regime.kind) + indices = list(regime.entry_indices or []) + if not indices and entries: + indices = [ + idx + for idx, entry in enumerate(entries) + if classify_page_number_kind(entry.get("page_number")) == kind + ] + if not indices: + return None + titles: set[str] = set() + for idx in indices: + if idx < 0 or idx >= len(entries): + continue + heading = entries[idx].get("heading") + title = normalize_heading_text(str(heading or "")) + if title: + titles.add(title) + return titles or None + + +def _leaf_in_regime( + node: TitleNode, + *, + kind: str, + entry_titles: set[str] | None, +) -> bool: + if entry_titles is not None: + return node.title in entry_titles + leaf_kind = normalize_kind(node.page_kind or classify_page_number_kind(node.printed_label)) + return leaf_kind == kind + + +def project_nodes_for_regime( + nodes: list[TitleNode], + *, + kind: str, + entry_titles: set[str] | None, +) -> list[TitleNode]: + """Copy tree: only this regime's leaves keep a parsed ``printed_page``.""" + + def walk(node: TitleNode) -> TitleNode: + children = [walk(child) for child in node.children] + if node.children: + return replace(node, children=children) + if not _leaf_in_regime(node, kind=kind, entry_titles=entry_titles): + return replace(node, printed_page=None, children=[]) + printed = node.printed_page + if printed is None and node.printed_label is not None: + printed = parse_printed_page(node.printed_label, kind=kind) + return replace( + node, + printed_page=printed, + page_kind=kind, + children=[], + ) + + return [walk(node) for node in nodes] + + +def prune_regime_out_of_scope( + nodes: list[TitleNode], + *, + kind: str, + entry_titles: set[str] | None, + offset: int, + page_count: int, +) -> tuple[list[TitleNode], int]: + """Drop only this regime's leaves whose printed+offset falls outside the PDF.""" + removed = 0 + + def prune(node: TitleNode) -> TitleNode | None: + nonlocal removed + if not node.children: + if _leaf_in_regime(node, kind=kind, entry_titles=entry_titles): + printed = node.printed_page + if printed is None and node.printed_label is not None: + printed = parse_printed_page(node.printed_label, kind=kind) + if printed is not None: + expected = printed + offset + if expected < 1 or expected > page_count: + removed += 1 + return None + return replace(node, printed_page=printed) + return node + children: list[TitleNode] = [] + for child in node.children: + kept = prune(child) + if kept is not None: + children.append(kept) + if not children: + removed += 1 + return None + return replace(node, children=children) + + out: list[TitleNode] = [] + for node in nodes: + kept = prune(node) + if kept is not None: + out.append(kept) + return out, removed + + +def anchor_hierarchy_from_regimes( + *, + nodes: list[TitleNode], + result: CalibrationResult, + entries: list[dict[str, Any]] | None, + page_texts: dict[int, str], + body_pages: list[int], + page_count: int, + ctx: ToolContext | None, +) -> tuple[list[TitleNode], SkeletonAnchor]: + """Phase-2: per-regime offset bulk/bisect, then merge physical overrides.""" + flat_entries = list(entries or []) + seed = seed_overrides_from_samples(result=result, nodes=nodes) + merged: dict[tuple[str, ...], TitleMatch] = dict(seed) + working = nodes + total_pruned = 0 + regime_bulk = 0 + + usable_regimes = [ + regime + for regime in result.regimes + if regime.offset is not None + ] + if not usable_regimes and result.offset is not None: + usable_regimes = [ + CalibrationRegime( + kind="decimal", + offset=int(result.offset), + offset_status="ok", + ) + ] + + for regime in usable_regimes: + kind = normalize_kind(regime.kind) + offset = int(regime.offset) # type: ignore[arg-type] + entry_titles = _entry_titles_for_regime(regime=regime, entries=flat_entries) + working, pruned = prune_regime_out_of_scope( + working, + kind=kind, + entry_titles=entry_titles, + offset=offset, + page_count=page_count, + ) + total_pruned += pruned + if not working: + continue + + projected = project_nodes_for_regime( + working, kind=kind, entry_titles=entry_titles + ) + regime_paths = { + path + for path, node in iter_leaf_title_nodes(projected) + if node.printed_page is not None + } + regime_seed = { + path: match + for path, match in seed.items() + if path in regime_paths + } + + if ctx is None: + # Offline: still apply deterministic printed+offset for this regime. + from app.services.document_agent.structure.structure_anchoring import ( + bulk_offset_matches, + ) + + leaves = [ + (path, node) + for path, node in iter_leaf_title_nodes(projected) + if node.printed_page is not None + ] + if leaves: + matches = bulk_offset_matches(leaves, offset) + matches.update(regime_seed) + merged.update(matches) + regime_bulk += len(matches) + continue + + matches = offset_guided_anchoring( + nodes=projected, + offset=offset, + ctx=ctx, + page_count=page_count, + calibration_overrides=regime_seed, + ) + if matches: + merged.update(matches) + regime_bulk += len( + { + path + for path in matches + if path in regime_paths or path in regime_seed + } + ) + logger.info( + "[calibration.phase2] regime={} offset={} anchored={}", + kind, + offset, + len(matches), + ) + elif regime_seed: + merged.update(regime_seed) + logger.info( + "[calibration.phase2] regime={} offset={} seed_only={}", + kind, + offset, + len(regime_seed), + ) + + # Failed suffix / never-confirmed printed leaves → drop from TOC tree. + working, unanchored_removed = prune_unanchored_toc_leaves( + working, match_overrides=merged + ) + total_pruned += unanchored_removed + if working: + surviving_paths = { + path + for path, _node in _iter_all_title_nodes(working) + } + merged = { + path: match + for path, match in merged.items() + if path in surviving_paths + } + + match_overrides, null_page_report = locate_null_page_parent_overrides( + nodes=working, + match_overrides=merged, + page_texts=page_texts, + body_pages=body_pages, + ctx=ctx, + ) + + primary = pick_primary_offset(result) + if primary is None and usable_regimes: + primary = int(usable_regimes[0].offset) # type: ignore[arg-type] + + if primary is None: + offset_status = "failed" if ctx is not None else "skipped" + else: + offset_status = "ok" + + locate_agent = ( + "offset_guided_bulk" + if match_overrides and (regime_bulk > 0 or seed) + else "offset_only" + ) + bulk_count = len(match_overrides) + + return working, SkeletonAnchor( + offset=primary, + offset_status=offset_status, + match_overrides=match_overrides, + null_page_report=null_page_report, + bulk_count=bulk_count, + pruned_count=total_pruned, + locate_agent=locate_agent, + ) + + def _annotate_regimes_from_anchor( *, result: CalibrationResult, anchor: SkeletonAnchor, entries: list[dict[str, Any]], + nodes: list[TitleNode], ) -> list[CalibrationRegime]: """Attach production segment view onto agent regimes for diagnostics.""" + path_by_title = { + node.title: path for path, node in iter_leaf_title_nodes(nodes) + } out: list[CalibrationRegime] = [] for regime in result.regimes: kind = normalize_kind(regime.kind) + entry_titles = _entry_titles_for_regime(regime=regime, entries=entries) indices = list(regime.entry_indices or []) if not indices: indices = [ @@ -173,33 +420,50 @@ def _annotate_regimes_from_anchor( and classify_page_number_kind(entry.get("page_number")) == kind ] - # Production trees only integer-print leaves enter bulk; decimal regime - # maps directly onto SkeletonAnchor bulk when Phase-2 succeeded. - if ( - kind == "decimal" - and anchor.offset is not None - and int(anchor.bulk_count or 0) > 0 - ): - ok_indices = indices - no_toc: list[int] = [] + ok_indices: list[int] = [] + no_toc: list[int] = [] + for idx in indices: + if idx < 0 or idx >= len(entries): + continue + heading = str(entries[idx].get("heading") or "") + from app.services.document_parser.structure.body_boundary import ( + normalize_heading_text, + ) + + title = normalize_heading_text(heading) + path = path_by_title.get(title) + if path is not None and path in (anchor.match_overrides or {}): + ok_indices.append(idx) + else: + no_toc.append(idx) + + # Fallback: kind-matched leaves present in overrides. + if not ok_indices and entry_titles is None: + for path, node in iter_leaf_title_nodes(nodes): + if _leaf_in_regime(node, kind=kind, entry_titles=None) and path in ( + anchor.match_overrides or {} + ): + # No stable entry index — treat as complete via offset status. + ok_indices = indices + no_toc = [] + break + + segments: list[CalibrationSegment] = [] + if ok_indices and regime.offset is not None: segments = [ CalibrationSegment( - offset=int(anchor.offset), + offset=int(regime.offset), leaf_start=0, leaf_end=max(0, len(ok_indices) - 1), entry_indices=ok_indices, status="ok", ) ] - else: - ok_indices = [] - no_toc = list(indices) - segments = [] out.append( CalibrationRegime( kind=kind, - offset=anchor.offset if kind == "decimal" else regime.offset, + offset=regime.offset, offset_status="ok" if segments else "failed", entry_indices=indices, samples=list(regime.samples), @@ -208,7 +472,8 @@ def _annotate_regimes_from_anchor( no_toc_entry_indices=no_toc, notes=( f"production_bulk={anchor.bulk_count}; " - f"locate_agent={anchor.locate_agent}" + f"locate_agent={anchor.locate_agent}; " + f"regime_anchored={len(ok_indices)}" ), ) ) @@ -220,45 +485,45 @@ def finalize_calibration_result( result: CalibrationResult, entries: list[dict[str, Any]], toc_hierarchies: list[dict[str, Any]], - ctx: ToolContext, + ctx: ToolContext | None, page_count: int, page_texts: dict[int, str] | None = None, body_pages: list[int] | None = None, -) -> tuple[SkeletonAnchor, CalibrationResult]: - """Run production Phase-2 from an agent candidate offset.""" - offset_hint = pick_primary_offset(result) + nodes: list[TitleNode] | None = None, +) -> tuple[list[TitleNode], SkeletonAnchor, CalibrationResult]: + """Run production multi-regime Phase-2 from agent candidate offsets.""" texts = dict(page_texts or {}) bodies = list(body_pages or sorted(texts.keys()) or list(range(1, page_count + 1))) - # Same TitleNode prep as C4 / extract_section_skeletons before anchoring. - from app.services.page_memory.skeleton_extractor import ( - _collapse_intermediate_single_child_chains, - ) + if nodes is None: + from app.services.page_memory.skeleton_extractor import ( + _collapse_intermediate_single_child_chains, + ) - nodes = _collapse_intermediate_single_child_chains( - extract_toc_nodes(toc_hierarchies) - ) - seed = _seed_overrides_from_samples(result=result, nodes=nodes) + nodes = _collapse_intermediate_single_child_chains( + extract_toc_nodes(toc_hierarchies) + ) - working, anchor = anchor_hierarchy_from_offset( + working, anchor = anchor_hierarchy_from_regimes( nodes=nodes, - offset_hint=offset_hint, - calibration_overrides=seed, + result=result, + entries=entries or flat_toc_entries(toc_hierarchies), page_texts=texts, body_pages=bodies, page_count=page_count, ctx=ctx, ) logger.info( - "[calibration.completion] offset={} status={} bulk={} pruned={} nodes={}", + "[calibration.completion] offset={} status={} bulk={} pruned={} nodes={} regimes={}", anchor.offset, anchor.offset_status, anchor.bulk_count, anchor.pruned_count, len(working), + len(result.regimes), ) regimes = _annotate_regimes_from_anchor( - result=result, anchor=anchor, entries=entries + result=result, anchor=anchor, entries=entries, nodes=working ) complete = sum(len(r.segments) for r in regimes) notes_parts = [result.notes] if result.notes else [] @@ -276,7 +541,7 @@ def finalize_calibration_result( region_index=result.region_index, history_tail=list(result.history_tail), ) - return anchor, finalized + return working, anchor, finalized def build_calibration_payload( @@ -292,10 +557,10 @@ def build_calibration_payload( payload.update( { "status": result.status, - "regimes": [regime.to_dict() if hasattr(regime, "to_dict") else regime for regime in ( - # dataclasses asdict via CalibrationResult - result.to_dict().get("regimes") or [] - )], + "regimes": [ + regime.to_dict() if hasattr(regime, "to_dict") else regime + for regime in (result.to_dict().get("regimes") or []) + ], "regions": list(region_payloads or []), "tool_calls": int(tool_calls if tool_calls is not None else result.tool_calls), "notes": result.notes, diff --git a/apps/worker/app/services/document_agent/agents/calibration/service.py b/apps/worker/app/services/document_agent/agents/calibration/service.py index c6e688022..d1f743767 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/service.py +++ b/apps/worker/app/services/document_agent/agents/calibration/service.py @@ -1,7 +1,7 @@ """Production calibration entry: Agent Phase-1 offset discovery. -Same return shape as the former ``calibrate_offset_via_vlm`` so callers can -swap without changing prune / bulk / null-page. +Returns a full ``CalibrationResult`` (all regimes). Callers run multi-regime +Phase-2 via ``finalize_calibration_result`` / ``anchor_hierarchy``. """ from __future__ import annotations @@ -11,12 +11,9 @@ from loguru import logger from app.services.document_agent.agents.calibration.loop import run_calibration_phase1 -from app.services.document_agent.agents.calibration.procedure import ( - pick_primary_offset, - _seed_overrides_from_samples, -) +from app.services.document_agent.agents.calibration.types import CalibrationResult from app.services.document_agent.manifest import ToolContext -from app.services.document_agent.structure.hierarchy_locator import TitleMatch, TitleNode +from app.services.document_agent.structure.hierarchy_locator import TitleNode def calibrate_offset( @@ -26,17 +23,19 @@ def calibrate_offset( ctx: ToolContext | None, page_texts: dict[int, str], page_count: int, -) -> tuple[int | None, dict[tuple[str, ...], TitleMatch]]: - """Discover printed→physical offset via the calibration SubAgent (Phase 1). +) -> CalibrationResult: + """Discover printed→physical offsets via the calibration SubAgent (Phase 1). - Returns ``(offset, seed_overrides)``. Phase 2 (tail / bisect / null-page) - stays in ``anchor_hierarchy_from_offset`` using the caller's node tree. + Returns the full Phase-1 ``CalibrationResult`` including every regime the + agent submitted. Phase-2 (per-regime bulk / bisect / null-page merge) is + owned by ``finalize_calibration_result`` / ``anchor_hierarchy``. """ + del nodes # Phase-1 works from toc_hierarchies entries; nodes used in Phase-2. if ctx is None: - return None, {} + return CalibrationResult(status="failed", notes="ctx missing") hierarchies = list(toc_hierarchies or []) if not hierarchies: - return None, {} + return CalibrationResult(status="failed", notes="toc_hierarchies empty") if page_count and not ctx.blackboard.page_count: ctx.blackboard.page_count = int(page_count) @@ -52,17 +51,12 @@ def calibrate_offset( ) except Exception as exc: logger.warning("[calibration] Phase-1 failed: {}", exc) - return None, {} + return CalibrationResult(status="failed", notes=str(exc)) - offset = pick_primary_offset(phase1) - if offset is None: - logger.info("[calibration] Phase-1 produced no primary offset") - return None, {} - - seed = _seed_overrides_from_samples(result=phase1, nodes=nodes) logger.info( - "[calibration] Phase-1 offset={} seed_overrides={}", - offset, - len(seed), + "[calibration] Phase-1 status={} regimes={} primary_offset={}", + phase1.status, + len(phase1.regimes), + phase1.offset, ) - return offset, seed + return phase1 diff --git a/apps/worker/app/services/document_agent/agents/calibration/tools.py b/apps/worker/app/services/document_agent/agents/calibration/tools.py index a2d7c44b3..37cf26c6a 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/tools.py +++ b/apps/worker/app/services/document_agent/agents/calibration/tools.py @@ -70,6 +70,7 @@ def _calibration_inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolRe merged.setdefault("folder_name", "calibration_inspect") merged.setdefault("prefix", "calib") merged.setdefault("usage_task", "calibration.inspect_pages") + merged.setdefault("visual_stage", "calibration") return inspect_pages(ctx, merged) diff --git a/apps/worker/app/services/document_agent/budget.py b/apps/worker/app/services/document_agent/budget.py index d0c91abc9..86058e27b 100644 --- a/apps/worker/app/services/document_agent/budget.py +++ b/apps/worker/app/services/document_agent/budget.py @@ -11,6 +11,7 @@ "toc_confirm", "coarse_planner", "structural_react", + "calibration", "page_locate", "page_tagging", ] @@ -174,6 +175,7 @@ def fork(self, ratio: float) -> "BudgetTracker": usage = self._visual_stage_usage.get(stage, StageUsage()) remaining_guarantee = max(envelope.min_guarantee - usage.committed, 0) if stage == "page_locate": + # Legacy uncapped locate stage: keep child uncapped. child_envelopes[stage] = StageEnvelope(min_guarantee=0, cap=0) continue child_envelopes[stage] = StageEnvelope( diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 2f8c699fb..983089d8e 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -46,7 +46,7 @@ def __init__( self.blackboard = AgentBlackboard() self.budget = BudgetTracker( plan_budget=int(os.environ.get("PARSE_AGENT_PLAN_BUDGET", "50000")), - visual_budget=int(os.environ.get("PARSE_AGENT_VISUAL_BUDGET", "80000")), + visual_budget=int(os.environ.get("PARSE_AGENT_VISUAL_BUDGET", "120000")), visual_stage_envelopes={ "toc_confirm": StageEnvelope( min_guarantee=int( @@ -66,6 +66,12 @@ def __init__( ), cap=int(os.environ.get("PARSE_AGENT_STRUCTURAL_REACT_CAP", "64000")), ), + "calibration": StageEnvelope( + min_guarantee=int( + os.environ.get("PARSE_AGENT_CALIBRATION_MIN_BUDGET", "12000") + ), + cap=int(os.environ.get("PARSE_AGENT_CALIBRATION_CAP", "40000")), + ), "page_locate": StageEnvelope( min_guarantee=int( os.environ.get("PARSE_AGENT_PAGE_LOCATE_MIN_BUDGET", "0") diff --git a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py index f33717875..46cace860 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -25,6 +25,71 @@ ] +_ROMAN_MAP = {"i": 1, "v": 5, "x": 10, "l": 50, "c": 100, "d": 500, "m": 1000} + + +def classify_page_number_kind(label: Any) -> str: + text = str(label or "").strip() + if not text: + return "other" + if re.fullmatch(r"\d+", text): + return "decimal" + if re.fullmatch(r"[ivxlcdm]+", text, flags=re.IGNORECASE): + return "roman" + if re.fullmatch(r"[A-Za-z]+-\d+", text): + return "prefixed" + return "other" + + +def normalize_page_kind(kind: str) -> str: + text = (kind or "other").strip().lower() + if text in {"arabic", "arabic_digits", "decimal"}: + return "decimal" + if text == "roman": + return "roman" + if text in {"prefixed", "folio"}: + return "prefixed" + return text or "other" + + +def parse_printed_page(label: Any, *, kind: str) -> int | None: + text = str(label or "").strip() + if not text: + return None + kind_l = normalize_page_kind(kind) + if kind_l == "decimal": + return int(text) if text.isdigit() else None + if kind_l == "roman": + return _roman_to_int(text) + if kind_l == "prefixed": + match = re.fullmatch(r"[A-Za-z]+-(\d+)", text) + return int(match.group(1)) if match else None + if text.isdigit(): + return int(text) + if re.fullmatch(r"[ivxlcdm]+", text, flags=re.IGNORECASE): + return _roman_to_int(text) + match = re.fullmatch(r"[A-Za-z]+-(\d+)", text) + return int(match.group(1)) if match else None + + +def _roman_to_int(text: str) -> int | None: + raw = text.strip().lower() + if not raw or not re.fullmatch(r"[ivxlcdm]+", raw): + return None + total = 0 + prev = 0 + for ch in reversed(raw): + value = _ROMAN_MAP.get(ch) + if value is None: + return None + if value < prev: + total -= value + else: + total += value + prev = value + return total if total > 0 else None + + @dataclass(frozen=True) class PageRange: start: int @@ -41,6 +106,8 @@ class TitleNode: title: str level: int printed_page: int | None = None + printed_label: str | None = None + page_kind: str | None = None physical_page_hint: int | None = None children: list["TitleNode"] = field(default_factory=list) @@ -623,11 +690,16 @@ def _parse_markdown_toc_entries(markdown: str) -> list[dict[str, Any]]: level = _safe_int(row.get("level")) heading = row.get("heading") if heading and level: + raw_page = row.get("page_number") entries.append( { "heading": heading, "level": level, - "page_number": _safe_int(row.get("page_number")), + # Preserve raw label (roman / prefixed / decimal); parsing is + # regime-aware at TitleNode construction time. + "page_number": raw_page.strip() + if isinstance(raw_page, str) + else raw_page, } ) return entries @@ -657,10 +729,24 @@ def _entries_to_tree(entries: list[dict[str, Any]]) -> list[TitleNode]: level = _safe_int(entry.get("level")) or 1 if not title or len(title) < 2: continue + raw_label = entry.get("page_number") + printed_label = ( + None + if raw_label is None or raw_label == "" + else str(raw_label).strip() + ) + page_kind = classify_page_number_kind(printed_label) if printed_label else None + printed_page = ( + parse_printed_page(printed_label, kind=page_kind or "other") + if printed_label + else None + ) node = TitleNode( title=title, level=level, - printed_page=_safe_int(entry.get("page_number")), + printed_page=printed_page, + printed_label=printed_label, + page_kind=page_kind, ) while stack and stack[-1][0] >= level: stack.pop() diff --git a/apps/worker/app/services/document_agent/structure/page_locate_agent.py b/apps/worker/app/services/document_agent/structure/page_locate_agent.py index f6aa312d4..f16f02735 100644 --- a/apps/worker/app/services/document_agent/structure/page_locate_agent.py +++ b/apps/worker/app/services/document_agent/structure/page_locate_agent.py @@ -77,7 +77,7 @@ def verify_section_page_choice( prompt = _build_verify_prompt(title=title, candidates=candidates) est = 800 * len(rendered) + 800 - stage = "page_locate" + stage = "calibration" if not ctx.budget.try_reserve("visual", est, stage=stage): best = candidates[0] return { @@ -85,7 +85,7 @@ def verify_section_page_choice( "candidate_pages": pages, "confidence": min(best.confidence, BUDGET_EXHAUSTED_GREP_CONFIDENCE_CAP), "source": "agent_heuristic", - "reason": "page_locate visual budget exhausted; selected top grep candidate", + "reason": "calibration visual budget exhausted; selected top grep candidate", } content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt}] diff --git a/apps/worker/app/services/document_agent/structure/structure_anchoring.py b/apps/worker/app/services/document_agent/structure/structure_anchoring.py index d6232fc1a..24fb2d229 100644 --- a/apps/worker/app/services/document_agent/structure/structure_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/structure_anchoring.py @@ -1,7 +1,7 @@ -"""Shared hierarchy anchoring: offset calibrate, null-page locate, bulk apply. +"""Shared hierarchy anchoring: Phase-2 bulk/bisect/null-page + SkeletonAnchor. -Extracted from page_memory.skeleton_extractor so profile-time skeleton phase -and page-memory C4 share one implementation. No page_memory imports. +Phase-1 offset discovery lives in ``document_agent.agents.calibration``. +``anchor_hierarchy`` composes Phase-1 + Phase-2 for production callers. """ from __future__ import annotations @@ -23,6 +23,7 @@ ) from loguru import logger + def prune_out_of_scope_nodes( nodes: list[TitleNode], *, @@ -73,6 +74,57 @@ def _prune(node: TitleNode) -> TitleNode | None: return pruned, removed +def prune_unanchored_toc_leaves( + nodes: list[TitleNode], + *, + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> tuple[list[TitleNode], int]: + """Remove TOC leaves that have no physical ``match_overrides`` entry. + + Implements Phase-2 ``suffix = no TOC``: after bulk/bisect/recalibrate, any + leaf that was not successfully anchored is dropped from the coarse tree + instead of sticky ``inherited_unlocated`` ranges. Childless parents are + removed unless they themselves have an override. + """ + removed = 0 + + def _prune( + node: TitleNode, parent_titles: tuple[str, ...] + ) -> TitleNode | None: + nonlocal removed + path = (*parent_titles, node.title) + if node.children: + children: list[TitleNode] = [] + for child in node.children: + kept = _prune(child, path) + if kept is not None: + children.append(kept) + if children: + return replace(node, children=children) + if path in match_overrides: + return replace(node, children=[]) + removed += 1 + return None + if path in match_overrides: + return node + removed += 1 + return None + + out: list[TitleNode] = [] + for node in nodes: + kept = _prune(node, ()) + if kept is not None: + out.append(kept) + + if removed: + logger.info( + "[structure_anchoring] pruned {} unanchored TOC nodes " + "(suffix / no match_overrides → no TOC)", + removed, + ) + return out, removed + + def toc_range_start(hierarchy: dict[str, Any]) -> int | None: toc_range = hierarchy.get("toc_range") if not isinstance(toc_range, (list, tuple)) or not toc_range: @@ -293,66 +345,7 @@ def _visual_rtl_locate_parent( return None, visual_calls -# ── Offset calibration (Agent Phase-1) ─────────────────────────────────────── - - -def calibrate_offset( - *, - nodes: list[TitleNode], - toc_hierarchies: list[dict[str, Any]] | None, - ctx: ToolContext | None, - page_texts: dict[int, str], - page_count: int, -) -> tuple[int | None, dict[tuple[str, ...], TitleMatch]]: - """Discover printed→physical offset (calibration SubAgent Phase-1). - - Production and debug share this entry. Phase-2 bulk/bisect stays in - ``anchor_hierarchy_from_offset``. - """ - from app.services.document_agent.agents.calibration.service import ( - calibrate_offset as _agent_calibrate_offset, - ) - - return _agent_calibrate_offset( - nodes=nodes, - toc_hierarchies=toc_hierarchies, - ctx=ctx, - page_texts=page_texts, - page_count=page_count, - ) - - -def calibrate_offset_via_vlm( - *, - nodes: list[TitleNode], - toc_hierarchies: list[dict[str, Any]] | None, - ctx: ToolContext | None, - page_texts: dict[int, str], - page_count: int, -) -> tuple[int | None, dict[tuple[str, ...], TitleMatch]]: - """Deprecated alias — use ``calibrate_offset`` (Agent Phase-1).""" - return calibrate_offset( - nodes=nodes, - toc_hierarchies=toc_hierarchies, - ctx=ctx, - page_texts=page_texts, - page_count=page_count, - ) - - -def toc_cluster_end_page(toc_hierarchies: list[dict[str, Any]] | None) -> int | None: - """Get the last physical page of the primary TOC cluster.""" - if not toc_hierarchies: - return None - end_pages: list[int] = [] - for hierarchy in toc_hierarchies: - end = toc_range_end(hierarchy) - if end is not None: - end_pages.append(end) - return max(end_pages) if end_pages else None - - -# ── Offset-guided bulk anchoring with recursive recalibrate (Phase A3) ─────── +# ── Offset-guided bulk anchoring with recursive recalibrate (Phase-2) ─────── _TAIL_VERIFY_CONFIDENCE_THRESHOLD = 0.6 _MAX_RECALIBRATE_DEPTH = 5 @@ -576,20 +569,25 @@ def offset_guided_anchoring( for path, node in iter_leaf_title_nodes(nodes) if node.printed_page is not None ] - if len(leaves) < 2: - return None + if not leaves: + return dict(calibration_overrides) or None all_matches: dict[tuple[str, ...], TitleMatch] = {} all_matches.update(calibration_overrides) - _anchor_segment_recursive( - leaves=leaves, - offset=offset, - ctx=ctx, - page_count=page_count, - matches=all_matches, - depth=0, - ) + # Single-leaf regimes (roman front-matter, F-1 appendix, …) still get a + # deterministic printed→physical override; Phase-1 already calibrated them. + if len(leaves) == 1: + all_matches.update(bulk_offset_matches(leaves, offset)) + else: + _anchor_segment_recursive( + leaves=leaves, + offset=offset, + ctx=ctx, + page_count=page_count, + matches=all_matches, + depth=0, + ) if not all_matches: return None @@ -740,7 +738,7 @@ def anchor_hierarchy_from_offset( ) -> tuple[list[TitleNode], SkeletonAnchor]: """Production prune → bulk → null-page given a precomputed offset. - Swap point for agent / VLM calibration: both feed ``offset_hint`` here. + Phase-2 entry after Agent ``calibrate_offset`` (Phase-1). """ seed_overrides = dict(calibration_overrides or {}) pruned_count = 0 @@ -769,6 +767,11 @@ def anchor_hierarchy_from_offset( locate_agent = "offset_only" bulk_count = 0 + working, unanchored_removed = prune_unanchored_toc_leaves( + working, match_overrides=match_overrides + ) + pruned_count += unanchored_removed + match_overrides, null_page_report = locate_null_page_parent_overrides( nodes=working, match_overrides=match_overrides, @@ -802,24 +805,45 @@ def anchor_hierarchy( page_count: int, ctx: ToolContext | None, ) -> tuple[list[TitleNode], SkeletonAnchor]: - """Run offset → prune → bulk → null-page in production order. + """Run Phase-1 calibrate_offset → multi-regime Phase-2 merge. Returns possibly-pruned nodes and the anchor payload. Caller owns resolve_hierarchy_page_ranges / skeleton assembly. """ - offset_hint, calibration_overrides = calibrate_offset( + from app.services.document_agent.agents.calibration.procedure import ( + finalize_calibration_result, + flat_toc_entries, + ) + from app.services.document_agent.agents.calibration.service import ( + calibrate_offset, + ) + + phase1 = calibrate_offset( nodes=nodes, toc_hierarchies=toc_hierarchies, ctx=ctx, page_texts=page_texts, page_count=page_count, ) - return anchor_hierarchy_from_offset( - nodes=nodes, - offset_hint=offset_hint, - calibration_overrides=calibration_overrides, + if phase1.status == "failed" and not phase1.regimes and phase1.offset is None: + return anchor_hierarchy_from_offset( + nodes=nodes, + offset_hint=None, + calibration_overrides={}, + page_texts=page_texts, + body_pages=body_pages, + page_count=page_count, + ctx=ctx, + ) + + working, anchor, _finalized = finalize_calibration_result( + result=phase1, + entries=flat_toc_entries(toc_hierarchies), + toc_hierarchies=list(toc_hierarchies or []), + ctx=ctx, + page_count=page_count, page_texts=page_texts, body_pages=body_pages, - page_count=page_count, - ctx=ctx, + nodes=nodes, ) + return working, anchor diff --git a/apps/worker/app/services/document_agent/tools/inspect_pages.py b/apps/worker/app/services/document_agent/tools/inspect_pages.py index 3946f71ea..7e1b2c7b3 100644 --- a/apps/worker/app/services/document_agent/tools/inspect_pages.py +++ b/apps/worker/app/services/document_agent/tools/inspect_pages.py @@ -17,7 +17,11 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: - """Open one or more physical PDF pages, render them, and answer ``question``.""" + """Open one or more physical PDF pages, render them, and answer ``question``. + + Hard limits are token/loop budgets (via ``BudgetTracker``), not a total page + counter. ``inspect_page_cap`` only caps pages **per call** (batch size). + """ start = time.monotonic() raw_pages = args.get("pages") or [] question = str(args.get("question") or "").strip() @@ -52,29 +56,6 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) - used = int(ctx.blackboard.global_signals.get("inspect_pages_used") or 0) - # Backward-compatible calibration counter. - used = max( - used, - int(ctx.blackboard.global_signals.get("calibration_inspect_pages_used") or 0), - ) - page_budget = int(ctx.settings.get("inspect_page_budget") or 0) - if page_budget > 0 and used >= page_budget: - return ToolResult( - status="error", - error="inspect page budget exhausted", - latency_ms=int((time.monotonic() - start) * 1000), - ) - if page_budget > 0: - remain = max(page_budget - used, 0) - pages = pages[:remain] - if not pages: - return ToolResult( - status="error", - error="inspect page budget exhausted", - latency_ms=int((time.monotonic() - start) * 1000), - ) - from app.services.document_agent.visual import render_pages folder_name = str(args.get("folder_name") or "inspect_pages") @@ -101,6 +82,18 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) + # Token budget (optional stage, e.g. calibration). No total page-count ledger. + stage = args.get("visual_stage") or ctx.settings.get("inspect_visual_stage") + stage_name = str(stage).strip() if stage else None + est = 800 * len(rendered) + 800 + if stage_name and ctx.budget is not None: + if not ctx.budget.try_reserve("visual", est, stage=stage_name): + return ToolResult( + status="error", + error="calibration visual budget exhausted", + latency_ms=int((time.monotonic() - start) * 1000), + ) + prompt = ( "Answer the question about the provided PDF page image(s). " "Return strict JSON object with keys: " @@ -135,7 +128,17 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: usage_task=usage_task, ) payload = json.loads(raw) if raw else {} + tokens_used = int((usage or {}).get("total_tokens") or 0) + if stage_name and ctx.budget is not None: + ctx.budget.commit( + "visual", + actual=tokens_used or est, + est=est, + stage=stage_name, + ) except Exception as exc: + if stage_name and ctx.budget is not None: + ctx.budget.refund("visual", est=est, stage=stage_name) logger.warning("[inspect.pages] VLM failed: {}", exc) return ToolResult( status="error", @@ -143,6 +146,12 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) + # Diagnostic counter only (not a hard budget). + used = int(ctx.blackboard.global_signals.get("inspect_pages_used") or 0) + used = max( + used, + int(ctx.blackboard.global_signals.get("calibration_inspect_pages_used") or 0), + ) next_used = used + len(pages) ctx.blackboard.global_signals["inspect_pages_used"] = next_used ctx.blackboard.global_signals["calibration_inspect_pages_used"] = next_used @@ -157,7 +166,6 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "raw": payload, }, latency_ms=int((time.monotonic() - start) * 1000), - tokens_used=int((usage or {}).get("total_tokens") or 0), + tokens_used=tokens_used, output_summary={"pages": pages, "answer": payload.get("answer")}, ) - 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 7aa3a57cf..c35cf2b3f 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 @@ -644,19 +644,21 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: offset_hint: int | None = None if ctx.blackboard.toc_hierarchies: - from app.services.document_agent.structure.hierarchy_locator import extract_toc_nodes - from app.services.document_agent.structure.structure_anchoring import ( - calibrate_offset, + 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) - offset_hint, _ = calibrate_offset( + 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 # Try TOC chapter-based planning first diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index abce3c08b..eda82b247 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -21,11 +21,9 @@ iter_leaf_title_nodes, resolve_hierarchy_page_ranges, ) +from app.services.document_agent.agents.calibration import calibrate_offset from app.services.document_agent.structure.structure_anchoring import ( anchor_hierarchy, - calibrate_offset, - locate_null_page_parent_overrides, - offset_guided_anchoring, toc_range_end, toc_range_start, ) @@ -279,6 +277,8 @@ def _collapse(node: TitleNode) -> TitleNode: if only_child.children: merged_title = f"{node.title} {only_child.title}" merged_printed_page = only_child.printed_page or node.printed_page + merged_printed_label = only_child.printed_label or node.printed_label + merged_page_kind = only_child.page_kind or node.page_kind merged_physical_hint = ( only_child.physical_page_hint or node.physical_page_hint ) @@ -291,6 +291,8 @@ def _collapse(node: TitleNode) -> TitleNode: node, title=merged_title, printed_page=merged_printed_page, + printed_label=merged_printed_label, + page_kind=merged_page_kind, physical_page_hint=merged_physical_hint, children=promoted, ) @@ -453,13 +455,19 @@ def _resolve_pending_tocs( if p <= toc_scope_end and (toc_scope_start is None or p >= toc_scope_start) ] - offset, cal_overrides = calibrate_offset( + 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( @@ -481,38 +489,25 @@ def _resolve_pending_tocs( ) continue - offset_matches = offset_guided_anchoring( - nodes=nodes, - offset=offset, + 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, - calibration_overrides=cal_overrides, - ) - - if offset_matches is not None: - match_overrides = offset_matches - locate_summary: dict[str, Any] = { - "agent": "offset_guided_bulk", - "offset": offset, - "bulk_count": len(offset_matches), - "toc_relationship": relationship, - } - else: - match_overrides = cal_overrides - locate_summary = { - "agent": "offset_only", - "offset": offset, - "toc_relationship": relationship, - "reason": "offset_guided_anchoring_skipped_or_empty", - } - - match_overrides, null_page_report = locate_null_page_parent_overrides( - nodes=nodes, - match_overrides=match_overrides, page_texts=page_texts, body_pages=toc_body_pages, - ctx=ctx, + nodes=nodes, ) + match_overrides = skeleton_anchor.match_overrides + null_page_report = skeleton_anchor.null_page_report + locate_summary: dict[str, Any] = { + "agent": skeleton_anchor.locate_agent, + "offset": skeleton_anchor.offset, + "bulk_count": skeleton_anchor.bulk_count, + "pruned_out_of_scope": skeleton_anchor.pruned_count, + "toc_relationship": relationship, + } 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), @@ -526,7 +521,7 @@ def _resolve_pending_tocs( } ranges = resolve_hierarchy_page_ranges( - nodes, + resolve_nodes, page_count=toc_scope_end, page_texts=page_texts, body_pages=toc_body_pages, diff --git a/apps/worker/tests/contract/test_structure_anchoring_contract.py b/apps/worker/tests/contract/test_structure_anchoring_contract.py index 84261f4fc..e17a7565d 100644 --- a/apps/worker/tests/contract/test_structure_anchoring_contract.py +++ b/apps/worker/tests/contract/test_structure_anchoring_contract.py @@ -1,4 +1,4 @@ -"""Contract tests for structure_anchoring (moved from skeleton_extractor).""" +"""Contract tests for structure_anchoring (Phase-2) + calibrate wiring.""" from __future__ import annotations @@ -96,7 +96,7 @@ def test_null_page_parent_located_via_compact_text() -> None: assert report[0]["page"] == 5 -def test_calibrate_and_bulk_via_mocked_offset() -> None: +def test_phase2_bulk_via_mocked_offset() -> None: leaves = [ _leaf("Intro", 3), _leaf("Body", 10), @@ -119,33 +119,17 @@ def fake_verify(**kwargs: Any) -> dict[str, Any]: expected = kwargs["candidate_matches"][0].page return {"selected_page": expected, "confidence": 0.9, "reason": "ok"} - with ( - patch.object( - anchoring, - "calibrate_offset", - return_value=(2, seed), - ), - patch.object( - anchoring, - "verify_section_page_choice", - side_effect=fake_verify, - ), + with patch.object( + anchoring, + "verify_section_page_choice", + side_effect=fake_verify, ): - offset, seed_overrides = anchoring.calibrate_offset( - nodes=leaves, - toc_hierarchies=[{"toc_range": [1, 2], "toc_tree": {}}], - ctx=ctx, - page_texts={}, - page_count=30, - ) - assert offset == 2 - assert seed_overrides matches = anchoring.offset_guided_anchoring( nodes=leaves, - offset=offset, + offset=2, ctx=ctx, page_count=30, - calibration_overrides=seed_overrides, + calibration_overrides=seed, ) assert matches is not None assert len(matches) >= 3 @@ -154,28 +138,52 @@ def fake_verify(**kwargs: Any) -> dict[str, Any]: assert matches[("End",)].page == 22 -def test_anchor_hierarchy_returns_skeleton_anchor_fields() -> None: +def test_anchor_hierarchy_uses_calibration_phase1() -> None: leaves = [_leaf("Only", 2)] - toc_hierarchies = [{"toc_range": [1, 1], "toc_tree": {}}] + toc_hierarchies = [ + { + "toc_range": [1, 1], + "toc_with_level": [ + {"heading": "Only", "level": 1, "page_number": 2}, + ], + } + ] ctx = _ctx() - seed = { - ("Only",): TitleMatch( - page=4, - confidence=0.95, - source="agent_vlm", - matched_line="", - score=0.95, - candidates=[4], - evidence={"calibration": True}, - ) - } + + from app.services.document_agent.agents.calibration.types import ( + CalibrationRegime, + CalibrationResult, + CalibrationSample, + ) + + phase1 = CalibrationResult( + status="ok", + offset=2, + offset_status="ok", + regimes=[ + CalibrationRegime( + kind="decimal", + offset=2, + offset_status="ok", + entry_indices=[0], + samples=[ + CalibrationSample( + title="Only", printed_label=2, physical=4 + ) + ], + ) + ], + ) def fake_verify(**kwargs: Any) -> dict[str, Any]: expected = kwargs["candidate_matches"][0].page return {"selected_page": expected, "confidence": 0.95, "reason": "ok"} with ( - patch.object(anchoring, "calibrate_offset", return_value=(2, seed)), + patch( + "app.services.document_agent.agents.calibration.service.calibrate_offset", + return_value=phase1, + ), patch.object( anchoring, "verify_section_page_choice", @@ -198,3 +206,191 @@ def fake_verify(**kwargs: Any) -> dict[str, Any]: assert isinstance(anchor.bulk_count, int) assert isinstance(anchor.pruned_count, int) assert nodes + assert ("Only",) in anchor.match_overrides + assert anchor.match_overrides[("Only",)].page == 4 + + +def test_prune_unanchored_suffix_removes_toc_leaves() -> None: + """Leaves without match_overrides are dropped (suffix = no TOC).""" + nodes = [ + _leaf("A", 1), + _leaf("B", 10), + _leaf("C", 20), + _leaf("D", 30), + ] + overrides = anchoring.bulk_offset_matches( + [(("A",), nodes[0]), (("B",), nodes[1])], + offset=5, + ) + pruned, removed = anchoring.prune_unanchored_toc_leaves( + nodes, match_overrides=overrides + ) + assert removed == 2 + assert [n.title for n in pruned] == ["A", "B"] + assert ("A",) in overrides and ("B",) in overrides + + +def test_phase2_recalibrate_miss_drops_suffix_from_tree() -> None: + """When suffix cannot be recalibrated, those leaves leave the TOC tree.""" + from app.services.document_agent.agents.calibration.procedure import ( + anchor_hierarchy_from_regimes, + ) + from app.services.document_agent.agents.calibration.types import ( + CalibrationRegime, + CalibrationResult, + CalibrationSample, + ) + + leaves = [ + _leaf("Ch1", 1), + _leaf("Ch2", 5), + _leaf("Ch3", 20), + _leaf("Ch4", 30), + ] + phase1 = CalibrationResult( + status="ok", + offset=10, + regimes=[ + CalibrationRegime( + kind="decimal", + offset=10, + offset_status="ok", + entry_indices=[0, 1, 2, 3], + samples=[ + CalibrationSample(title="Ch1", printed_label=1, physical=11) + ], + ) + ], + ) + ctx = _ctx() + + def fake_verify(**kwargs: Any) -> dict[str, Any]: + expected = int(kwargs["candidate_matches"][0].page) + title = str(kwargs.get("title") or "") + # Prefix Ch1/Ch2 at offset=10 confirm; Ch3/Ch4 and recalibrate (+1..+5) miss. + ok_pages = {11, 15} # 1+10, 5+10 + if expected in ok_pages and title in {"Ch1", "Ch2"}: + return {"selected_page": expected, "confidence": 0.95, "reason": "ok"} + # Tail / bisect mid / recalibrate probes for Ch3/Ch4 all fail. + return {"selected_page": None, "confidence": 0.1, "reason": "miss"} + + with patch.object( + anchoring, + "verify_section_page_choice", + side_effect=fake_verify, + ): + working, anchor = anchor_hierarchy_from_regimes( + nodes=leaves, + result=phase1, + entries=[ + {"heading": "Ch1", "level": 1, "page_number": 1}, + {"heading": "Ch2", "level": 1, "page_number": 5}, + {"heading": "Ch3", "level": 1, "page_number": 20}, + {"heading": "Ch4", "level": 1, "page_number": 30}, + ], + page_texts={11: "Ch1", 15: "Ch2", 30: "noise", 40: "noise"}, + body_pages=list(range(1, 50)), + page_count=50, + ctx=ctx, + ) + + titles = [n.title for n in working] + assert titles == ["Ch1", "Ch2"] + assert ("Ch1",) in anchor.match_overrides + assert ("Ch2",) in anchor.match_overrides + assert ("Ch3",) not in anchor.match_overrides + assert ("Ch4",) not in anchor.match_overrides + assert anchor.pruned_count >= 2 + + +def test_multi_regime_phase2_merges_physical_overrides() -> None: + """Roman + decimal + prefixed each apply their own offset → physical pages.""" + from app.services.document_agent.agents.calibration.procedure import ( + anchor_hierarchy_from_regimes, + ) + from app.services.document_agent.agents.calibration.types import ( + CalibrationRegime, + CalibrationResult, + CalibrationSample, + ) + from app.services.document_agent.structure.hierarchy_locator import extract_toc_nodes + + toc = [ + { + "toc_range": [1, 1], + "toc_with_level": [ + {"heading": "Glossary", "level": 1, "page_number": "iv"}, + {"heading": "Summary", "level": 1, "page_number": 1}, + {"heading": "Risks", "level": 1, "page_number": 10}, + {"heading": "Financials", "level": 1, "page_number": "F-1"}, + ], + } + ] + nodes = extract_toc_nodes(toc) + assert nodes[0].page_kind == "roman" + assert nodes[0].printed_page == 4 + assert nodes[1].page_kind == "decimal" + assert nodes[3].page_kind == "prefixed" + assert nodes[3].printed_page == 1 + + phase1 = CalibrationResult( + status="ok", + offset=20, + regimes=[ + CalibrationRegime( + kind="roman", + offset=16, + offset_status="ok", + entry_indices=[0], + samples=[ + CalibrationSample(title="Glossary", printed_label="iv", physical=20) + ], + ), + CalibrationRegime( + kind="decimal", + offset=20, + offset_status="ok", + entry_indices=[1, 2], + samples=[ + CalibrationSample(title="Summary", printed_label=1, physical=21) + ], + ), + CalibrationRegime( + kind="prefixed", + offset=300, + offset_status="ok", + entry_indices=[3], + samples=[ + CalibrationSample( + title="Financials", printed_label="F-1", physical=301 + ) + ], + ), + ], + ) + ctx = _ctx() + + def fake_verify(**kwargs: Any) -> dict[str, Any]: + expected = kwargs["candidate_matches"][0].page + return {"selected_page": expected, "confidence": 0.95, "reason": "ok"} + + with patch.object( + anchoring, + "verify_section_page_choice", + side_effect=fake_verify, + ): + _working, anchor = anchor_hierarchy_from_regimes( + nodes=nodes, + result=phase1, + entries=toc[0]["toc_with_level"], + page_texts={20: "Glossary", 21: "Summary", 30: "Risks", 301: "Financials"}, + body_pages=list(range(1, 320)), + page_count=320, + ctx=ctx, + ) + + assert anchor.match_overrides[("Glossary",)].page == 20 + assert anchor.match_overrides[("Summary",)].page == 21 + assert anchor.match_overrides[("Risks",)].page == 30 + assert anchor.match_overrides[("Financials",)].page == 301 + assert anchor.offset == 20 # primary decimal \ No newline at end of file From 5ff2246fd0638cd8f5fc525c1dcdab6ee4111406 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 12 Aug 2026 17:49:44 +0800 Subject: [PATCH 05/12] fix: drop unused AsyncSession import in retrieval contract test Co-authored-by: Cursor --- apps/api/tests/contract/test_retrieval_contract.py | 1 - 1 file changed, 1 deletion(-) diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 032e1a8a6..78dffed98 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -6,7 +6,6 @@ import pytest from httpx import AsyncClient from pytest import MonkeyPatch -from sqlalchemy.ext.asyncio import AsyncSession from tests.support.contract_database import ContractDatabase From 22bb71d8c37a089a9084f21ee84bdad4c92b1870 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 12 Aug 2026 17:52:42 +0800 Subject: [PATCH 06/12] refactor: require calibrated TOC offset and share hierarchy collapse Drop arithmetic offset fallback and legacy page_locate budget so shard planning only uses VLM-calibrated offsets, and reuse one collapse helper across calibration and page_memory. Co-authored-by: Cursor --- apps/worker/.env.example | 2 - .../agents/calibration/procedure.py | 6 +- .../app/services/document_agent/budget.py | 5 -- .../services/document_agent/coordinator.py | 6 -- .../structure/hierarchy_locator.py | 42 ++++++++++ .../tools/propose_shard_plan.py | 76 ++++++------------- .../page_memory/skeleton_extractor.py | 40 +--------- 7 files changed, 74 insertions(+), 103 deletions(-) diff --git a/apps/worker/.env.example b/apps/worker/.env.example index 7c140aa80..5e4d118e7 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -141,8 +141,6 @@ PARSE_AGENT_CALIBRATION_MIN_BUDGET=12000 PARSE_AGENT_CALIBRATION_CAP=40000 PARSE_AGENT_PAGE_TAGGING_MIN_BUDGET=0 PARSE_AGENT_PAGE_TAGGING_CAP=0 -PARSE_AGENT_PAGE_LOCATE_MIN_BUDGET=0 -PARSE_AGENT_PAGE_LOCATE_CAP=0 # Parser row schema. `entities` (JSON typed entities, §4.4) and `asset_title` # (asset caption/label, §4.5) are additive trailing columns; the parser layer diff --git a/apps/worker/app/services/document_agent/agents/calibration/procedure.py b/apps/worker/app/services/document_agent/agents/calibration/procedure.py index a27126f63..40642bf78 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/procedure.py +++ b/apps/worker/app/services/document_agent/agents/calibration/procedure.py @@ -495,11 +495,11 @@ def finalize_calibration_result( texts = dict(page_texts or {}) bodies = list(body_pages or sorted(texts.keys()) or list(range(1, page_count + 1))) if nodes is None: - from app.services.page_memory.skeleton_extractor import ( - _collapse_intermediate_single_child_chains, + from app.services.document_agent.structure.hierarchy_locator import ( + collapse_intermediate_single_child_chains, ) - nodes = _collapse_intermediate_single_child_chains( + nodes = collapse_intermediate_single_child_chains( extract_toc_nodes(toc_hierarchies) ) diff --git a/apps/worker/app/services/document_agent/budget.py b/apps/worker/app/services/document_agent/budget.py index 86058e27b..3ecdf867f 100644 --- a/apps/worker/app/services/document_agent/budget.py +++ b/apps/worker/app/services/document_agent/budget.py @@ -12,7 +12,6 @@ "coarse_planner", "structural_react", "calibration", - "page_locate", "page_tagging", ] @@ -174,10 +173,6 @@ def fork(self, ratio: float) -> "BudgetTracker": for stage, envelope in self._visual_stage_envelopes.items(): usage = self._visual_stage_usage.get(stage, StageUsage()) remaining_guarantee = max(envelope.min_guarantee - usage.committed, 0) - if stage == "page_locate": - # Legacy uncapped locate stage: keep child uncapped. - child_envelopes[stage] = StageEnvelope(min_guarantee=0, cap=0) - continue child_envelopes[stage] = StageEnvelope( min_guarantee=int(remaining_guarantee * ratio), cap=int(envelope.cap * ratio) if envelope.cap is not None else None, diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 983089d8e..00caba70a 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -72,12 +72,6 @@ def __init__( ), cap=int(os.environ.get("PARSE_AGENT_CALIBRATION_CAP", "40000")), ), - "page_locate": StageEnvelope( - min_guarantee=int( - os.environ.get("PARSE_AGENT_PAGE_LOCATE_MIN_BUDGET", "0") - ), - cap=int(os.environ.get("PARSE_AGENT_PAGE_LOCATE_CAP", "0")) or None, - ), "page_tagging": StageEnvelope( min_guarantee=int( os.environ.get("PARSE_AGENT_PAGE_TAGGING_MIN_BUDGET", "0") diff --git a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py index 46cace860..86719c1d1 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -766,3 +766,45 @@ def _safe_int(value: Any) -> int | None: return int(value) except (TypeError, ValueError): return None + + +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. + Shared by page_memory C4 and calibration finalize. + """ + from dataclasses import replace as _replace + + def _collapse(node: TitleNode) -> TitleNode: + collapsed_children = [_collapse(c) for c in node.children] + + if len(collapsed_children) == 1: + only_child = collapsed_children[0] + if only_child.children: + merged_title = f"{node.title} {only_child.title}" + merged_printed_page = only_child.printed_page or node.printed_page + merged_printed_label = only_child.printed_label or node.printed_label + merged_page_kind = only_child.page_kind or node.page_kind + merged_physical_hint = ( + only_child.physical_page_hint or node.physical_page_hint + ) + promoted = [ + _replace(gc, level=max(1, gc.level - 1)) + for gc in only_child.children + ] + return _replace( + node, + title=merged_title, + printed_page=merged_printed_page, + printed_label=merged_printed_label, + page_kind=merged_page_kind, + physical_page_hint=merged_physical_hint, + children=promoted, + ) + + return _replace(node, children=collapsed_children) + + return [_collapse(n) for n in nodes] 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 c35cf2b3f..4f5b91c5b 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 @@ -15,7 +15,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 loguru import logger from shared.utils.token_estimate import estimate_tokens @@ -27,11 +26,10 @@ def derive_leaf_cut_pages( """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). The offset from printed page to physical page is - either provided via offset_override (VLM-calibrated) or computed arithmetically - from toc_range and the first entry's page_number as a fallback. + (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: + if not toc_hierarchies or offset_override is None: return [] all_pages: list[int] = [] @@ -47,28 +45,10 @@ def derive_leaf_cut_pages( if not entries: continue - if offset_override is not None: - offset = offset_override - else: - toc_end_page = toc_range[1] if isinstance(toc_range, list) else toc_range - first_printed = next( - (e.get("page_number") for e in entries if e.get("page_number") is not None), - None, - ) - if first_printed is None: - continue - offset = (toc_end_page + 1) - first_printed - logger.warning( - "[propose_shard_plan] using arithmetic offset fallback: " - "toc_end={} first_printed={} offset={}", - toc_end_page, - first_printed, - offset, - ) - + offset = offset_override for i, entry in enumerate(entries): pn = entry.get("page_number") - if pn is None: + if not isinstance(pn, int): continue is_leaf = ( i == len(entries) - 1 @@ -95,8 +75,10 @@ def derive_chapter_boundaries( 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: + if not toc_hierarchies or offset_override is None: return [] all_entries: list[dict[str, Any]] = [] @@ -112,23 +94,14 @@ def derive_chapter_boundaries( if not entries: continue - if offset_override is not None: - offset = offset_override - else: - toc_end_page = toc_range[1] if isinstance(toc_range, list) else toc_range - first_printed = next( - (e.get("page_number") for e in entries if e.get("page_number") is not None), - None, - ) - if first_printed is None: - continue - offset = (toc_end_page + 1) - first_printed + offset = offset_override - # Collect all entries with physical pages + # 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 pn is None: + if not isinstance(pn, int): continue physical = pn + offset if physical < 1 or physical > page_count: @@ -202,9 +175,20 @@ def split_toc_for_shard( For continuation shards (not starting at page 1), the ancestor chain of the first entry is prepended so downstream heading prediction has the full structural context. + + Requires calibrated ``offset_override`` for page-unit TOC regions. """ if not toc_hierarchies: return None + if offset_override is None: + # Without a calibrated offset, keep non-page TOC payloads as-is and + # skip page-unit hierarchies rather than inventing arithmetic offsets. + kept = [ + hier + for hier in toc_hierarchies + if hier.get("toc_range_unit") != "page" + ] + return kept or None result: list[dict[str, Any]] = [] for hier in toc_hierarchies: @@ -220,23 +204,13 @@ def split_toc_for_shard( if not entries: continue - if offset_override is not None: - offset = offset_override - else: - toc_end_page = toc_range[1] if isinstance(toc_range, list) else toc_range - first_printed = next( - (e.get("page_number") for e in entries if e.get("page_number") is not None), - None, - ) - if first_printed is None: - continue - offset = (toc_end_page + 1) - first_printed + offset = offset_override shard_entries: list[dict[str, Any]] = [] first_idx: int | None = None for idx, entry in enumerate(entries): pn = entry.get("page_number") - if pn is None: + if not isinstance(pn, int): continue physical = pn + offset if shard_page_start <= physical <= shard_page_end: diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index eda82b247..6a35e63df 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -264,43 +264,11 @@ def _collapse_intermediate_single_child_chains( Leaf nodes (children=[]) are never absorbed into their parent title. """ - from dataclasses import replace as _replace - - def _collapse(node: TitleNode) -> TitleNode: - # Recurse first (bottom-up), so grand-children are already collapsed. - collapsed_children = [_collapse(c) for c in node.children] - - if len(collapsed_children) == 1: - only_child = collapsed_children[0] - # Only fold when the child is itself an intermediate node - # (i.e. still has children). Leaf nodes are left intact. - if only_child.children: - merged_title = f"{node.title} {only_child.title}" - merged_printed_page = only_child.printed_page or node.printed_page - merged_printed_label = only_child.printed_label or node.printed_label - merged_page_kind = only_child.page_kind or node.page_kind - merged_physical_hint = ( - only_child.physical_page_hint or node.physical_page_hint - ) - # Promote grandchildren one level up (close the level gap). - promoted = [ - _replace(gc, level=max(1, gc.level - 1)) - for gc in only_child.children - ] - return _replace( - node, - title=merged_title, - printed_page=merged_printed_page, - printed_label=merged_printed_label, - page_kind=merged_page_kind, - physical_page_hint=merged_physical_hint, - children=promoted, - ) - - return _replace(node, children=collapsed_children) - - return [_collapse(n) for n in nodes] + from app.services.document_agent.structure.hierarchy_locator import ( + collapse_intermediate_single_child_chains, + ) + return collapse_intermediate_single_child_chains(nodes) def _root_skeleton( *, From c4baf6ccab27efb227981497e7ed33b86a405429 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 12 Aug 2026 18:15:02 +0800 Subject: [PATCH 07/12] fix: keep native Excel datetime cells until HTML render Avoid early Timestamp/NaT stringification in postprocess_tb so sparse datetime columns still render clean HTML without NaT leakage. Co-authored-by: Cursor --- .../tables/table_frame_parser.py | 16 +----- .../contract/test_excel_parser_contract.py | 55 +++++++++++++++++++ 2 files changed, 57 insertions(+), 14 deletions(-) diff --git a/apps/worker/app/services/document_parser/tables/table_frame_parser.py b/apps/worker/app/services/document_parser/tables/table_frame_parser.py index 38aa92b68..cc3b6f683 100644 --- a/apps/worker/app/services/document_parser/tables/table_frame_parser.py +++ b/apps/worker/app/services/document_parser/tables/table_frame_parser.py @@ -1,7 +1,6 @@ # pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalOperand=false, reportOptionalSubscript=false, reportReturnType=false, reportOperatorIssue=false, reportIndexIssue=false, reportAssignmentType=false, reportGeneralTypeIssues=false from __future__ import annotations -import datetime import os import uuid from collections import OrderedDict @@ -372,21 +371,10 @@ def make_padded(name: object) -> object: for column in table_frame.columns ] - table_frame = table_frame.map( + # Keep native temporals (incl. NaT); stringify at df2html / to_html only. + return table_frame.map( lambda value: value.replace("\n", "") if isinstance(value, str) else value ) - return process_datetime_cells(table_frame) - - -def process_datetime_cells(table_frame: pd.DataFrame) -> pd.DataFrame: - table_frame = table_frame.copy() - - def convert(value: object) -> object: - if isinstance(value, (pd.Timestamp, datetime.datetime)): - return value.strftime("%Y-%m-%d %H:%M:%S") - return value - - return table_frame.apply(lambda column: column.map(convert)) def process_duplicate_cols(columns: object) -> list[object]: diff --git a/apps/worker/tests/contract/test_excel_parser_contract.py b/apps/worker/tests/contract/test_excel_parser_contract.py index cdaf54d9d..2b17d5f90 100644 --- a/apps/worker/tests/contract/test_excel_parser_contract.py +++ b/apps/worker/tests/contract/test_excel_parser_contract.py @@ -99,3 +99,58 @@ def test_parser_maps_document_name_to_task_local_path_segment( assert full_output_dir.endswith("images.xlsx") assert parsed_df is not None assert parsed_df["path"].tolist() == ["images.xlsx/Visible"] + + +def _write_workbook_with_sparse_datetimes(workbook_path: Path) -> None: + """Datetime column with blanks → pandas NaT in-frame; HTML must still render.""" + import openpyxl + from datetime import datetime + + workbook = openpyxl.Workbook() + sheet = workbook.active + sheet.title = "Roster" + sheet["A1"] = "Name" + sheet["B1"] = "SignupDate" + sheet["A2"] = "Alice" + sheet["B2"] = datetime(2026, 6, 27) + sheet["A3"] = "Bob" + sheet["B3"] = None + sheet["A4"] = "Carol" + sheet["B4"] = datetime(2026, 7, 4) + workbook.save(workbook_path) + + +def test_xlsx_parser_tolerates_nat_in_datetime_column( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.parse_service import checkerboard_parse_output + + workbook_path = tmp_path / "roster.xlsx" + output_root = tmp_path / "parser-output" + _write_workbook_with_sparse_datetimes(workbook_path) + + parse_output = checkerboard_parse_output( + file_full_path=str(workbook_path), + filename="roster.xlsx", + output_dir=str(output_root), + internal_output_filename="roster.xlsx", + summary_image=False, + summary_table=False, + summary_txt=False, + smart_title_parse=False, + stopwords=[], + ) + + parsed_df = parse_output.parsed_df + assert parsed_df is not None + assert len(parsed_df) == 1 + + table_html = ( + Path(parse_output.output_dir) / "tables" / "table-Roster.html" + ).read_text(encoding="utf-8") + assert "Alice" in table_html + assert "2026-06-27" in table_html + assert "Carol" in table_html + assert "2026-07-04" in table_html + assert "NaT" not in table_html From 2a6122db7b4e9c0f431ad16f3c60b6b19609cbe4 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 12 Aug 2026 18:27:12 +0800 Subject: [PATCH 08/12] test: stabilize structure anchoring verify mocks under contract suite Patch verify on the live anchoring globals and drop brittle SkeletonAnchor isinstance checks so contract autouse rebinds cannot miss the mock. Co-authored-by: Cursor --- .../test_structure_anchoring_contract.py | 55 +++++++++++-------- 1 file changed, 33 insertions(+), 22 deletions(-) diff --git a/apps/worker/tests/contract/test_structure_anchoring_contract.py b/apps/worker/tests/contract/test_structure_anchoring_contract.py index e17a7565d..d81465f65 100644 --- a/apps/worker/tests/contract/test_structure_anchoring_contract.py +++ b/apps/worker/tests/contract/test_structure_anchoring_contract.py @@ -3,6 +3,8 @@ from __future__ import annotations import os +from collections.abc import Callable, Iterator +from contextlib import contextmanager from typing import Any from unittest.mock import patch @@ -20,6 +22,31 @@ from app.services.document_agent.structure import structure_anchoring as anchoring +@contextmanager +def _patch_verify(fake_verify: Callable[..., dict[str, Any]]) -> Iterator[None]: + """Patch verify on the module dict closed over by live anchoring code.""" + from app.services.document_agent.agents.calibration import procedure + + dicts = [procedure.offset_guided_anchoring.__globals__, anchoring.__dict__] + seen: set[int] = set() + originals: list[tuple[dict[str, Any], Any]] = [] + for module_dict in dicts: + dict_id = id(module_dict) + if dict_id in seen: + continue + seen.add(dict_id) + originals.append((module_dict, module_dict.get("verify_section_page_choice"))) + module_dict["verify_section_page_choice"] = fake_verify + try: + yield + finally: + for module_dict, original in originals: + if original is None: + module_dict.pop("verify_section_page_choice", None) + else: + module_dict["verify_section_page_choice"] = original + + def _ctx() -> ToolContext: return ToolContext( pdf_path="/tmp/doc.pdf", @@ -27,7 +54,7 @@ def _ctx() -> ToolContext: blackboard=AgentBlackboard(), budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), trace=None, - settings={}, + settings={"vlm_model": "test-vlm"}, ) @@ -119,11 +146,7 @@ def fake_verify(**kwargs: Any) -> dict[str, Any]: expected = kwargs["candidate_matches"][0].page return {"selected_page": expected, "confidence": 0.9, "reason": "ok"} - with patch.object( - anchoring, - "verify_section_page_choice", - side_effect=fake_verify, - ): + with _patch_verify(fake_verify): matches = anchoring.offset_guided_anchoring( nodes=leaves, offset=2, @@ -184,11 +207,7 @@ def fake_verify(**kwargs: Any) -> dict[str, Any]: "app.services.document_agent.agents.calibration.service.calibrate_offset", return_value=phase1, ), - patch.object( - anchoring, - "verify_section_page_choice", - side_effect=fake_verify, - ), + _patch_verify(fake_verify), ): nodes, anchor = anchoring.anchor_hierarchy( nodes=leaves, @@ -198,7 +217,7 @@ def fake_verify(**kwargs: Any) -> dict[str, Any]: page_count=10, ctx=ctx, ) - assert isinstance(anchor, anchoring.SkeletonAnchor) + # Duck-type: contract conftest can leave a stale SkeletonAnchor class identity. assert anchor.offset == 2 assert anchor.offset_status == "ok" assert isinstance(anchor.match_overrides, dict) @@ -274,11 +293,7 @@ def fake_verify(**kwargs: Any) -> dict[str, Any]: # Tail / bisect mid / recalibrate probes for Ch3/Ch4 all fail. return {"selected_page": None, "confidence": 0.1, "reason": "miss"} - with patch.object( - anchoring, - "verify_section_page_choice", - side_effect=fake_verify, - ): + with _patch_verify(fake_verify): working, anchor = anchor_hierarchy_from_regimes( nodes=leaves, result=phase1, @@ -374,11 +389,7 @@ def fake_verify(**kwargs: Any) -> dict[str, Any]: expected = kwargs["candidate_matches"][0].page return {"selected_page": expected, "confidence": 0.95, "reason": "ok"} - with patch.object( - anchoring, - "verify_section_page_choice", - side_effect=fake_verify, - ): + with _patch_verify(fake_verify): _working, anchor = anchor_hierarchy_from_regimes( nodes=nodes, result=phase1, From 97878fd5d25599d2b28cd2a056ecec0079ef9fd3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 13 Aug 2026 01:55:40 +0800 Subject: [PATCH 09/12] fix: support ECS task role for S3 storage --- .../shared/core/config/storage.py | 40 ++++++-- .../tests/test_storage_config_contract.py | 96 +++++++++++++++++++ 2 files changed, 128 insertions(+), 8 deletions(-) create mode 100644 packages/shared-python/shared/tests/test_storage_config_contract.py diff --git a/packages/shared-python/shared/core/config/storage.py b/packages/shared-python/shared/core/config/storage.py index f8a31c393..0fc8571eb 100644 --- a/packages/shared-python/shared/core/config/storage.py +++ b/packages/shared-python/shared/core/config/storage.py @@ -6,7 +6,7 @@ import boto3 from botocore.client import BaseClient from botocore.config import Config -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator from shared.core.exceptions.domain_exceptions import ( DependencyMissingException, @@ -30,8 +30,8 @@ class StorageConfig(BaseModel): # Shared S3-style configuration used by S3, OSS, and MinIO. S3_BUCKET_NAME: str = Field(..., description="Bucket name") - S3_ACCESS_KEY_ID: str = Field(..., description="Access key ID") - S3_SECRET_ACCESS_KEY: str = Field(..., description="Secret access key") + S3_ACCESS_KEY_ID: str = Field(default="", description="Access key ID") + S3_SECRET_ACCESS_KEY: str = Field(default="", description="Secret access key") S3_ENDPOINT_URL: str = Field( default="", description="Endpoint URL for S3-compatible services such as MinIO" ) @@ -113,6 +113,27 @@ class StorageConfig(BaseModel): default=True, description="Verify OSS event signatures" ) + @model_validator(mode="after") + def validate_storage_credentials(self) -> "StorageConfig": + """Validate credentials according to the selected storage backend.""" + storage_type = self.S3_TYPE.lower() + has_access_key = bool(self.S3_ACCESS_KEY_ID) + has_secret_key = bool(self.S3_SECRET_ACCESS_KEY) + + if has_access_key != has_secret_key: + raise ValueError( + "S3_ACCESS_KEY_ID and S3_SECRET_ACCESS_KEY must be configured together" + ) + + if storage_type in {"oss", "minio"} and not ( + has_access_key and has_secret_key + ): + raise ValueError( + f"Explicit storage credentials are required when S3_TYPE={storage_type}" + ) + + return self + def get_s3_client(self) -> BaseClient: """Return an S3 client for S3-compatible backends.""" # Build the client config. @@ -128,11 +149,14 @@ def get_s3_client(self) -> BaseClient: config = Config(**config_kwargs) if config_kwargs else None # Build client kwargs. - client_kwargs: dict[str, object] = { - "service_name": "s3", - "aws_access_key_id": self.S3_ACCESS_KEY_ID, - "aws_secret_access_key": self.S3_SECRET_ACCESS_KEY, - } + client_kwargs: dict[str, object] = {"service_name": "s3"} + + # When explicit keys are omitted, boto3 automatically retrieves temporary + # authenticated credentials from the ECS task role. Explicit credentials + # remain supported for local and legacy deployments. + if self.S3_ACCESS_KEY_ID and self.S3_SECRET_ACCESS_KEY: + client_kwargs["aws_access_key_id"] = self.S3_ACCESS_KEY_ID + client_kwargs["aws_secret_access_key"] = self.S3_SECRET_ACCESS_KEY # Add endpoint_url for MinIO or custom S3-compatible services. if self.S3_ENDPOINT_URL: diff --git a/packages/shared-python/shared/tests/test_storage_config_contract.py b/packages/shared-python/shared/tests/test_storage_config_contract.py new file mode 100644 index 000000000..d7dce29ef --- /dev/null +++ b/packages/shared-python/shared/tests/test_storage_config_contract.py @@ -0,0 +1,96 @@ +"""Contracts for storage credentials at the boto3 boundary.""" + +import os +from unittest.mock import Mock + +import pytest +from pydantic import ValidationError + +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-bucket") +os.environ.setdefault("S3_ACCESS_KEY_ID", "") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "") +os.environ.setdefault("S3_TEMP_PATH", "/tmp/knowhere-storage-contract") + +from shared.core.config.storage import StorageConfig + + +def create_storage_config(**overrides: str) -> StorageConfig: + """Create the smallest valid storage configuration for a contract test.""" + values: dict[str, str] = { + "S3_BUCKET_NAME": "test-bucket", + "S3_TEMP_PATH": "/tmp/knowhere-storage-contract", + } + values.update(overrides) + return StorageConfig(**values) + + +def test_aws_s3_uses_default_credential_chain_when_keys_are_empty( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """AWS S3 must allow ECS task-role credentials through boto3's chain.""" + boto3_client: Mock = Mock() + monkeypatch.setattr("shared.core.config.storage.boto3.client", boto3_client) + + config: StorageConfig = create_storage_config(S3_TYPE="s3", S3_REGION="us-east-1") + + config.get_s3_client() + + boto3_client.assert_called_once() + client_arguments: dict[str, object] = dict(boto3_client.call_args.kwargs) + assert client_arguments["service_name"] == "s3" + assert client_arguments["region_name"] == "us-east-1" + assert "aws_access_key_id" not in client_arguments + assert "aws_secret_access_key" not in client_arguments + + +def test_aws_s3_passes_complete_explicit_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Existing deployments may continue to provide an explicit key pair.""" + boto3_client: Mock = Mock() + monkeypatch.setattr("shared.core.config.storage.boto3.client", boto3_client) + config: StorageConfig = create_storage_config( + S3_TYPE="s3", + S3_ACCESS_KEY_ID="access-key", + S3_SECRET_ACCESS_KEY="secret-key", + ) + + config.get_s3_client() + + client_arguments: dict[str, object] = dict(boto3_client.call_args.kwargs) + assert client_arguments["aws_access_key_id"] == "access-key" + assert client_arguments["aws_secret_access_key"] == "secret-key" + + +@pytest.mark.parametrize("storage_type", ["oss", "minio"]) +def test_s3_compatible_backends_require_explicit_credentials( + storage_type: str, +) -> None: + """OSS and MinIO must not silently fall back to an AWS identity chain.""" + with pytest.raises( + ValidationError, + match=f"Explicit storage credentials are required when S3_TYPE={storage_type}", + ): + create_storage_config(S3_TYPE=storage_type) + + +@pytest.mark.parametrize( + ("access_key_id", "secret_access_key"), + [("access-key", ""), ("", "secret-key")], +) +def test_storage_rejects_partial_explicit_credentials( + access_key_id: str, + secret_access_key: str, +) -> None: + """A partial key pair must fail before an unusable client is created.""" + with pytest.raises( + ValidationError, + match="S3_ACCESS_KEY_ID and S3_SECRET_ACCESS_KEY must be configured together", + ): + create_storage_config( + S3_TYPE="s3", + S3_ACCESS_KEY_ID=access_key_id, + S3_SECRET_ACCESS_KEY=secret_access_key, + ) From a119ed14b1d1acd73d442ad82ef93e3d974c604f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 13 Aug 2026 02:35:15 +0800 Subject: [PATCH 10/12] fix: break calibration import cycles --- .../document_agent/agents/calibration/loop.py | 2 +- .../agents/calibration/orchestrator.py | 57 ++ .../agents/calibration/procedure.py | 32 +- .../structure/anchoring_primitives.py | 795 +++++++++++++++++ .../structure/structure_anchoring.py | 827 +----------------- .../page_memory/skeleton_extractor.py | 4 +- 6 files changed, 907 insertions(+), 810 deletions(-) create mode 100644 apps/worker/app/services/document_agent/agents/calibration/orchestrator.py create mode 100644 apps/worker/app/services/document_agent/structure/anchoring_primitives.py diff --git a/apps/worker/app/services/document_agent/agents/calibration/loop.py b/apps/worker/app/services/document_agent/agents/calibration/loop.py index 852c9991f..0ebb26039 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/loop.py +++ b/apps/worker/app/services/document_agent/agents/calibration/loop.py @@ -25,7 +25,7 @@ from app.services.document_agent.budget import BudgetTracker, StageEnvelope from app.services.document_agent.manifest import ToolContext, ToolResult from app.services.document_agent.state import AgentBlackboard -from app.services.document_agent.structure.structure_anchoring import ( +from app.services.document_agent.structure.anchoring_primitives import ( deserialize_skeleton_anchor, serialize_skeleton_anchor, ) diff --git a/apps/worker/app/services/document_agent/agents/calibration/orchestrator.py b/apps/worker/app/services/document_agent/agents/calibration/orchestrator.py new file mode 100644 index 000000000..88736bf5d --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/orchestrator.py @@ -0,0 +1,57 @@ +"""Calibration orchestration across Phase-1 and structure Phase-2.""" + +from __future__ import annotations + +from typing import Any + +from app.services.document_agent.agents.calibration.procedure import ( + finalize_calibration_result, + flat_toc_entries, +) +from app.services.document_agent.agents.calibration import service +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.structure.hierarchy_locator import TitleNode +from app.services.document_agent.structure.anchoring_primitives import ( + SkeletonAnchor, + anchor_hierarchy_from_offset, +) + + +def anchor_hierarchy( + *, + nodes: list[TitleNode], + toc_hierarchies: list[dict[str, Any]] | None, + page_texts: dict[int, str], + body_pages: list[int], + page_count: int, + ctx: ToolContext | None, +) -> tuple[list[TitleNode], SkeletonAnchor]: + """Run calibration Phase-1 and the production Phase-2 completion.""" + phase1 = service.calibrate_offset( + nodes=nodes, + toc_hierarchies=toc_hierarchies, + ctx=ctx, + page_texts=page_texts, + page_count=page_count, + ) + if phase1.status == "failed" and not phase1.regimes and phase1.offset is None: + return anchor_hierarchy_from_offset( + nodes=nodes, + offset_hint=None, + calibration_overrides={}, + page_texts=page_texts, + body_pages=body_pages, + page_count=page_count, + ctx=ctx, + ) + working, anchor, _finalized = finalize_calibration_result( + result=phase1, + entries=flat_toc_entries(toc_hierarchies), + toc_hierarchies=list(toc_hierarchies or []), + ctx=ctx, + page_count=page_count, + page_texts=page_texts, + body_pages=body_pages, + nodes=nodes, + ) + return working, anchor diff --git a/apps/worker/app/services/document_agent/agents/calibration/procedure.py b/apps/worker/app/services/document_agent/agents/calibration/procedure.py index 40642bf78..a3ed401d4 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/procedure.py +++ b/apps/worker/app/services/document_agent/agents/calibration/procedure.py @@ -31,18 +31,44 @@ normalize_page_kind, parse_printed_page, ) -from app.services.document_agent.structure.structure_anchoring import ( +from app.services.document_agent.structure.anchoring_primitives import ( SkeletonAnchor, locate_null_page_parent_overrides, - offset_guided_anchoring, prune_unanchored_toc_leaves, serialize_skeleton_anchor, ) +from app.services.document_agent.structure import anchoring_primitives as _anchoring +from app.services.document_agent.structure.page_locate_agent import ( + verify_section_page_choice, +) # Re-export under prior names so existing imports keep working. normalize_kind = normalize_page_kind +def offset_guided_anchoring( + *, + nodes: list[TitleNode], + offset: int, + ctx: ToolContext, + page_count: int, + calibration_overrides: dict[tuple[str, ...], TitleMatch], +) -> dict[tuple[str, ...], TitleMatch] | None: + """Forward phase-2 anchoring while preserving the historical patch seam.""" + original = _anchoring.verify_section_page_choice + _anchoring.verify_section_page_choice = verify_section_page_choice + try: + return _anchoring.offset_guided_anchoring( + nodes=nodes, + offset=offset, + ctx=ctx, + page_count=page_count, + calibration_overrides=calibration_overrides, + ) + finally: + _anchoring.verify_section_page_choice = original + + def pick_primary_offset(result: CalibrationResult) -> int | None: """Prefer decimal-regime candidate offset; else first regime with an offset.""" for regime in result.regimes: @@ -298,7 +324,7 @@ def anchor_hierarchy_from_regimes( if ctx is None: # Offline: still apply deterministic printed+offset for this regime. - from app.services.document_agent.structure.structure_anchoring import ( + from app.services.document_agent.structure.anchoring_primitives import ( bulk_offset_matches, ) diff --git a/apps/worker/app/services/document_agent/structure/anchoring_primitives.py b/apps/worker/app/services/document_agent/structure/anchoring_primitives.py new file mode 100644 index 000000000..390e8ba33 --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/anchoring_primitives.py @@ -0,0 +1,795 @@ +"""Shared hierarchy anchoring: Phase-2 bulk/bisect/null-page + SkeletonAnchor. +Phase-1 offset discovery lives in ``document_agent.agents.calibration``. +``anchor_hierarchy`` composes Phase-1 + Phase-2 for production callers. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any + +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.structure.hierarchy_locator import ( + TitleMatch, + TitleNode, + first_leaf_start_under, + iter_leaf_title_nodes, + last_leaf_start_under, + locate_title_compact_strict, +) +from app.services.document_agent.structure.page_locate_agent import ( + verify_section_page_choice, +) +from loguru import logger + + +def prune_out_of_scope_nodes( + nodes: list[TitleNode], + *, + offset: int, + page_count: int, +) -> tuple[list[TitleNode], int]: + """Remove leaf nodes whose printed_page + offset exceeds page_count. + + Bottom-up: prune out-of-scope leaves, then remove intermediate nodes + that become childless after pruning. Returns (pruned_tree, removed_count). + """ + removed = 0 + + def _prune(node: TitleNode) -> TitleNode | None: + nonlocal removed + if not node.children: + if node.printed_page is not None: + expected = node.printed_page + offset + if expected > page_count or expected < 1: + removed += 1 + return None + return node + pruned_children = [] + for child in node.children: + result = _prune(child) + if result is not None: + pruned_children.append(result) + if not pruned_children: + removed += 1 + return None + return replace(node, children=pruned_children) + + pruned = [] + for node in nodes: + result = _prune(node) + if result is not None: + pruned.append(result) + + if removed: + logger.info( + "[structure_anchoring] pruned {} out-of-scope TOC nodes " + "(printed_page + offset={} exceeds page_count={})", + removed, + offset, + page_count, + ) + + return pruned, removed + + +def prune_unanchored_toc_leaves( + nodes: list[TitleNode], + *, + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> tuple[list[TitleNode], int]: + """Remove TOC leaves that have no physical ``match_overrides`` entry. + + Implements Phase-2 ``suffix = no TOC``: after bulk/bisect/recalibrate, any + leaf that was not successfully anchored is dropped from the coarse tree + instead of sticky ``inherited_unlocated`` ranges. Childless parents are + removed unless they themselves have an override. + """ + removed = 0 + + def _prune( + node: TitleNode, parent_titles: tuple[str, ...] + ) -> TitleNode | None: + nonlocal removed + path = (*parent_titles, node.title) + if node.children: + children: list[TitleNode] = [] + for child in node.children: + kept = _prune(child, path) + if kept is not None: + children.append(kept) + if children: + return replace(node, children=children) + if path in match_overrides: + return replace(node, children=[]) + removed += 1 + return None + if path in match_overrides: + return node + removed += 1 + return None + + out: list[TitleNode] = [] + for node in nodes: + kept = _prune(node, ()) + if kept is not None: + out.append(kept) + + if removed: + logger.info( + "[structure_anchoring] pruned {} unanchored TOC nodes " + "(suffix / no match_overrides → no TOC)", + removed, + ) + return out, removed + + +def toc_range_start(hierarchy: dict[str, Any]) -> int | None: + toc_range = hierarchy.get("toc_range") + if not isinstance(toc_range, (list, tuple)) or not toc_range: + return None + try: + return int(toc_range[0]) + except (TypeError, ValueError): + return None + + +def toc_range_end(hierarchy: dict[str, Any]) -> int | None: + toc_range = hierarchy.get("toc_range") + if not isinstance(toc_range, (list, tuple)) or not toc_range: + return None + try: + return int(toc_range[-1]) + except (TypeError, ValueError): + return None + + +# ── Null-page parent locate (compact-strict + RTL visual) ─────────────────── + +_NULL_PARENT_VISUAL_CONFIDENCE = 0.6 + + +def locate_null_page_parent_overrides( + *, + nodes: list[TitleNode], + match_overrides: dict[tuple[str, ...], TitleMatch], + page_texts: dict[int, str], + body_pages: list[int], + ctx: ToolContext | None, +) -> tuple[dict[tuple[str, ...], TitleMatch], list[dict[str, Any]]]: + """Locate TOC parents with ``printed_page=None`` into ``match_overrides``. + + Window for parent P: ``[last leaf start under previous same-level sibling, + first leaf start under P]``. Text path is compact→strict unique page; on + miss/ambiguity, scan right→left with ``verify_section_page_choice``. + + Returns ``(overrides, report)`` where *report* lists every null-page parent + attempt (for debug / LLM-call accounting). + """ + if not nodes or not body_pages: + return dict(match_overrides), [] + + out = dict(match_overrides) + body_set = set(body_pages) + parent_scope_start = body_pages[0] + report: list[dict[str, Any]] = [] + + def walk( + sibling_nodes: list[TitleNode], + parent_titles: tuple[str, ...], + scope_start: int, + ) -> None: + for index, node in enumerate(sibling_nodes): + path_titles = (*parent_titles, node.title) + if ( + node.children + and node.printed_page is None + and path_titles not in out + ): + if index > 0: + left = last_leaf_start_under( + sibling_nodes[index - 1], parent_titles, out + ) + if left is None: + left = scope_start + else: + left = scope_start + right = first_leaf_start_under(node, parent_titles, out) + entry: dict[str, Any] = { + "path_titles": list(path_titles), + "title": node.title, + "printed_page": None, + "window": None, + "result": "skipped_no_right", + "page": None, + "accept": None, + "visual_verify_calls": 0, + } + if right is None or right < left: + report.append(entry) + logger.info( + "[structure_anchoring] null-page parent skipped: " + "title={!r} reason=no_located_first_child left={}", + node.title, + left, + ) + else: + entry["window"] = [left, right] + scope_pages = [ + page for page in body_pages if left <= page <= right + ] + match = locate_title_compact_strict( + node.title, + scope_pages=scope_pages, + page_texts=page_texts, + ) + visual_calls = 0 + if match is None and ctx is not None: + match, visual_calls = _visual_rtl_locate_parent( + title=node.title, + left=left, + right=right, + body_set=body_set, + ctx=ctx, + ) + entry["visual_verify_calls"] = visual_calls + if match is not None and match.page in body_set: + out[path_titles] = match + entry["result"] = str(match.evidence.get("accept") or match.source) + entry["page"] = match.page + entry["accept"] = match.evidence.get("accept") + logger.info( + "[structure_anchoring] null-page parent located: " + "title={!r} page={} window={} accept={} visual_calls={}", + node.title, + match.page, + [left, right], + match.evidence.get("accept"), + visual_calls, + ) + else: + entry["result"] = "unresolved" + logger.info( + "[structure_anchoring] null-page parent unresolved: " + "title={!r} window={} visual_calls={}", + node.title, + [left, right], + visual_calls, + ) + report.append(entry) + if node.children: + child_scope_start = ( + out[path_titles].page if path_titles in out else scope_start + ) + walk(node.children, path_titles, child_scope_start) + + walk(nodes, (), parent_scope_start) + logger.info( + "[structure_anchoring] null-page parent locate summary: " + "attempted={} located={} unresolved={} visual_verify_calls={}", + len(report), + sum(1 for row in report if row.get("page") is not None), + sum(1 for row in report if row.get("result") == "unresolved"), + sum(int(row.get("visual_verify_calls") or 0) for row in report), + ) + return out, report + + +def _visual_rtl_locate_parent( + *, + title: str, + left: int, + right: int, + body_set: set[int], + ctx: ToolContext, +) -> tuple[TitleMatch | None, int]: + """Confirm parent title from right boundary toward left via VLM verify.""" + visual_calls = 0 + for page in range(right, left - 1, -1): + if page not in body_set: + continue + candidate = TitleMatch( + page=page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[page], + evidence={"null_page_parent_probe": True}, + ) + visual_calls += 1 + result = verify_section_page_choice( + ctx=ctx, + title=title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + selected = result.get("selected_page") + confidence = float(result.get("confidence") or 0.0) + if selected != page or confidence < _NULL_PARENT_VISUAL_CONFIDENCE: + continue + if result.get("source") == "agent_vlm": + return ( + TitleMatch( + page=page, + confidence=confidence, + source="agent_vlm", + matched_line="", + score=confidence, + candidates=[page], + evidence={ + "accept": "visual_rtl", + "reason": result.get("reason", ""), + "visual_verify_calls": visual_calls, + }, + ), + visual_calls, + ) + return ( + TitleMatch( + page=page, + confidence=confidence, + source="agent_heuristic", + matched_line="", + score=confidence, + candidates=[page], + evidence={ + "accept": "visual_rtl", + "reason": result.get("reason", ""), + "visual_verify_calls": visual_calls, + }, + ), + visual_calls, + ) + return None, visual_calls + + +# ── Offset-guided bulk anchoring with recursive recalibrate (Phase-2) ─────── + +_TAIL_VERIFY_CONFIDENCE_THRESHOLD = 0.6 +_MAX_RECALIBRATE_DEPTH = 5 +_MAX_RECALIBRATE_DELTA = 5 + + +def _verify_offset_tail( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, +) -> bool: + """VLM-verify that the offset holds for the last leaf entry (Theorem 1). + + If head offset == tail offset, monotonicity guarantees all intermediate + entries share the same offset. + + Prefers a tail leaf whose expected page is strictly less than page_count + (boundary pages are unreliable for VLM verification). + """ + tail_leaves = [ + (path, node) for path, node in reversed(leaves) if node.printed_page is not None + ] + if not tail_leaves: + return True + + # Prefer non-boundary: printed_page + offset < page_count + selected = None + for path, node in tail_leaves: + pp = node.printed_page + if pp is None: + continue + expected = pp + offset + if 1 <= expected < page_count: + selected = (path, node) + break + if selected is None: + # All leaves are at the boundary; fall back to the last one + selected = tail_leaves[0] + + path, node = selected + printed_page = node.printed_page + if printed_page is None: + return True + expected_page = printed_page + offset + if expected_page < 1 or expected_page > page_count: + return False + + candidate = TitleMatch( + page=expected_page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[expected_page], + evidence={"tail_verify_probe": True}, + ) + result = verify_section_page_choice( + ctx=ctx, + title=node.title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + confirmed = ( + result.get("selected_page") == expected_page + and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD + ) + logger.info( + "[structure_anchoring] tail verify: title={!r} expected_page={} confirmed={} confidence={}", + node.title, + expected_page, + confirmed, + result.get("confidence", 0), + ) + return confirmed + + +def _vlm_confirm_single_page( + *, + ctx: ToolContext, + title: str, + expected_page: int, + page_count: int, +) -> bool: + """Single-page VLM confirmation for binary search steps.""" + if expected_page < 1 or expected_page > page_count: + return False + candidate = TitleMatch( + page=expected_page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[expected_page], + evidence={"bisect_probe": True}, + ) + result = verify_section_page_choice( + ctx=ctx, + title=title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + return ( + result.get("selected_page") == expected_page + and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD + ) + + +def _bisect_offset_breakpoint( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, +) -> int: + """Binary search for the last leaf index where offset is valid. O(log n) VLM calls.""" + lo, hi = 0, len(leaves) - 1 + while lo < hi: + mid = (lo + hi + 1) // 2 + _, node = leaves[mid] + if node.printed_page is None: + hi = mid - 1 + continue + expected = node.printed_page + offset + if _vlm_confirm_single_page( + ctx=ctx, title=node.title, expected_page=expected, page_count=page_count + ): + lo = mid + else: + hi = mid - 1 + logger.info( + "[structure_anchoring] bisect breakpoint: last_valid_index={} / total={}", + lo, + len(leaves), + ) + return lo + + +def bulk_offset_matches( + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, +) -> dict[tuple[str, ...], TitleMatch]: + """Generate TitleMatch overrides for all leaves using offset. No VLM calls.""" + matches: dict[tuple[str, ...], TitleMatch] = {} + for path_titles, node in leaves: + if node.printed_page is None: + continue + page = node.printed_page + offset + matches[path_titles] = TitleMatch( + page=page, + confidence=0.88, + source="agent_vlm", + matched_line="", + score=0.88, + candidates=[page], + evidence={ + "bulk_offset": True, + "offset": offset, + "printed_page": node.printed_page, + }, + ) + return matches + + +def _recalibrate_after_breakpoint( + *, + entry_node: TitleNode, + old_offset: int, + ctx: ToolContext, + page_count: int, +) -> int | None: + """Probe offsets old_offset+1, +2, ... to find new offset after breakpoint. + + Monotonicity guarantees new offset > old offset, so search space is tiny. + """ + entry_printed_page = entry_node.printed_page + if entry_printed_page is None: + return None + for delta in range(1, _MAX_RECALIBRATE_DELTA + 1): + new_offset = old_offset + delta + if _vlm_confirm_single_page( + ctx=ctx, + title=entry_node.title, + expected_page=entry_printed_page + new_offset, + page_count=page_count, + ): + logger.info( + "[structure_anchoring] recalibrate: title={!r} new_offset={} (delta=+{})", + entry_node.title, + new_offset, + delta, + ) + return new_offset + return None + + +def offset_guided_anchoring( + *, + nodes: list[TitleNode], + offset: int, + ctx: ToolContext, + page_count: int, + calibration_overrides: dict[tuple[str, ...], TitleMatch], +) -> dict[tuple[str, ...], TitleMatch] | None: + """Offset-guided bulk anchoring with recursive recalibrate on breakpoints. + + Strategy: + 1. Tail verify last leaf with current offset + 2. If pass → bulk apply all leaves (Theorem 1) + 3. If fail → binary search for breakpoint + 4. Bulk apply leaves before breakpoint + 5. Recalibrate: probe remaining[0] with offset+1, +2, ... (monotonicity) + 6. Recurse on remaining segment with new offset + 7. If recalibrate fails → return partial (caller falls back for remainder) + + Returns match_overrides for all anchored leaves, or None for full fallback. + """ + leaves = [ + (path, node) + for path, node in iter_leaf_title_nodes(nodes) + if node.printed_page is not None + ] + if not leaves: + return dict(calibration_overrides) or None + + all_matches: dict[tuple[str, ...], TitleMatch] = {} + all_matches.update(calibration_overrides) + + # Single-leaf regimes (roman front-matter, F-1 appendix, …) still get a + # deterministic printed→physical override; Phase-1 already calibrated them. + if len(leaves) == 1: + all_matches.update(bulk_offset_matches(leaves, offset)) + else: + _anchor_segment_recursive( + leaves=leaves, + offset=offset, + ctx=ctx, + page_count=page_count, + matches=all_matches, + depth=0, + ) + + if not all_matches: + return None + + logger.info( + "[structure_anchoring] offset bulk anchoring: {} / {} leaves anchored", + len(all_matches), + len(leaves), + ) + return all_matches + + +def _anchor_segment_recursive( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, + matches: dict[tuple[str, ...], TitleMatch], + depth: int, +) -> None: + """Recursively anchor a segment of leaves, handling multiple breakpoints.""" + if not leaves or depth >= _MAX_RECALIBRATE_DEPTH: + return + + if _verify_offset_tail(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count): + bulk = bulk_offset_matches(leaves, offset) + matches.update(bulk) + return + + bp = _bisect_offset_breakpoint(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count) + confirmed_leaves = leaves[: bp + 1] + if confirmed_leaves: + bulk = bulk_offset_matches(confirmed_leaves, offset) + matches.update(bulk) + + remaining = leaves[bp + 1:] + if not remaining: + return + + _, first_remaining_node = remaining[0] + new_offset = _recalibrate_after_breakpoint( + entry_node=first_remaining_node, + old_offset=offset, + ctx=ctx, + page_count=page_count, + ) + if new_offset is None: + return + + _anchor_segment_recursive( + leaves=remaining, + offset=new_offset, + ctx=ctx, + page_count=page_count, + matches=matches, + depth=depth + 1, + ) + + +@dataclass +class SkeletonAnchor: + offset: int | None + offset_status: str + match_overrides: dict[tuple[str, ...], TitleMatch] + null_page_report: list[dict[str, Any]] + bulk_count: int + pruned_count: int = 0 + locate_agent: str = "offset_only" + + +def serialize_title_match(match: TitleMatch) -> dict[str, Any]: + return { + "page": match.page, + "confidence": match.confidence, + "source": match.source, + "matched_line": match.matched_line, + "score": match.score, + "candidates": list(match.candidates), + "evidence": dict(match.evidence or {}), + } + + +def serialize_skeleton_anchor(anchor: SkeletonAnchor) -> dict[str, Any]: + """JSON-friendly SkeletonAnchor (path tuples joined by ' / ').""" + overrides: dict[str, Any] = {} + for path, match in (anchor.match_overrides or {}).items(): + key = " / ".join(str(part) for part in path) + overrides[key] = serialize_title_match(match) + return { + "offset": anchor.offset, + "offset_status": anchor.offset_status, + "match_overrides": overrides, + "null_page_report": list(anchor.null_page_report or []), + "bulk_count": int(anchor.bulk_count or 0), + "pruned_count": int(anchor.pruned_count or 0), + "locate_agent": anchor.locate_agent, + } + + +def deserialize_title_match(data: dict[str, Any]) -> TitleMatch: + return TitleMatch( + page=int(data["page"]), + confidence=float(data.get("confidence") or 0.0), + source=data.get("source") or "agent_vlm", # type: ignore[arg-type] + matched_line=str(data.get("matched_line") or ""), + score=float(data.get("score") or 0.0), + candidates=[int(p) for p in (data.get("candidates") or [])], + evidence=dict(data.get("evidence") or {}), + ) + + +def deserialize_skeleton_anchor(data: dict[str, Any]) -> SkeletonAnchor: + raw_overrides = data.get("match_overrides") or {} + overrides: dict[tuple[str, ...], TitleMatch] = {} + if isinstance(raw_overrides, dict): + for key, value in raw_overrides.items(): + if not isinstance(value, dict): + continue + if isinstance(key, str): + path = tuple(part.strip() for part in key.split(" / ") if part.strip()) + elif isinstance(key, (list, tuple)): + path = tuple(str(part) for part in key) + else: + continue + if path: + overrides[path] = deserialize_title_match(value) + return SkeletonAnchor( + offset=data.get("offset") if data.get("offset") is None else int(data["offset"]), + offset_status=str(data.get("offset_status") or "failed"), + match_overrides=overrides, + null_page_report=list(data.get("null_page_report") or []), + bulk_count=int(data.get("bulk_count") or 0), + pruned_count=int(data.get("pruned_count") or 0), + locate_agent=str(data.get("locate_agent") or "offset_only"), + ) + + +def anchor_hierarchy_from_offset( + *, + nodes: list[TitleNode], + offset_hint: int | None, + calibration_overrides: dict[tuple[str, ...], TitleMatch] | None = None, + page_texts: dict[int, str], + body_pages: list[int], + page_count: int, + ctx: ToolContext | None, +) -> tuple[list[TitleNode], SkeletonAnchor]: + """Production prune → bulk → null-page given a precomputed offset. + + Phase-2 entry after Agent ``calibrate_offset`` (Phase-1). + """ + seed_overrides = dict(calibration_overrides or {}) + pruned_count = 0 + working = nodes + if offset_hint is not None: + working, pruned_count = prune_out_of_scope_nodes( + working, offset=offset_hint, page_count=page_count + ) + + offset_matches: dict[tuple[str, ...], TitleMatch] | None = None + if offset_hint is not None and ctx is not None and working: + offset_matches = offset_guided_anchoring( + nodes=working, + offset=offset_hint, + ctx=ctx, + page_count=page_count, + calibration_overrides=seed_overrides, + ) + + if offset_matches is not None: + match_overrides = offset_matches + locate_agent = "offset_guided_bulk" + bulk_count = len(offset_matches) + else: + match_overrides = seed_overrides + locate_agent = "offset_only" + bulk_count = 0 + + working, unanchored_removed = prune_unanchored_toc_leaves( + working, match_overrides=match_overrides + ) + pruned_count += unanchored_removed + + match_overrides, null_page_report = locate_null_page_parent_overrides( + nodes=working, + match_overrides=match_overrides, + page_texts=page_texts, + body_pages=body_pages, + ctx=ctx, + ) + + if offset_hint is None: + offset_status = "failed" if ctx is not None else "skipped" + else: + offset_status = "ok" + + return working, SkeletonAnchor( + offset=offset_hint, + offset_status=offset_status, + match_overrides=match_overrides, + null_page_report=null_page_report, + bulk_count=bulk_count, + pruned_count=pruned_count, + locate_agent=locate_agent, + ) diff --git a/apps/worker/app/services/document_agent/structure/structure_anchoring.py b/apps/worker/app/services/document_agent/structure/structure_anchoring.py index 24fb2d229..257ecb28f 100644 --- a/apps/worker/app/services/document_agent/structure/structure_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/structure_anchoring.py @@ -1,546 +1,25 @@ -"""Shared hierarchy anchoring: Phase-2 bulk/bisect/null-page + SkeletonAnchor. +"""Compatibility exports for hierarchy anchoring. -Phase-1 offset discovery lives in ``document_agent.agents.calibration``. -``anchor_hierarchy`` composes Phase-1 + Phase-2 for production callers. +Low-level anchoring primitives live in :mod:`anchoring_primitives`; the +calibration-owned orchestrator is resolved lazily to keep imports acyclic. """ from __future__ import annotations -from dataclasses import dataclass, replace +from importlib import import_module from typing import Any -from app.services.document_agent.manifest import ToolContext -from app.services.document_agent.structure.hierarchy_locator import ( +from app.services.document_agent.structure.anchoring_primitives import * # noqa: F403 +from app.services.document_agent.structure.anchoring_primitives import ( + SkeletonAnchor, TitleMatch, TitleNode, - first_leaf_start_under, - iter_leaf_title_nodes, - last_leaf_start_under, - locate_title_compact_strict, + ToolContext, ) +from app.services.document_agent.structure import anchoring_primitives as _anchoring from app.services.document_agent.structure.page_locate_agent import ( verify_section_page_choice, ) -from loguru import logger - - -def prune_out_of_scope_nodes( - nodes: list[TitleNode], - *, - offset: int, - page_count: int, -) -> tuple[list[TitleNode], int]: - """Remove leaf nodes whose printed_page + offset exceeds page_count. - - Bottom-up: prune out-of-scope leaves, then remove intermediate nodes - that become childless after pruning. Returns (pruned_tree, removed_count). - """ - removed = 0 - - def _prune(node: TitleNode) -> TitleNode | None: - nonlocal removed - if not node.children: - if node.printed_page is not None: - expected = node.printed_page + offset - if expected > page_count or expected < 1: - removed += 1 - return None - return node - pruned_children = [] - for child in node.children: - result = _prune(child) - if result is not None: - pruned_children.append(result) - if not pruned_children: - removed += 1 - return None - return replace(node, children=pruned_children) - - pruned = [] - for node in nodes: - result = _prune(node) - if result is not None: - pruned.append(result) - - if removed: - logger.info( - "[structure_anchoring] pruned {} out-of-scope TOC nodes " - "(printed_page + offset={} exceeds page_count={})", - removed, - offset, - page_count, - ) - - return pruned, removed - - -def prune_unanchored_toc_leaves( - nodes: list[TitleNode], - *, - match_overrides: dict[tuple[str, ...], TitleMatch], -) -> tuple[list[TitleNode], int]: - """Remove TOC leaves that have no physical ``match_overrides`` entry. - - Implements Phase-2 ``suffix = no TOC``: after bulk/bisect/recalibrate, any - leaf that was not successfully anchored is dropped from the coarse tree - instead of sticky ``inherited_unlocated`` ranges. Childless parents are - removed unless they themselves have an override. - """ - removed = 0 - - def _prune( - node: TitleNode, parent_titles: tuple[str, ...] - ) -> TitleNode | None: - nonlocal removed - path = (*parent_titles, node.title) - if node.children: - children: list[TitleNode] = [] - for child in node.children: - kept = _prune(child, path) - if kept is not None: - children.append(kept) - if children: - return replace(node, children=children) - if path in match_overrides: - return replace(node, children=[]) - removed += 1 - return None - if path in match_overrides: - return node - removed += 1 - return None - - out: list[TitleNode] = [] - for node in nodes: - kept = _prune(node, ()) - if kept is not None: - out.append(kept) - - if removed: - logger.info( - "[structure_anchoring] pruned {} unanchored TOC nodes " - "(suffix / no match_overrides → no TOC)", - removed, - ) - return out, removed - - -def toc_range_start(hierarchy: dict[str, Any]) -> int | None: - toc_range = hierarchy.get("toc_range") - if not isinstance(toc_range, (list, tuple)) or not toc_range: - return None - try: - return int(toc_range[0]) - except (TypeError, ValueError): - return None - - -def toc_range_end(hierarchy: dict[str, Any]) -> int | None: - toc_range = hierarchy.get("toc_range") - if not isinstance(toc_range, (list, tuple)) or not toc_range: - return None - try: - return int(toc_range[-1]) - except (TypeError, ValueError): - return None - - -# ── Null-page parent locate (compact-strict + RTL visual) ─────────────────── - -_NULL_PARENT_VISUAL_CONFIDENCE = 0.6 - - -def locate_null_page_parent_overrides( - *, - nodes: list[TitleNode], - match_overrides: dict[tuple[str, ...], TitleMatch], - page_texts: dict[int, str], - body_pages: list[int], - ctx: ToolContext | None, -) -> tuple[dict[tuple[str, ...], TitleMatch], list[dict[str, Any]]]: - """Locate TOC parents with ``printed_page=None`` into ``match_overrides``. - - Window for parent P: ``[last leaf start under previous same-level sibling, - first leaf start under P]``. Text path is compact→strict unique page; on - miss/ambiguity, scan right→left with ``verify_section_page_choice``. - - Returns ``(overrides, report)`` where *report* lists every null-page parent - attempt (for debug / LLM-call accounting). - """ - if not nodes or not body_pages: - return dict(match_overrides), [] - - out = dict(match_overrides) - body_set = set(body_pages) - parent_scope_start = body_pages[0] - report: list[dict[str, Any]] = [] - - def walk( - sibling_nodes: list[TitleNode], - parent_titles: tuple[str, ...], - scope_start: int, - ) -> None: - for index, node in enumerate(sibling_nodes): - path_titles = (*parent_titles, node.title) - if ( - node.children - and node.printed_page is None - and path_titles not in out - ): - if index > 0: - left = last_leaf_start_under( - sibling_nodes[index - 1], parent_titles, out - ) - if left is None: - left = scope_start - else: - left = scope_start - right = first_leaf_start_under(node, parent_titles, out) - entry: dict[str, Any] = { - "path_titles": list(path_titles), - "title": node.title, - "printed_page": None, - "window": None, - "result": "skipped_no_right", - "page": None, - "accept": None, - "visual_verify_calls": 0, - } - if right is None or right < left: - report.append(entry) - logger.info( - "[structure_anchoring] null-page parent skipped: " - "title={!r} reason=no_located_first_child left={}", - node.title, - left, - ) - else: - entry["window"] = [left, right] - scope_pages = [ - page for page in body_pages if left <= page <= right - ] - match = locate_title_compact_strict( - node.title, - scope_pages=scope_pages, - page_texts=page_texts, - ) - visual_calls = 0 - if match is None and ctx is not None: - match, visual_calls = _visual_rtl_locate_parent( - title=node.title, - left=left, - right=right, - body_set=body_set, - ctx=ctx, - ) - entry["visual_verify_calls"] = visual_calls - if match is not None and match.page in body_set: - out[path_titles] = match - entry["result"] = str(match.evidence.get("accept") or match.source) - entry["page"] = match.page - entry["accept"] = match.evidence.get("accept") - logger.info( - "[structure_anchoring] null-page parent located: " - "title={!r} page={} window={} accept={} visual_calls={}", - node.title, - match.page, - [left, right], - match.evidence.get("accept"), - visual_calls, - ) - else: - entry["result"] = "unresolved" - logger.info( - "[structure_anchoring] null-page parent unresolved: " - "title={!r} window={} visual_calls={}", - node.title, - [left, right], - visual_calls, - ) - report.append(entry) - if node.children: - child_scope_start = ( - out[path_titles].page if path_titles in out else scope_start - ) - walk(node.children, path_titles, child_scope_start) - - walk(nodes, (), parent_scope_start) - logger.info( - "[structure_anchoring] null-page parent locate summary: " - "attempted={} located={} unresolved={} visual_verify_calls={}", - len(report), - sum(1 for row in report if row.get("page") is not None), - sum(1 for row in report if row.get("result") == "unresolved"), - sum(int(row.get("visual_verify_calls") or 0) for row in report), - ) - return out, report - - -def _visual_rtl_locate_parent( - *, - title: str, - left: int, - right: int, - body_set: set[int], - ctx: ToolContext, -) -> tuple[TitleMatch | None, int]: - """Confirm parent title from right boundary toward left via VLM verify.""" - visual_calls = 0 - for page in range(right, left - 1, -1): - if page not in body_set: - continue - candidate = TitleMatch( - page=page, - confidence=0.4, - source="agent_heuristic", - matched_line="", - score=0.4, - candidates=[page], - evidence={"null_page_parent_probe": True}, - ) - visual_calls += 1 - result = verify_section_page_choice( - ctx=ctx, - title=title, - candidate_matches=[candidate], - candidate_page_cap=1, - ) - selected = result.get("selected_page") - confidence = float(result.get("confidence") or 0.0) - if selected != page or confidence < _NULL_PARENT_VISUAL_CONFIDENCE: - continue - if result.get("source") == "agent_vlm": - return ( - TitleMatch( - page=page, - confidence=confidence, - source="agent_vlm", - matched_line="", - score=confidence, - candidates=[page], - evidence={ - "accept": "visual_rtl", - "reason": result.get("reason", ""), - "visual_verify_calls": visual_calls, - }, - ), - visual_calls, - ) - return ( - TitleMatch( - page=page, - confidence=confidence, - source="agent_heuristic", - matched_line="", - score=confidence, - candidates=[page], - evidence={ - "accept": "visual_rtl", - "reason": result.get("reason", ""), - "visual_verify_calls": visual_calls, - }, - ), - visual_calls, - ) - return None, visual_calls - - -# ── Offset-guided bulk anchoring with recursive recalibrate (Phase-2) ─────── - -_TAIL_VERIFY_CONFIDENCE_THRESHOLD = 0.6 -_MAX_RECALIBRATE_DEPTH = 5 -_MAX_RECALIBRATE_DELTA = 5 - - -def _verify_offset_tail( - *, - leaves: list[tuple[tuple[str, ...], TitleNode]], - offset: int, - ctx: ToolContext, - page_count: int, -) -> bool: - """VLM-verify that the offset holds for the last leaf entry (Theorem 1). - - If head offset == tail offset, monotonicity guarantees all intermediate - entries share the same offset. - - Prefers a tail leaf whose expected page is strictly less than page_count - (boundary pages are unreliable for VLM verification). - """ - tail_leaves = [ - (path, node) for path, node in reversed(leaves) if node.printed_page is not None - ] - if not tail_leaves: - return True - - # Prefer non-boundary: printed_page + offset < page_count - selected = None - for path, node in tail_leaves: - pp = node.printed_page - if pp is None: - continue - expected = pp + offset - if 1 <= expected < page_count: - selected = (path, node) - break - if selected is None: - # All leaves are at the boundary; fall back to the last one - selected = tail_leaves[0] - - path, node = selected - printed_page = node.printed_page - if printed_page is None: - return True - expected_page = printed_page + offset - if expected_page < 1 or expected_page > page_count: - return False - - candidate = TitleMatch( - page=expected_page, - confidence=0.4, - source="agent_heuristic", - matched_line="", - score=0.4, - candidates=[expected_page], - evidence={"tail_verify_probe": True}, - ) - result = verify_section_page_choice( - ctx=ctx, - title=node.title, - candidate_matches=[candidate], - candidate_page_cap=1, - ) - confirmed = ( - result.get("selected_page") == expected_page - and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD - ) - logger.info( - "[structure_anchoring] tail verify: title={!r} expected_page={} confirmed={} confidence={}", - node.title, - expected_page, - confirmed, - result.get("confidence", 0), - ) - return confirmed - - -def _vlm_confirm_single_page( - *, - ctx: ToolContext, - title: str, - expected_page: int, - page_count: int, -) -> bool: - """Single-page VLM confirmation for binary search steps.""" - if expected_page < 1 or expected_page > page_count: - return False - candidate = TitleMatch( - page=expected_page, - confidence=0.4, - source="agent_heuristic", - matched_line="", - score=0.4, - candidates=[expected_page], - evidence={"bisect_probe": True}, - ) - result = verify_section_page_choice( - ctx=ctx, - title=title, - candidate_matches=[candidate], - candidate_page_cap=1, - ) - return ( - result.get("selected_page") == expected_page - and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD - ) - - -def _bisect_offset_breakpoint( - *, - leaves: list[tuple[tuple[str, ...], TitleNode]], - offset: int, - ctx: ToolContext, - page_count: int, -) -> int: - """Binary search for the last leaf index where offset is valid. O(log n) VLM calls.""" - lo, hi = 0, len(leaves) - 1 - while lo < hi: - mid = (lo + hi + 1) // 2 - _, node = leaves[mid] - if node.printed_page is None: - hi = mid - 1 - continue - expected = node.printed_page + offset - if _vlm_confirm_single_page( - ctx=ctx, title=node.title, expected_page=expected, page_count=page_count - ): - lo = mid - else: - hi = mid - 1 - logger.info( - "[structure_anchoring] bisect breakpoint: last_valid_index={} / total={}", - lo, - len(leaves), - ) - return lo - - -def bulk_offset_matches( - leaves: list[tuple[tuple[str, ...], TitleNode]], - offset: int, -) -> dict[tuple[str, ...], TitleMatch]: - """Generate TitleMatch overrides for all leaves using offset. No VLM calls.""" - matches: dict[tuple[str, ...], TitleMatch] = {} - for path_titles, node in leaves: - if node.printed_page is None: - continue - page = node.printed_page + offset - matches[path_titles] = TitleMatch( - page=page, - confidence=0.88, - source="agent_vlm", - matched_line="", - score=0.88, - candidates=[page], - evidence={ - "bulk_offset": True, - "offset": offset, - "printed_page": node.printed_page, - }, - ) - return matches - - -def _recalibrate_after_breakpoint( - *, - entry_node: TitleNode, - old_offset: int, - ctx: ToolContext, - page_count: int, -) -> int | None: - """Probe offsets old_offset+1, +2, ... to find new offset after breakpoint. - - Monotonicity guarantees new offset > old offset, so search space is tiny. - """ - entry_printed_page = entry_node.printed_page - if entry_printed_page is None: - return None - for delta in range(1, _MAX_RECALIBRATE_DELTA + 1): - new_offset = old_offset + delta - if _vlm_confirm_single_page( - ctx=ctx, - title=entry_node.title, - expected_page=entry_printed_page + new_offset, - page_count=page_count, - ): - logger.info( - "[structure_anchoring] recalibrate: title={!r} new_offset={} (delta=+{})", - entry_node.title, - new_offset, - delta, - ) - return new_offset - return None def offset_guided_anchoring( @@ -551,249 +30,19 @@ def offset_guided_anchoring( page_count: int, calibration_overrides: dict[tuple[str, ...], TitleMatch], ) -> dict[tuple[str, ...], TitleMatch] | None: - """Offset-guided bulk anchoring with recursive recalibrate on breakpoints. - - Strategy: - 1. Tail verify last leaf with current offset - 2. If pass → bulk apply all leaves (Theorem 1) - 3. If fail → binary search for breakpoint - 4. Bulk apply leaves before breakpoint - 5. Recalibrate: probe remaining[0] with offset+1, +2, ... (monotonicity) - 6. Recurse on remaining segment with new offset - 7. If recalibrate fails → return partial (caller falls back for remainder) - - Returns match_overrides for all anchored leaves, or None for full fallback. - """ - leaves = [ - (path, node) - for path, node in iter_leaf_title_nodes(nodes) - if node.printed_page is not None - ] - if not leaves: - return dict(calibration_overrides) or None - - all_matches: dict[tuple[str, ...], TitleMatch] = {} - all_matches.update(calibration_overrides) - - # Single-leaf regimes (roman front-matter, F-1 appendix, …) still get a - # deterministic printed→physical override; Phase-1 already calibrated them. - if len(leaves) == 1: - all_matches.update(bulk_offset_matches(leaves, offset)) - else: - _anchor_segment_recursive( - leaves=leaves, + """Forward phase-2 anchoring while preserving the historical patch seam.""" + original = _anchoring.verify_section_page_choice + _anchoring.verify_section_page_choice = verify_section_page_choice + try: + return _anchoring.offset_guided_anchoring( + nodes=nodes, offset=offset, ctx=ctx, page_count=page_count, - matches=all_matches, - depth=0, - ) - - if not all_matches: - return None - - logger.info( - "[structure_anchoring] offset bulk anchoring: {} / {} leaves anchored", - len(all_matches), - len(leaves), - ) - return all_matches - - -def _anchor_segment_recursive( - *, - leaves: list[tuple[tuple[str, ...], TitleNode]], - offset: int, - ctx: ToolContext, - page_count: int, - matches: dict[tuple[str, ...], TitleMatch], - depth: int, -) -> None: - """Recursively anchor a segment of leaves, handling multiple breakpoints.""" - if not leaves or depth >= _MAX_RECALIBRATE_DEPTH: - return - - if _verify_offset_tail(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count): - bulk = bulk_offset_matches(leaves, offset) - matches.update(bulk) - return - - bp = _bisect_offset_breakpoint(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count) - confirmed_leaves = leaves[: bp + 1] - if confirmed_leaves: - bulk = bulk_offset_matches(confirmed_leaves, offset) - matches.update(bulk) - - remaining = leaves[bp + 1:] - if not remaining: - return - - _, first_remaining_node = remaining[0] - new_offset = _recalibrate_after_breakpoint( - entry_node=first_remaining_node, - old_offset=offset, - ctx=ctx, - page_count=page_count, - ) - if new_offset is None: - return - - _anchor_segment_recursive( - leaves=remaining, - offset=new_offset, - ctx=ctx, - page_count=page_count, - matches=matches, - depth=depth + 1, - ) - - -@dataclass -class SkeletonAnchor: - offset: int | None - offset_status: str - match_overrides: dict[tuple[str, ...], TitleMatch] - null_page_report: list[dict[str, Any]] - bulk_count: int - pruned_count: int = 0 - locate_agent: str = "offset_only" - - -def serialize_title_match(match: TitleMatch) -> dict[str, Any]: - return { - "page": match.page, - "confidence": match.confidence, - "source": match.source, - "matched_line": match.matched_line, - "score": match.score, - "candidates": list(match.candidates), - "evidence": dict(match.evidence or {}), - } - - -def serialize_skeleton_anchor(anchor: SkeletonAnchor) -> dict[str, Any]: - """JSON-friendly SkeletonAnchor (path tuples joined by ' / ').""" - overrides: dict[str, Any] = {} - for path, match in (anchor.match_overrides or {}).items(): - key = " / ".join(str(part) for part in path) - overrides[key] = serialize_title_match(match) - return { - "offset": anchor.offset, - "offset_status": anchor.offset_status, - "match_overrides": overrides, - "null_page_report": list(anchor.null_page_report or []), - "bulk_count": int(anchor.bulk_count or 0), - "pruned_count": int(anchor.pruned_count or 0), - "locate_agent": anchor.locate_agent, - } - - -def deserialize_title_match(data: dict[str, Any]) -> TitleMatch: - return TitleMatch( - page=int(data["page"]), - confidence=float(data.get("confidence") or 0.0), - source=data.get("source") or "agent_vlm", # type: ignore[arg-type] - matched_line=str(data.get("matched_line") or ""), - score=float(data.get("score") or 0.0), - candidates=[int(p) for p in (data.get("candidates") or [])], - evidence=dict(data.get("evidence") or {}), - ) - - -def deserialize_skeleton_anchor(data: dict[str, Any]) -> SkeletonAnchor: - raw_overrides = data.get("match_overrides") or {} - overrides: dict[tuple[str, ...], TitleMatch] = {} - if isinstance(raw_overrides, dict): - for key, value in raw_overrides.items(): - if not isinstance(value, dict): - continue - if isinstance(key, str): - path = tuple(part.strip() for part in key.split(" / ") if part.strip()) - elif isinstance(key, (list, tuple)): - path = tuple(str(part) for part in key) - else: - continue - if path: - overrides[path] = deserialize_title_match(value) - return SkeletonAnchor( - offset=data.get("offset") if data.get("offset") is None else int(data["offset"]), - offset_status=str(data.get("offset_status") or "failed"), - match_overrides=overrides, - null_page_report=list(data.get("null_page_report") or []), - bulk_count=int(data.get("bulk_count") or 0), - pruned_count=int(data.get("pruned_count") or 0), - locate_agent=str(data.get("locate_agent") or "offset_only"), - ) - - -def anchor_hierarchy_from_offset( - *, - nodes: list[TitleNode], - offset_hint: int | None, - calibration_overrides: dict[tuple[str, ...], TitleMatch] | None = None, - page_texts: dict[int, str], - body_pages: list[int], - page_count: int, - ctx: ToolContext | None, -) -> tuple[list[TitleNode], SkeletonAnchor]: - """Production prune → bulk → null-page given a precomputed offset. - - Phase-2 entry after Agent ``calibrate_offset`` (Phase-1). - """ - seed_overrides = dict(calibration_overrides or {}) - pruned_count = 0 - working = nodes - if offset_hint is not None: - working, pruned_count = prune_out_of_scope_nodes( - working, offset=offset_hint, page_count=page_count - ) - - offset_matches: dict[tuple[str, ...], TitleMatch] | None = None - if offset_hint is not None and ctx is not None and working: - offset_matches = offset_guided_anchoring( - nodes=working, - offset=offset_hint, - ctx=ctx, - page_count=page_count, - calibration_overrides=seed_overrides, + calibration_overrides=calibration_overrides, ) - - if offset_matches is not None: - match_overrides = offset_matches - locate_agent = "offset_guided_bulk" - bulk_count = len(offset_matches) - else: - match_overrides = seed_overrides - locate_agent = "offset_only" - bulk_count = 0 - - working, unanchored_removed = prune_unanchored_toc_leaves( - working, match_overrides=match_overrides - ) - pruned_count += unanchored_removed - - match_overrides, null_page_report = locate_null_page_parent_overrides( - nodes=working, - match_overrides=match_overrides, - page_texts=page_texts, - body_pages=body_pages, - ctx=ctx, - ) - - if offset_hint is None: - offset_status = "failed" if ctx is not None else "skipped" - else: - offset_status = "ok" - - return working, SkeletonAnchor( - offset=offset_hint, - offset_status=offset_status, - match_overrides=match_overrides, - null_page_report=null_page_report, - bulk_count=bulk_count, - pruned_count=pruned_count, - locate_agent=locate_agent, - ) + finally: + _anchoring.verify_section_page_choice = original def anchor_hierarchy( @@ -805,45 +54,15 @@ def anchor_hierarchy( page_count: int, ctx: ToolContext | None, ) -> tuple[list[TitleNode], SkeletonAnchor]: - """Run Phase-1 calibrate_offset → multi-regime Phase-2 merge. - - Returns possibly-pruned nodes and the anchor payload. Caller owns - resolve_hierarchy_page_ranges / skeleton assembly. - """ - from app.services.document_agent.agents.calibration.procedure import ( - finalize_calibration_result, - flat_toc_entries, - ) - from app.services.document_agent.agents.calibration.service import ( - calibrate_offset, + """Resolve the calibration-owned orchestration entry point on demand.""" + orchestrator = import_module( + "app.services.document_agent.agents.calibration.orchestrator" ) - - phase1 = calibrate_offset( + return orchestrator.anchor_hierarchy( nodes=nodes, toc_hierarchies=toc_hierarchies, - ctx=ctx, page_texts=page_texts, + body_pages=body_pages, page_count=page_count, - ) - if phase1.status == "failed" and not phase1.regimes and phase1.offset is None: - return anchor_hierarchy_from_offset( - nodes=nodes, - offset_hint=None, - calibration_overrides={}, - page_texts=page_texts, - body_pages=body_pages, - page_count=page_count, - ctx=ctx, - ) - - working, anchor, _finalized = finalize_calibration_result( - result=phase1, - entries=flat_toc_entries(toc_hierarchies), - toc_hierarchies=list(toc_hierarchies or []), ctx=ctx, - page_count=page_count, - page_texts=page_texts, - body_pages=body_pages, - nodes=nodes, ) - return working, anchor diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 6a35e63df..3dfd19c80 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -22,8 +22,8 @@ resolve_hierarchy_page_ranges, ) from app.services.document_agent.agents.calibration import calibrate_offset -from app.services.document_agent.structure.structure_anchoring import ( - anchor_hierarchy, +from app.services.document_agent.agents.calibration.orchestrator import anchor_hierarchy +from app.services.document_agent.structure.anchoring_primitives import ( toc_range_end, toc_range_start, ) From 4a14a611df0ff3a394460c1f8bdf0320f84b7208 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 13 Aug 2026 02:38:57 +0800 Subject: [PATCH 11/12] fix: make anchoring compatibility exports explicit --- .../structure/structure_anchoring.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/apps/worker/app/services/document_agent/structure/structure_anchoring.py b/apps/worker/app/services/document_agent/structure/structure_anchoring.py index 257ecb28f..76472445b 100644 --- a/apps/worker/app/services/document_agent/structure/structure_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/structure_anchoring.py @@ -9,7 +9,19 @@ from importlib import import_module from typing import Any -from app.services.document_agent.structure.anchoring_primitives import * # noqa: F403 +from app.services.document_agent.structure.anchoring_primitives import ( + anchor_hierarchy_from_offset, + bulk_offset_matches, + deserialize_skeleton_anchor, + deserialize_title_match, + locate_null_page_parent_overrides, + prune_out_of_scope_nodes, + prune_unanchored_toc_leaves, + serialize_skeleton_anchor, + serialize_title_match, + toc_range_end, + toc_range_start, +) from app.services.document_agent.structure.anchoring_primitives import ( SkeletonAnchor, TitleMatch, From 4f6a10b9efb08ff9af480dc863fef6080c271d84 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 13 Aug 2026 02:42:49 +0800 Subject: [PATCH 12/12] fix: satisfy lint for anchoring exports --- .../structure/structure_anchoring.py | 48 +++++++++++++------ 1 file changed, 33 insertions(+), 15 deletions(-) diff --git a/apps/worker/app/services/document_agent/structure/structure_anchoring.py b/apps/worker/app/services/document_agent/structure/structure_anchoring.py index 76472445b..b32a68494 100644 --- a/apps/worker/app/services/document_agent/structure/structure_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/structure_anchoring.py @@ -9,30 +9,48 @@ from importlib import import_module from typing import Any -from app.services.document_agent.structure.anchoring_primitives import ( - anchor_hierarchy_from_offset, - bulk_offset_matches, - deserialize_skeleton_anchor, - deserialize_title_match, - locate_null_page_parent_overrides, - prune_out_of_scope_nodes, - prune_unanchored_toc_leaves, - serialize_skeleton_anchor, - serialize_title_match, - toc_range_end, - toc_range_start, -) +from app.services.document_agent.structure import anchoring_primitives as _anchoring from app.services.document_agent.structure.anchoring_primitives import ( SkeletonAnchor, TitleMatch, TitleNode, - ToolContext, ) -from app.services.document_agent.structure import anchoring_primitives as _anchoring +from app.services.document_agent.manifest import ToolContext from app.services.document_agent.structure.page_locate_agent import ( verify_section_page_choice, ) +__all__ = [ + "SkeletonAnchor", + "TitleMatch", + "TitleNode", + "anchor_hierarchy", + "anchor_hierarchy_from_offset", + "bulk_offset_matches", + "deserialize_skeleton_anchor", + "deserialize_title_match", + "locate_null_page_parent_overrides", + "offset_guided_anchoring", + "prune_out_of_scope_nodes", + "prune_unanchored_toc_leaves", + "serialize_skeleton_anchor", + "serialize_title_match", + "toc_range_end", + "toc_range_start", +] + +anchor_hierarchy_from_offset = _anchoring.anchor_hierarchy_from_offset +bulk_offset_matches = _anchoring.bulk_offset_matches +deserialize_skeleton_anchor = _anchoring.deserialize_skeleton_anchor +deserialize_title_match = _anchoring.deserialize_title_match +locate_null_page_parent_overrides = _anchoring.locate_null_page_parent_overrides +prune_out_of_scope_nodes = _anchoring.prune_out_of_scope_nodes +prune_unanchored_toc_leaves = _anchoring.prune_unanchored_toc_leaves +serialize_skeleton_anchor = _anchoring.serialize_skeleton_anchor +serialize_title_match = _anchoring.serialize_title_match +toc_range_end = _anchoring.toc_range_end +toc_range_start = _anchoring.toc_range_start + def offset_guided_anchoring( *,