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..5e4d118e7 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -130,17 +130,17 @@ 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 -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/__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..9910ac271 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/SKILL.md @@ -0,0 +1,101 @@ +# 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`. +- 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 + +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. 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. +5. Submit **candidate** offsets. Do not treat Phase 1 alone as a finished + coarse-structure calibration. + +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. + +## Phase 2 — Completion (deterministic after submit; production path) + +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 + +- 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 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`). + +## 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 token / round budgets 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..852c9991f --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/loop.py @@ -0,0 +1,492 @@ +"""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, 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 ( + 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. +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") + + +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, +) -> 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: + 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_visual_stage", "calibration") + + 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) + 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, + "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 + ), + "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" + ), + }, + "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 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): + 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, + 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 _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_visual_stage": "calibration", + }, + ) + + 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, + ) + if phase1.status == "failed" and not phase1.regimes: + return phase1, {} + + _working, 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..40642bf78 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/procedure.py @@ -0,0 +1,570 @@ +"""Phase-2 completion aligned with production anchoring. + +After the agent submits candidate regime offsets, this module: +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 + +Returns production ``SkeletonAnchor`` plus regime diagnostics for debug payloads. +""" + +from __future__ import annotations + +from dataclasses import replace +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, + 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, + locate_null_page_parent_overrides, + offset_guided_anchoring, + prune_unanchored_toc_leaves, + serialize_skeleton_anchor, +) + +# Re-export under prior names so existing imports keep working. +normalize_kind = normalize_page_kind + + +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: + 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 _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 = [ + idx + for idx, entry in enumerate(entries) + if isinstance(entry, dict) + and classify_page_number_kind(entry.get("page_number")) == kind + ] + + 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(regime.offset), + leaf_start=0, + leaf_end=max(0, len(ok_indices) - 1), + entry_indices=ok_indices, + status="ok", + ) + ] + + out.append( + CalibrationRegime( + kind=kind, + offset=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}; " + f"regime_anchored={len(ok_indices)}" + ), + ) + ) + return out + + +def finalize_calibration_result( + *, + result: CalibrationResult, + entries: list[dict[str, Any]], + toc_hierarchies: list[dict[str, Any]], + ctx: ToolContext | None, + page_count: int, + page_texts: dict[int, str] | None = None, + body_pages: list[int] | None = None, + 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))) + if nodes is None: + from app.services.document_agent.structure.hierarchy_locator import ( + collapse_intermediate_single_child_chains, + ) + + nodes = collapse_intermediate_single_child_chains( + extract_toc_nodes(toc_hierarchies) + ) + + working, anchor = anchor_hierarchy_from_regimes( + nodes=nodes, + 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={} 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, nodes=working + ) + 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 working, 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 (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..d1f743767 --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/service.py @@ -0,0 +1,62 @@ +"""Production calibration entry: Agent Phase-1 offset discovery. + +Returns a full ``CalibrationResult`` (all regimes). Callers run multi-regime +Phase-2 via ``finalize_calibration_result`` / ``anchor_hierarchy``. +""" + +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.types import CalibrationResult +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.structure.hierarchy_locator import 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, +) -> CalibrationResult: + """Discover printed→physical offsets via the calibration SubAgent (Phase 1). + + 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 CalibrationResult(status="failed", notes="ctx missing") + hierarchies = list(toc_hierarchies or []) + if not hierarchies: + return CalibrationResult(status="failed", notes="toc_hierarchies empty") + + 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 CalibrationResult(status="failed", notes=str(exc)) + + logger.info( + "[calibration] Phase-1 status={} regimes={} primary_offset={}", + phase1.status, + len(phase1.regimes), + phase1.offset, + ) + 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 new file mode 100644 index 000000000..37cf26c6a --- /dev/null +++ b/apps/worker/app/services/document_agent/agents/calibration/tools.py @@ -0,0 +1,141 @@ +"""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") + merged.setdefault("visual_stage", "calibration") + 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/budget.py b/apps/worker/app/services/document_agent/budget.py index d0c91abc9..3ecdf867f 100644 --- a/apps/worker/app/services/document_agent/budget.py +++ b/apps/worker/app/services/document_agent/budget.py @@ -11,7 +11,7 @@ "toc_confirm", "coarse_planner", "structural_react", - "page_locate", + "calibration", "page_tagging", ] @@ -173,9 +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": - 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 2f8c699fb..00caba70a 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,11 +66,11 @@ def __init__( ), cap=int(os.environ.get("PARSE_AGENT_STRUCTURAL_REACT_CAP", "64000")), ), - "page_locate": StageEnvelope( + "calibration": StageEnvelope( min_guarantee=int( - os.environ.get("PARSE_AGENT_PAGE_LOCATE_MIN_BUDGET", "0") + os.environ.get("PARSE_AGENT_CALIBRATION_MIN_BUDGET", "12000") ), - cap=int(os.environ.get("PARSE_AGENT_PAGE_LOCATE_CAP", "0")) or None, + cap=int(os.environ.get("PARSE_AGENT_CALIBRATION_CAP", "40000")), ), "page_tagging": StageEnvelope( min_guarantee=int( 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/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/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/structure/hierarchy_locator.py b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py index f33717875..86719c1d1 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() @@ -680,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/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 new file mode 100644 index 000000000..24fb2d229 --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/structure_anchoring.py @@ -0,0 +1,849 @@ +"""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, + ) + + +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 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, + ) + + phase1 = 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/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/__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 index 2254dbe47..7e1b2c7b3 100644 --- a/apps/worker/app/services/document_agent/tools/inspect_pages.py +++ b/apps/worker/app/services/document_agent/tools/inspect_pages.py @@ -1,4 +1,4 @@ -"""Generic VLM inspection tool for selected PDF pages.""" +"""Generic inspect.pages tool: open physical pages, render, answer a question.""" from __future__ import annotations @@ -8,108 +8,164 @@ import time from typing import Any, cast +from loguru import logger + 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,), -) + + +_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``. + + 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() - pages = sorted( - { - int(page) - for page in (args.get("pages") or []) - if 1 <= int(page) <= ctx.blackboard.page_count - } - )[:10] + 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="inspect.pages requires at least one valid page", + error="no valid physical pages in range", 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}" + + 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, ) - 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) + if not rendered: return ToolResult( - status="ok", - payload=payload, + status="error", + error="render failed", 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): + + model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") + if not model: return ToolResult( status="error", - error="insufficient visual budget", + error="vlm_model missing", 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: " + '{"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 pngs: - with open(str(item["png_path"]), "rb") as f: - img_b64 = base64.b64encode(f.read()).decode() + 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}"}} + { + "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=model) + 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=1200, + max_tokens=800, 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, + usage_task=usage_task, ) - 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) + 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="ok", - payload=payload, + status="error", + error=f"vlm failed: {exc}", 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 + + # 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 + 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=tokens_used, + output_summary={"pages": pages, "answer": payload.get("answer")}, + ) 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/app/services/document_agent/tools/propose_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index 1c2737db7..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: @@ -367,61 +341,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, @@ -699,17 +618,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.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 - from app.services.page_memory.skeleton_extractor import _calibrate_offset_via_vlm nodes = extract_toc_nodes(ctx.blackboard.toc_hierarchies) - offset_hint, _ = _calibrate_offset_via_vlm( + 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/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/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/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index ce835659d..6a35e63df 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -14,20 +14,19 @@ 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.agents.calibration import calibrate_offset +from app.services.document_agent.structure.structure_anchoring import ( + anchor_hierarchy, + toc_range_end, + toc_range_start, +) from loguru import logger from shared.services.chunks.path_segments import ( append_document_path, @@ -37,56 +36,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 +110,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), @@ -341,39 +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_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, - 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( *, @@ -423,19 +318,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 +367,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 +376,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 +410,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,13 +423,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_via_vlm( + 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( @@ -1174,38 +457,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), @@ -1219,7 +489,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_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 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/contract/test_structure_anchoring_contract.py b/apps/worker/tests/contract/test_structure_anchoring_contract.py new file mode 100644 index 000000000..d81465f65 --- /dev/null +++ b/apps/worker/tests/contract/test_structure_anchoring_contract.py @@ -0,0 +1,407 @@ +"""Contract tests for structure_anchoring (Phase-2) + calibrate wiring.""" + +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 + +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 + + +@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", + job_id="job-anchor", + blackboard=AgentBlackboard(), + budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), + trace=None, + settings={"vlm_model": "test-vlm"}, + ) + + +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_phase2_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_verify(fake_verify): + matches = anchoring.offset_guided_anchoring( + nodes=leaves, + offset=2, + ctx=ctx, + page_count=30, + calibration_overrides=seed, + ) + 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_uses_calibration_phase1() -> None: + leaves = [_leaf("Only", 2)] + toc_hierarchies = [ + { + "toc_range": [1, 1], + "toc_with_level": [ + {"heading": "Only", "level": 1, "page_number": 2}, + ], + } + ] + ctx = _ctx() + + 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( + "app.services.document_agent.agents.calibration.service.calibrate_offset", + return_value=phase1, + ), + _patch_verify(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, + ) + # 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) + assert isinstance(anchor.null_page_report, list) + 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_verify(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_verify(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 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/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/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/ +

+
+ + +
+ + + + 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.