From c9de1d0cd73482fa7047e3275a3abe3855366af5 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 21 May 2026 21:32:41 +0800 Subject: [PATCH 01/11] feat: add DocumentAnatomyAgent and PDF shard splitter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce Phase-0 anatomy infrastructure for large-PDF splitting: - page_map.py: data contracts — PageMap, PageFeature, Shard, CutPoint, H1BoundaryResult, H1Match - agent.py: DocumentAnatomyAgent — LLM tool-calling loop (scan → find_h1 → propose_cuts → finalize) with deterministic fallback - tools/scan_all_page_features.py: full-page PyMuPDF structural feature extraction (text density, image coverage, table count, orientation) in an isolated child process - tools/find_h1_boundaries.py: TOC-page grep + body grep to locate level-1 heading physical pages; falls back to preview grep when no TOC is detected - shard_splitter.py: split_pdf_by_shards() and merge_shard_dataframes() with page_nums offset correction Integration into parse_pdfs() is the next step. --- .../app/services/document_agent/__init__.py | 27 +- .../app/services/document_agent/agent.py | 674 ++++++++++++++++++ .../app/services/document_agent/page_map.py | 152 ++++ .../services/document_agent/tools/__init__.py | 20 +- .../tools/find_h1_boundaries.py | 352 +++++++++ .../tools/scan_all_page_features.py | 244 +++++++ .../formats/pdf/shard_splitter.py | 208 ++++++ 7 files changed, 1675 insertions(+), 2 deletions(-) create mode 100644 apps/worker/app/services/document_agent/agent.py create mode 100644 apps/worker/app/services/document_agent/page_map.py create mode 100644 apps/worker/app/services/document_agent/tools/find_h1_boundaries.py create mode 100644 apps/worker/app/services/document_agent/tools/scan_all_page_features.py create mode 100644 apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py diff --git a/apps/worker/app/services/document_agent/__init__.py b/apps/worker/app/services/document_agent/__init__.py index 2e36f6f2e..9c8475d9a 100644 --- a/apps/worker/app/services/document_agent/__init__.py +++ b/apps/worker/app/services/document_agent/__init__.py @@ -1,13 +1,38 @@ -"""Phase 1 document split-agent utilities.""" +"""Document Anatomy Agent package. +Phase 0 (Anatomy): ``DocumentAnatomyAgent`` produces a ``PageMap`` before +any PDF parsing begins, enabling semantically-correct shard decisions. + +Phase 1 (Shard planning): ``ShardManifest`` / ``ShardSignal`` / … are the +original shard-planning primitives retained for backward compatibility. +""" + +from app.services.document_agent.agent import DocumentAnatomyAgent from app.services.document_agent.manifest import ( GlobalSignals, ShardManifest, ShardSignal, SpecialPage, ) +from app.services.document_agent.page_map import ( + CutPoint, + H1BoundaryResult, + H1Match, + PageFeature, + PageMap, + Shard, +) __all__ = [ + # Phase 0 — Anatomy Agent + "DocumentAnatomyAgent", + "CutPoint", + "H1BoundaryResult", + "H1Match", + "PageFeature", + "PageMap", + "Shard", + # Phase 1 — Shard planning primitives "GlobalSignals", "ShardManifest", "ShardSignal", diff --git a/apps/worker/app/services/document_agent/agent.py b/apps/worker/app/services/document_agent/agent.py new file mode 100644 index 000000000..abc0b3931 --- /dev/null +++ b/apps/worker/app/services/document_agent/agent.py @@ -0,0 +1,674 @@ +"""DocumentAnatomyAgent — produce a PageMap for a PDF via LLM tool-calling. + +Architecture +------------ +The agent holds a **minimal tool set** — only operations that genuinely require +structured data from the PDF are exposed as tools. Classification and +heuristic logic run deterministically inside the agent; only the final +cut-point decision is delegated to the LLM (which has full context by then). + +Two-phase tool calling: +1. ``scan_all_page_features`` — collects per-page structural signals. +2. ``find_h1_boundaries`` — locates level-1 headings via text search. + +The LLM then reasons over the collected evidence to call +``propose_cut_points``, producing the final shard plan. + +Design constraints +------------------ +- No hardcoded page counts, thresholds, or prompt examples. +- Every parameter that influences splitting comes from ``settings`` or is + passed explicitly by the caller. +- Graceful degradation at every step: tools return empty/safe results rather + than raising; the agent falls back to deterministic cuts if the LLM fails. +""" + +from __future__ import annotations + +import json +import os +import traceback +from datetime import datetime, timezone +from typing import Any + +from app.services.document_agent.page_map import ( + CutPoint, + H1BoundaryResult, + H1Match, + PageFeature, + PageMap, + Shard, +) +from app.services.document_agent.tools.classify_special_pages import ( + heuristic_classify_special_pages, +) +from app.services.document_agent.tools.find_h1_boundaries import find_h1_boundaries +from app.services.document_agent.tools.scan_all_page_features import ( + scan_all_page_features, +) +from loguru import logger + + +# ── Tool schemas ─────────────────────────────────────────────────────────────── + +_TOOL_SCHEMAS: list[dict[str, Any]] = [ + { + "type": "function", + "function": { + "name": "scan_all_page_features", + "description": ( + "Perform a full structural scan of every page in the PDF. " + "Returns per-page signals: text length, image coverage, table count, " + "orientation, blank-page flag, and a text preview. " + "Always call this first." + ), + "parameters": {"type": "object", "properties": {}, "required": []}, + }, + }, + { + "type": "function", + "function": { + "name": "find_h1_boundaries", + "description": ( + "Locate level-1 headings in the document body by grepping page " + "texts against TOC entries. Returns the page numbers where each " + "level-1 heading physically starts. Call after scan_all_page_features." + ), + "parameters": {"type": "object", "properties": {}, "required": []}, + }, + }, + { + "type": "function", + "function": { + "name": "propose_cut_points", + "description": ( + "Propose a list of shard boundaries for the document. " + "Each cut_point marks where one shard ends (the next starts at " + "cut_after_page + 1). " + "Constraints you must respect:\n" + "- Shards may not exceed max_pages_per_shard pages.\n" + "- Shards should be at least min_pages_per_shard pages.\n" + "- Do not cut through landscape blocks or table-heavy pages.\n" + "- Prefer h1_heading boundaries; fall back to blank/sparse pages; " + "use forced cuts only as a last resort.\n" + "- If the document is short enough that no split is needed, " + "return an empty cut_points list." + ), + "parameters": { + "type": "object", + "properties": { + "cut_points": { + "type": "array", + "items": { + "type": "object", + "properties": { + "cut_after_page": {"type": "integer"}, + "anchor_type": { + "type": "string", + "enum": [ + "h1_heading", + "blank", + "sparse", + "forced", + ], + }, + "rationale": {"type": "string"}, + "confidence": {"type": "number"}, + }, + "required": [ + "cut_after_page", + "anchor_type", + "rationale", + ], + }, + } + }, + "required": ["cut_points"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "finalize", + "description": "Signal that cut planning is complete. Call last.", + "parameters": {"type": "object", "properties": {}, "required": []}, + }, + }, +] + + +def _build_system_prompt( + split_threshold: int, + max_pages_per_shard: int, + min_pages_per_shard: int, +) -> str: + return ( + "You are a Document Anatomy Agent. Your task is to analyse a PDF and " + "decide how to split it into semantically coherent shards for downstream " + "hierarchical heading extraction.\n\n" + "Call tools in this order:\n" + "1. scan_all_page_features — always first.\n" + "2. find_h1_boundaries — always second.\n" + "3. propose_cut_points — reason carefully over the evidence:\n" + f" - Only propose cuts when the document exceeds {split_threshold} pages.\n" + f" - Each shard: between {min_pages_per_shard} and {max_pages_per_shard} pages.\n" + " - Prefer h1_heading anchors (level-1 heading start pages).\n" + " - Fall back to blank or sparse pages when headings are ambiguous.\n" + " - Use forced cuts only when no semantic boundary is available; " + "avoid cutting through landscape blocks or dense table regions.\n" + "4. finalize — always last.\n\n" + "Be concise. Do not repeat data already returned by tools." + ) + + +# ── Agent ────────────────────────────────────────────────────────────────────── + + +class DocumentAnatomyAgent: + """LLM-orchestrated agent producing a ``PageMap`` for a PDF. + + Parameters + ---------- + model_name: + Override the LLM model. Falls back to ``settings.HIERARCHY_LLM_MODEL`` + then ``settings.NORMOL_MODEL``. + split_threshold: + Minimum page count before physical splitting is considered. If + ``None``, read from ``settings.PDF_ANATOMY_SPLIT_THRESHOLD`` (default + behaviour); the caller can pass an explicit value for testing. + max_pages_per_shard: + Hard cap on shard size. Same settings fallback pattern. + min_pages_per_shard: + Prevent micro-shards that would cause overhead without benefit. + max_iterations: + Hard limit on the LLM tool-calling loop to prevent runaway calls. + """ + + def __init__( + self, + model_name: str | None = None, + split_threshold: int | None = None, + max_pages_per_shard: int | None = None, + min_pages_per_shard: int | None = None, + max_iterations: int = 10, + ) -> None: + self._model_name = model_name + self._split_threshold = split_threshold + self._max_pages_per_shard = max_pages_per_shard + self._min_pages_per_shard = min_pages_per_shard + self._max_iterations = max_iterations + + # Per-run state (reset by run()) + self._pdf_path: str = "" + self._page_features: list[PageFeature] = [] + self._page_labels: list[dict[str, Any]] = [] + self._h1_result: H1BoundaryResult | None = None + self._cut_points: list[CutPoint] = [] + self._decision_log: list[str] = [] + self._finalized: bool = False + + # ── Threshold resolution (defer to settings to avoid hardcoding) ─────────── + + def _resolve_thresholds(self) -> tuple[int, int, int]: + """Return (split_threshold, max_per_shard, min_per_shard) from settings.""" + try: + from shared.core.config import settings + + split = self._split_threshold or getattr( + settings, "PDF_ANATOMY_SPLIT_THRESHOLD", 200 + ) + max_s = self._max_pages_per_shard or getattr( + settings, "PDF_ANATOMY_MAX_PAGES_PER_SHARD", 200 + ) + min_s = self._min_pages_per_shard or getattr( + settings, "PDF_ANATOMY_MIN_PAGES_PER_SHARD", 20 + ) + except Exception: + split = self._split_threshold or 200 + max_s = self._max_pages_per_shard or 200 + min_s = self._min_pages_per_shard or 20 + return int(split), int(max_s), int(min_s) + + def _resolve_model(self) -> str: + try: + from shared.core.config import settings + + return ( + self._model_name + or getattr(settings, "HIERARCHY_LLM_MODEL", None) + or getattr(settings, "NORMOL_MODEL", None) + or "deepseek-chat" + ) + except Exception: + return self._model_name or "deepseek-chat" + + # ── Public entry point ───────────────────────────────────────────────────── + + def run(self, pdf_path: str, job_id: str) -> PageMap: + """Run the anatomy agent; always returns a valid ``PageMap``.""" + self._reset(pdf_path) + split_threshold, max_per_shard, min_per_shard = self._resolve_thresholds() + logger.info( + f"[DocumentAnatomyAgent] start '{os.path.basename(pdf_path)}' " + f"split_threshold={split_threshold} max_shard={max_per_shard}" + ) + try: + page_map = self._run_loop( + job_id, split_threshold, max_per_shard, min_per_shard + ) + except Exception as exc: + logger.error( + f"[DocumentAnatomyAgent] unrecoverable error: {exc}\n" + + traceback.format_exc() + ) + page_map = self._fallback_page_map(job_id, reason=str(exc)) + + logger.info( + f"[DocumentAnatomyAgent] done: {page_map.page_count} pages, " + f"needs_split={page_map.needs_split}, shards={len(page_map.shards)}" + ) + return page_map + + # ── Internal state ───────────────────────────────────────────────────────── + + def _reset(self, pdf_path: str) -> None: + self._pdf_path = pdf_path + self._page_features = [] + self._page_labels = [] + self._h1_result = None + self._cut_points = [] + self._decision_log = [] + self._finalized = False + + def _log(self, msg: str) -> None: + logger.info(f"[DocumentAnatomyAgent] {msg}") + self._decision_log.append(msg) + + # ── Tool dispatch ────────────────────────────────────────────────────────── + + def _tool_scan_all_page_features(self, _args: dict) -> dict[str, Any]: + self._log("→ scan_all_page_features") + features = scan_all_page_features(self._pdf_path) + self._page_features = features + + # Run heuristic classification immediately (deterministic, no LLM cost) + probe_fmt = [ + { + "page_number": f.page, + "text_length": f.text_length, + "image_coverage": f.image_coverage, + "table_count": f.table_count, + "drawings_count": f.drawings_count, + "orientation": f.orientation, + "is_blank_like": f.is_blank_like, + "text_preview": f.text_preview, + } + for f in features + ] + labels = heuristic_classify_special_pages(probe_fmt) + self._page_labels = labels.get("pages") or [] + + # Return a compact summary for the LLM context (not the full feature list) + counts: dict[str, int] = {} + for lbl in self._page_labels: + kind = str(lbl.get("special_kind") or "normal") + counts[kind] = counts.get(kind, 0) + 1 + self._log( + f" {len(features)} pages | " + + ", ".join(f"{k}={v}" for k, v in sorted(counts.items())) + ) + return { + "page_count": len(features), + "page_type_counts": counts, + # Pass sparse structural summary so LLM can reason without wall of text + "page_summary": [ + { + "page": f.page, + "kind": next( + ( + lbl.get("special_kind", "normal") + for lbl in self._page_labels + if lbl.get("page") == f.page + ), + "normal", + ), + "orientation": f.orientation, + "is_blank": f.is_blank_like, + } + for f in features + ], + } + + def _tool_find_h1_boundaries(self, _args: dict) -> dict[str, Any]: + self._log("→ find_h1_boundaries") + result = find_h1_boundaries(self._pdf_path, self._page_features) + self._h1_result = result + self._log( + f" method={result.method}, " + f"toc_pages={result.toc_pages}, " + f"h1_matches={len(result.h1_matches)}" + ) + return result.to_dict() + + def _tool_propose_cut_points(self, args: dict) -> dict[str, Any]: + self._log("→ propose_cut_points") + raw_cuts = args.get("cut_points") or [] + parsed: list[CutPoint] = [] + for raw in raw_cuts: + try: + parsed.append( + CutPoint( + cut_after_page=int(raw["cut_after_page"]), + anchor_type=raw.get("anchor_type", "forced"), + rationale=str(raw.get("rationale", ""))[:400], + confidence=float(raw.get("confidence", 1.0)), + ) + ) + except (KeyError, TypeError, ValueError): + continue + self._cut_points = sorted(parsed, key=lambda c: c.cut_after_page) + self._log(f" {len(self._cut_points)} cuts: {[c.cut_after_page for c in self._cut_points]}") + return {"accepted": len(self._cut_points)} + + def _tool_finalize(self, _args: dict) -> dict[str, Any]: + self._log("→ finalize") + self._finalized = True + return {"status": "ok"} + + _DISPATCH: dict[str, Any] = { + "scan_all_page_features": _tool_scan_all_page_features, + "find_h1_boundaries": _tool_find_h1_boundaries, + "propose_cut_points": _tool_propose_cut_points, + "finalize": _tool_finalize, + } + + def _dispatch(self, name: str, args: dict) -> dict[str, Any]: + handler = self._DISPATCH.get(name) + if handler is None: + return {"error": f"unknown tool: {name}"} + try: + return handler(self, args) + except Exception as exc: + logger.warning(f"[DocumentAnatomyAgent] tool '{name}' error: {exc}") + return {"error": str(exc)} + + # ── LLM loop ─────────────────────────────────────────────────────────────── + + def _run_loop( + self, + job_id: str, + split_threshold: int, + max_per_shard: int, + min_per_shard: int, + ) -> PageMap: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + model = self._resolve_model() + client = get_openai_client(model=model) + system_prompt = _build_system_prompt(split_threshold, max_per_shard, min_per_shard) + + messages: list[dict[str, Any]] = [ + {"role": "system", "content": system_prompt}, + { + "role": "user", + "content": ( + f"Analyse and plan shards for: {os.path.basename(self._pdf_path)}" + ), + }, + ] + + for iteration in range(self._max_iterations): + self._log(f"iteration {iteration + 1}/{self._max_iterations}") + try: + response = client.chat_completion( + messages=messages, + model=model, + temperature=0.0, + max_tokens=2000, + tools=_TOOL_SCHEMAS, + tool_choice="auto", + ) + except Exception as exc: + logger.error(f"[DocumentAnatomyAgent] LLM call failed: {exc}") + break + + tool_calls = _parse_tool_calls(response) + if not tool_calls: + self._log("no tool calls — exiting loop") + break + + assistant_msg: dict[str, Any] = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": tc["id"], + "type": "function", + "function": { + "name": tc["name"], + "arguments": json.dumps(tc["args"], ensure_ascii=False), + }, + } + for tc in tool_calls + ], + } + messages.append(assistant_msg) + + for tc in tool_calls: + result = self._dispatch(tc["name"], tc["args"]) + messages.append( + { + "role": "tool", + "tool_call_id": tc["id"], + "content": json.dumps(result, ensure_ascii=False, default=str), + } + ) + + if self._finalized: + self._log("finalized — exiting loop") + break + else: + self._log(f"reached max_iterations={self._max_iterations}") + + # Ensure data collection happened even if LLM skipped a step + if not self._page_features: + self._log("fallback: direct scan_all_page_features") + self._tool_scan_all_page_features({}) + + if self._h1_result is None: + self._log("fallback: direct find_h1_boundaries") + self._tool_find_h1_boundaries({}) + + page_count = len(self._page_features) + + # Deterministic fallback cuts if LLM proposed nothing and doc is long + if not self._cut_points and page_count > split_threshold: + self._log("no cuts proposed — applying deterministic fallback") + self._cut_points = _deterministic_cuts( + page_count=page_count, + page_labels=self._page_labels, + h1_result=self._h1_result, + max_per_shard=max_per_shard, + min_per_shard=min_per_shard, + ) + + return self._assemble_page_map(job_id, page_count) + + # ── PageMap assembly ─────────────────────────────────────────────────────── + + def _assemble_page_map(self, job_id: str, page_count: int) -> PageMap: + shards = _cuts_to_shards(self._cut_points, page_count) + h1_result = self._h1_result or H1BoundaryResult( + toc_pages=[], h1_matches=[], method="none" + ) + return PageMap( + job_id=job_id, + file_path=self._pdf_path, + page_count=page_count, + h1_result=h1_result, + page_features=self._page_features, + shards=shards, + needs_split=len(shards) > 1, + global_signals=_global_signals(self._page_features, self._page_labels, h1_result), + agent_decision_log=list(self._decision_log), + created_at=datetime.now(timezone.utc), + ) + + def _fallback_page_map(self, job_id: str, reason: str) -> PageMap: + page_count = len(self._page_features) + shards = [Shard(page_start=1, page_end=max(page_count, 1), page_offset=0)] + return PageMap( + job_id=job_id, + file_path=self._pdf_path, + page_count=page_count, + h1_result=H1BoundaryResult(toc_pages=[], h1_matches=[], method="none"), + page_features=self._page_features, + shards=shards, + needs_split=False, + global_signals={}, + agent_decision_log=self._decision_log + [f"FALLBACK: {reason}"], + created_at=datetime.now(timezone.utc), + ) + + +# ── Module helpers ───────────────────────────────────────────────────────────── + + +def _parse_tool_calls(response: Any) -> list[dict[str, Any]]: + """Normalise an LLM response into [{id, name, args}] dicts.""" + if isinstance(response, str): + return [] + choices = getattr(response, "choices", None) + if not choices: + return [] + tc_list = getattr(choices[0].message, "tool_calls", None) or [] + result = [] + for tc in tc_list: + fn = getattr(tc, "function", None) + if not fn: + continue + try: + args = json.loads(getattr(fn, "arguments", "{}") or "{}") + except (json.JSONDecodeError, TypeError): + args = {} + result.append({"id": getattr(tc, "id", ""), "name": getattr(fn, "name", ""), "args": args}) + return result + + +def _deterministic_cuts( + page_count: int, + page_labels: list[dict[str, Any]], + h1_result: H1BoundaryResult | None, + max_per_shard: int, + min_per_shard: int, +) -> list[CutPoint]: + """Produce cuts without LLM: h1 pages → blank/sparse pages → forced.""" + by_page: dict[int, str] = { + int(p.get("page") or 0): str(p.get("special_kind") or "normal") + for p in page_labels + if p.get("page") + } + avoid = {"table_heavy", "landscape"} + cuts: list[CutPoint] = [] + + # Option A: h1 heading boundaries + h1_pages = sorted({m.page for m in (h1_result.h1_matches if h1_result else [])}) + if h1_pages: + prev = 0 + for h1_page in h1_pages: + cut_page = h1_page - 1 + if cut_page <= 0 or cut_page <= prev: + continue + shard_len = cut_page - prev + if shard_len < min_per_shard: + continue + if shard_len > max_per_shard: + # Need an intermediate forced cut + forced = prev + max_per_shard + while forced < cut_page: + actual = forced + for offset in range(5): + cand = forced - offset + if cand > prev and by_page.get(cand, "normal") not in avoid: + actual = cand + break + cuts.append(CutPoint(cut_after_page=actual, anchor_type="forced", + rationale="intermediate forced cut before h1", confidence=0.5)) + prev = actual + forced = prev + max_per_shard + cuts.append(CutPoint(cut_after_page=cut_page, anchor_type="h1_heading", + rationale=f"h1 heading starts at page {h1_page}", confidence=0.9)) + prev = cut_page + if cuts: + return cuts + + # Option B: blank/sparse pages near shard boundaries + prev = 0 + while prev + max_per_shard < page_count: + target = prev + max_per_shard + actual = target + for offset in range(min(10, max_per_shard // 2)): + for cand in [target - offset, target + offset]: + if prev < cand <= page_count: + kind = by_page.get(cand, "normal") + if kind in {"blank", "sparse"} and cand - prev >= min_per_shard: + actual = cand + break + else: + continue + break + anchor = "blank" if by_page.get(actual, "normal") == "blank" else ( + "sparse" if by_page.get(actual, "normal") == "sparse" else "forced" + ) + cuts.append(CutPoint(cut_after_page=actual, anchor_type=anchor, + rationale=f"deterministic cut near shard boundary", confidence=0.6)) + prev = actual + + return cuts + + +def _cuts_to_shards(cuts: list[CutPoint], page_count: int) -> list[Shard]: + if page_count <= 0: + return [] + if not cuts: + return [Shard(page_start=1, page_end=page_count, page_offset=0)] + shards: list[Shard] = [] + prev = 0 + for cut in sorted(cuts, key=lambda c: c.cut_after_page): + cap = min(int(cut.cut_after_page), page_count) + if cap <= prev: + continue + shards.append(Shard(page_start=prev + 1, page_end=cap, page_offset=prev)) + prev = cap + if prev < page_count: + shards.append(Shard(page_start=prev + 1, page_end=page_count, page_offset=prev)) + return shards + + +def _global_signals( + features: list[PageFeature], + labels: list[dict[str, Any]], + h1_result: H1BoundaryResult, +) -> dict[str, Any]: + total = len(features) + if not total: + return {} + label_counts: dict[str, int] = {} + for lbl in labels: + k = str(lbl.get("special_kind") or "normal") + label_counts[k] = label_counts.get(k, 0) + 1 + return { + "total_pages": total, + "h1_method": h1_result.method, + "toc_pages": h1_result.toc_pages, + "h1_match_count": len(h1_result.h1_matches), + "h1_cut_candidates": h1_result.cut_candidate_pages(), + "landscape_ratio": round( + sum(1 for f in features if f.orientation == "landscape") / total, 3 + ), + "blank_ratio": round( + sum(1 for f in features if f.is_blank_like) / total, 3 + ), + "page_type_counts": label_counts, + } diff --git a/apps/worker/app/services/document_agent/page_map.py b/apps/worker/app/services/document_agent/page_map.py new file mode 100644 index 000000000..cff00cb32 --- /dev/null +++ b/apps/worker/app/services/document_agent/page_map.py @@ -0,0 +1,152 @@ +"""PageMap — the output contract of the Document Anatomy Agent. + +Keeps anatomy-specific models separate from the shard-planning primitives in +``manifest.py`` so they can evolve independently. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from typing import Any, Literal + + +# ── Page-level feature snapshot ─────────────────────────────────────────────── + + +@dataclass +class PageFeature: + """Raw structural measurements for a single PDF page (1-based numbering).""" + + page: int + text_length: int + text_density: float # chars per 10k pt² of page area + image_coverage: float # fraction [0, 1] + image_count: int + table_count: int + drawings_count: int + orientation: Literal["portrait", "landscape"] + width: float # points + height: float # points + is_blank_like: bool + text_preview: str # first N chars of page text (configurable) + + # Reserved for future Page Memory integration + embedding_ref: str | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +# ── H1 boundary detection result ───────────────────────────────────────────── + + +@dataclass +class H1Match: + """A single level-1 heading found via text search.""" + + title: str # normalized heading text from TOC + page: int # 1-based page where heading was found in the body + confidence: float # 1.0 = exact match, < 1.0 = fuzzy + match_text: str # the actual text snippet that matched on the page + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class H1BoundaryResult: + """Output of the ``find_h1_boundaries`` tool.""" + + toc_pages: list[int] # pages identified as TOC (by text content) + h1_matches: list[H1Match] # h1 headings and their body pages + method: Literal["toc_grep", "heading_grep", "none"] + notes: str = "" + + def cut_candidate_pages(self) -> list[int]: + """Pages just before each h1 heading page — natural cut points.""" + pages = sorted({m.page for m in self.h1_matches if m.page > 1}) + # Cut before the chapter starts (i.e., end the previous shard at page-1) + return [p - 1 for p in pages if p > 1] + + def to_dict(self) -> dict[str, Any]: + data = asdict(self) + data["h1_matches"] = [m.to_dict() for m in self.h1_matches] + return data + + +# ── Cut-point ───────────────────────────────────────────────────────────────── + + +@dataclass +class CutPoint: + """A proposed shard boundary produced by the agent.""" + + cut_after_page: int + rationale: str + anchor_type: Literal["h1_heading", "blank", "sparse", "forced"] + confidence: float = 1.0 + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +# ── Shard ────────────────────────────────────────────────────────────────────── + + +@dataclass +class Shard: + """Lightweight shard descriptor. ``page_offset`` is used to correct + ``page_nums`` in sub-PDF DataFrames back to the original document's + page numbering.""" + + page_start: int # 1-based, inclusive + page_end: int # 1-based, inclusive + page_offset: int # = page_start - 1 + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +# ── PageMap ─────────────────────────────────────────────────────────────────── + + +@dataclass +class PageMap: + """Complete anatomy report produced by ``DocumentAnatomyAgent``. + + Downstream consumers: + - ``formats/pdf/parser.py`` — decides whether to physically split. + - ``structure/layout_parser.py`` (future) — page labels annotate headings. + - Page Memory (future) — ``page_features`` become the structural component. + """ + + job_id: str + file_path: str + page_count: int + h1_result: H1BoundaryResult + page_features: list[PageFeature] + shards: list[Shard] + needs_split: bool + global_signals: dict[str, Any] + agent_decision_log: list[str] = field(default_factory=list) + created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + version: str = "1.0" + + def page_feature_map(self) -> dict[int, PageFeature]: + return {pf.page: pf for pf in self.page_features} + + def to_dict(self) -> dict[str, Any]: + return { + "version": self.version, + "job_id": self.job_id, + "file_path": self.file_path, + "page_count": self.page_count, + "needs_split": self.needs_split, + "h1_result": self.h1_result.to_dict(), + "shards": [s.to_dict() for s in self.shards], + "global_signals": self.global_signals, + "agent_decision_log": self.agent_decision_log, + "page_features": [pf.to_dict() for pf in self.page_features], + "created_at": self.created_at.isoformat(), + } diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index 8f54e85fa..7f0d81d23 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -1,15 +1,33 @@ -"""Document split-agent Phase 1 tools.""" +"""Document Anatomy Agent tools. + +Minimal set — only tools with genuine value are exported: + +- ``scan_all_page_features``: full-page structural feature extraction (required). +- ``find_h1_boundaries``: locates level-1 headings via text search (required). +- ``classify_special_pages`` / ``heuristic_classify_special_pages``: page + classification used internally by the agent (not called by LLM as a tool). +- ``propose_shard_plan`` / ``sample_pages`` / ``vlm_inspect_pages``: retained + for the existing shard-planning path and optional VLM inspection. +""" from app.services.document_agent.tools.classify_special_pages import ( classify_special_pages, + heuristic_classify_special_pages, ) +from app.services.document_agent.tools.find_h1_boundaries import find_h1_boundaries from app.services.document_agent.tools.probe_sample_pages import sample_pages from app.services.document_agent.tools.probe_vlm_inspect import vlm_inspect_pages from app.services.document_agent.tools.propose_shard_plan import propose_shard_plan +from app.services.document_agent.tools.scan_all_page_features import ( + scan_all_page_features, +) __all__ = [ "classify_special_pages", + "find_h1_boundaries", + "heuristic_classify_special_pages", "propose_shard_plan", "sample_pages", + "scan_all_page_features", "vlm_inspect_pages", ] diff --git a/apps/worker/app/services/document_agent/tools/find_h1_boundaries.py b/apps/worker/app/services/document_agent/tools/find_h1_boundaries.py new file mode 100644 index 000000000..e730b621d --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/find_h1_boundaries.py @@ -0,0 +1,352 @@ +"""find_h1_boundaries — locate level-1 headings in a PDF via text search. + +Why not use PyMuPDF ``doc.get_toc()`` page numbers? +--------------------------------------------------- +PDF bookmark page references encode physical page offsets, while printed page +numbers in a TOC reflect logical numbering that often includes unnumbered front +matter. For Chinese documents and scanned PDFs the offsets frequently disagree +by several pages or are entirely absent. + +This tool instead: +1. Identifies TOC pages by detecting TOC-marker text in page features already + gathered by ``scan_all_page_features``. +2. Reads those TOC pages in full to extract level-1 heading candidate strings. +3. Searches every page's full text for those candidates (fuzzy, after + normalisation) to find where they actually appear in the body. + +The result is a set of ``H1Match`` records linking each heading title to the +page where it physically starts — a reliable basis for shard cut decisions. +""" + +from __future__ import annotations + +import gc +import re +import unicodedata +from typing import Any + +from app.services.document_agent.page_map import H1BoundaryResult, H1Match, PageFeature +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) +from loguru import logger + + +# ── Text normalisation helpers ───────────────────────────────────────────────── + +# Patterns that prefix chapter/section numbers in various languages +_LEADING_NUMBER_RE = re.compile( + r"""^ + (?: + 第\s*[零一二三四五六七八九十百千\d]+\s*[章节篇部分] # 第X章/节 + | [零一二三四五六七八九十百千]+\s*[、。，,] # 一、 + | [（(]\s*[零一二三四五六七八九十百千\d]+\s*[）)] # （一） + | \d+(?:\.\d+)*\.?\s* # 1. / 1.2 / 1.2.3. + | [IVXLCDM]+\.?\s* # Roman I. II. + | [A-Za-z]\.\s* # A. B. + | Chapter\s+\w+\s* # Chapter N + ) + """, + re.VERBOSE | re.IGNORECASE, +) + + +def _normalise_heading(text: str) -> str: + """Strip leading numbers/labels and normalise whitespace for matching.""" + text = unicodedata.normalize("NFKC", text) + text = re.sub(r"\s+", " ", text).strip() + # Strip leading numbering patterns + stripped = _LEADING_NUMBER_RE.sub("", text).strip() + # Keep the original if stripping removed everything (guard against over-stripping) + return stripped if stripped else text + + +def _fuzzy_contains(needle: str, haystack: str, min_len: int = 4) -> bool: + """Check whether ``needle`` (normalised heading) appears in ``haystack``. + + Matching strategy (in order of strictness): + 1. Exact substring after normalisation. + 2. Stripped-number variant of needle appears in normalised haystack. + """ + if not needle or len(needle) < min_len: + return False + norm_hay = re.sub(r"\s+", " ", unicodedata.normalize("NFKC", haystack)) + if needle in norm_hay: + return True + stripped = _normalise_heading(needle) + if stripped and len(stripped) >= min_len and stripped in norm_hay: + return True + return False + + +# ── TOC page detection ───────────────────────────────────────────────────────── + +_TOC_MARKERS = frozenset(["目录", "目次", "contents", "tableofcontents"]) + + +def _is_toc_page(feature: PageFeature) -> bool: + """Heuristic: is this page a Table of Contents page?""" + text = re.sub(r"\s+", "", feature.text_preview.lower()) + return any(marker in text for marker in _TOC_MARKERS) + + +# ── Child-process worker: read full page texts ───────────────────────────────── + + +@worker +def _read_page_texts_worker( + queue, + pdf_path: str, + page_indices: list[int], # 0-based +) -> None: + """Read full text for the requested pages; runs in an isolated process.""" + import pymupdf # type: ignore[import] + + results: dict[int, str] = {} + try: + doc = pymupdf.open(pdf_path) + for idx in page_indices: + if 0 <= idx < doc.page_count: + try: + results[idx] = doc[idx].get_text() or "" + except Exception: + results[idx] = "" + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + + queue.put({"ok": True, "texts": results}) + + +def _load_full_page_texts( + pdf_path: str, + page_indices: list[int], + timeout: int = 120, +) -> dict[int, str]: + """Return {0-based-index: full_text} for the requested pages.""" + if not page_indices: + return {} + try: + result = run_in_child_process( + _read_page_texts_worker, pdf_path, page_indices, timeout=timeout + ) + return {int(k): str(v) for k, v in (result.get("texts") or {}).items()} + except Exception as exc: + logger.warning(f"[find_h1_boundaries] full-text load failed: {exc}") + return {} + + +# ── TOC text → h1 candidate extraction ──────────────────────────────────────── + +# Patterns that suggest a TOC line is a level-1 heading +# (numbered first-level or occupies a prominent position in the TOC) +_H1_TOC_LINE_RE = re.compile( + r""" + ^\s* + (?: + 第\s*[零一二三四五六七八九十百千\d]+\s*[章篇部] # 第X章 + | [零一二三四五六七八九十百千]+\s*[、。] # 一、 + | \d+\s*[\.\s] # 1. / 1 (level-1 only) + | Chapter\s+\w+ # Chapter … + | [IVXLCDM]+\.?\s+\w # Roman numeral + ) + """, + re.VERBOSE | re.IGNORECASE, +) + +# A line is clearly a sub-heading if it starts with at least two level numbers +_SUBHEADING_RE = re.compile(r"^\s*\d+\.\d+", re.IGNORECASE) + + +def _extract_h1_candidates_from_toc_text(toc_text: str) -> list[str]: + """Parse raw TOC page text and return likely level-1 heading strings.""" + candidates: list[str] = [] + seen: set[str] = set() + + for raw_line in toc_text.splitlines(): + line = raw_line.strip() + if not line: + continue + # Skip sub-headings (e.g. "1.2 something") + if _SUBHEADING_RE.match(line): + continue + # Skip lines that look like page numbers only + if re.fullmatch(r"[\d\s\.\-·…]+", line): + continue + # Must match a level-1 pattern + if not _H1_TOC_LINE_RE.match(line): + continue + # Strip trailing page-number suffix common in TOCs: " ……… 12" + cleaned = re.sub(r"[\s\.\-·…]+\d+\s*$", "", line).strip() + if not cleaned: + continue + norm = _normalise_heading(cleaned) + if norm and norm not in seen and len(norm) >= 2: + candidates.append(cleaned) # keep original for matching fidelity + seen.add(norm) + + return candidates + + +# ── Public API ───────────────────────────────────────────────────────────────── + + +def find_h1_boundaries( + pdf_path: str, + page_features: list[PageFeature], + *, + timeout: int = 180, +) -> H1BoundaryResult: + """Locate level-1 headings in a PDF by grepping full page texts. + + Algorithm + --------- + 1. Identify TOC pages from already-collected ``page_features``. + 2. Read those TOC pages in full; extract level-1 heading candidates. + 3. If no TOC found, fall back to grepping markdown-style headings + (``# Title`` or lines that match level-1 numbering patterns) from + each page's text_preview across all pages. + 4. For each candidate heading, search all pages' full text for the + heading text (after normalisation). Record the first matching page. + + Returns + ------- + ``H1BoundaryResult`` — always non-raising; ``method="none"`` when + nothing useful was found. + """ + if not page_features: + return H1BoundaryResult(toc_pages=[], h1_matches=[], method="none", + notes="no page features provided") + + # Step 1: find TOC pages + toc_pages = [pf.page for pf in page_features if _is_toc_page(pf)] + logger.info(f"[find_h1_boundaries] TOC pages (heuristic): {toc_pages}") + + # Step 2: extract h1 candidates from TOC page full text + h1_candidates: list[str] = [] + if toc_pages: + toc_indices = [p - 1 for p in toc_pages] # 0-based + toc_texts = _load_full_page_texts(pdf_path, toc_indices, timeout=timeout) + for idx, text in sorted(toc_texts.items()): + page_candidates = _extract_h1_candidates_from_toc_text(text) + logger.info( + f"[find_h1_boundaries] TOC page {idx + 1}: " + f"{len(page_candidates)} h1 candidates" + ) + h1_candidates.extend(page_candidates) + + # Step 3: fallback — grep heading-like lines from all page previews + using_fallback = not h1_candidates + if using_fallback: + logger.info( + "[find_h1_boundaries] no TOC or no candidates — " + "grepping all page text_previews for heading-like lines" + ) + for pf in page_features: + for line in pf.text_preview.splitlines(): + line = line.strip() + if _H1_TOC_LINE_RE.match(line) and not _SUBHEADING_RE.match(line): + cleaned = re.sub(r"[\s\.\-·…]+\d+\s*$", "", line).strip() + if cleaned and len(_normalise_heading(cleaned)) >= 2: + h1_candidates.append(cleaned) + + # Deduplicate candidates preserving order + seen_norms: set[str] = set() + unique_candidates: list[str] = [] + for c in h1_candidates: + norm = _normalise_heading(c) + if norm and norm not in seen_norms: + unique_candidates.append(c) + seen_norms.add(norm) + h1_candidates = unique_candidates + + if not h1_candidates: + return H1BoundaryResult( + toc_pages=toc_pages, + h1_matches=[], + method="none", + notes="no h1 candidates extracted", + ) + + logger.info( + f"[find_h1_boundaries] {len(h1_candidates)} unique h1 candidates to search" + ) + + # Step 4: search all pages for each candidate + # Load full text for all pages (excluding confirmed TOC pages to avoid + # matching the TOC entry itself instead of the body heading) + non_toc_indices = [ + pf.page - 1 for pf in page_features if pf.page not in set(toc_pages) + ] + all_texts = _load_full_page_texts(pdf_path, non_toc_indices, timeout=timeout) + + # Build a sorted list of (page_number, full_text) for ordered search + page_text_pairs: list[tuple[int, str]] = sorted( + ((idx + 1, text) for idx, text in all_texts.items()), + key=lambda x: x[0], + ) + + h1_matches: list[H1Match] = [] + for candidate in h1_candidates: + norm_candidate = _normalise_heading(candidate) + matched_page: int | None = None + match_text: str = "" + confidence: float = 0.0 + + for page_num, full_text in page_text_pairs: + # Try exact match first (higher confidence) + if candidate in full_text: + matched_page = page_num + match_text = candidate + confidence = 1.0 + break + # Try normalised match + if norm_candidate and _fuzzy_contains(norm_candidate, full_text): + matched_page = page_num + match_text = norm_candidate + confidence = 0.85 + break + + if matched_page is not None: + h1_matches.append( + H1Match( + title=candidate, + page=matched_page, + confidence=confidence, + match_text=match_text, + ) + ) + logger.debug( + f"[find_h1_boundaries] '{candidate[:40]}' → page {matched_page} " + f"(conf={confidence})" + ) + else: + logger.debug( + f"[find_h1_boundaries] '{candidate[:40]}' → not found in body" + ) + + method: str + if h1_matches and not using_fallback: + method = "toc_grep" + elif h1_matches: + method = "heading_grep" + else: + method = "none" + + logger.info( + f"[find_h1_boundaries] result: method={method}, " + f"{len(h1_matches)}/{len(h1_candidates)} headings matched" + ) + return H1BoundaryResult( + toc_pages=toc_pages, + h1_matches=h1_matches, + method=method, + notes=( + f"{len(h1_matches)} of {len(h1_candidates)} candidates matched; " + f"fallback={using_fallback}" + ), + ) diff --git a/apps/worker/app/services/document_agent/tools/scan_all_page_features.py b/apps/worker/app/services/document_agent/tools/scan_all_page_features.py new file mode 100644 index 000000000..a5a97846c --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/scan_all_page_features.py @@ -0,0 +1,244 @@ +"""scan_all_page_features — Tool #2 for DocumentAnatomyAgent. + +Performs a **full-page traversal** (not sampling) of a PDF, extracting +structural features for every page. Runs inside an isolated PyMuPDF child +process to ensure memory is freed after extraction. + +Contrast with ``probe_sample_pages.py`` (which samples up to 25 pages for a +quick diagnostic): this tool is used when the agent needs precise per-page +labels to find semantic cut points. + +For a 200-page A4 PDF the child process typically completes in < 8 s. +""" + +from __future__ import annotations + +import gc +import statistics +from typing import Any + +from app.services.document_agent.page_map import PageFeature +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) +from loguru import logger + + +# ── Low-level helpers (same logic as probe_sample_pages, kept local) ────────── + + +def _rect_area(rect: Any) -> float: + w = max(float(getattr(rect, "width", 0.0) or 0.0), 0.0) + h = max(float(getattr(rect, "height", 0.0) or 0.0), 0.0) + return w * h + + +def _measure_image_coverage(page: Any, page_area: float) -> tuple[float, int]: + if page_area <= 0: + return 0.0, 0 + image_area = 0.0 + images = page.get_images(full=True) or [] + seen_rects: set[tuple] = set() + for image in images: + if not image: + continue + xref = image[0] + try: + rects = page.get_image_rects(xref) or [] + except Exception: + rects = [] + for rect in rects: + key = ( + round(float(getattr(rect, "x0", 0.0) or 0.0), 2), + round(float(getattr(rect, "y0", 0.0) or 0.0), 2), + round(float(getattr(rect, "x1", 0.0) or 0.0), 2), + round(float(getattr(rect, "y1", 0.0) or 0.0), 2), + ) + if key in seen_rects: + continue + seen_rects.add(key) + image_area += _rect_area(rect) + return min(image_area / page_area, 1.0), len(images) + + +def _table_count(page: Any) -> int: + try: + finder = page.find_tables() + return len(getattr(finder, "tables", []) or []) + except Exception: + return 0 + + +def _extract_single_page_feature(page: Any, page_index: int) -> dict[str, Any]: + """Extract features for one page; returns a plain dict for IPC.""" + rect = page.rect + page_area = max(_rect_area(rect), 1.0) + text = page.get_text() or "" + text_len = len(text.strip()) + image_coverage, image_count = _measure_image_coverage(page, page_area) + try: + drawings_count = len(page.get_drawings() or []) + except Exception: + drawings_count = 0 + orientation = "landscape" if float(rect.width) > float(rect.height) else "portrait" + text_density = text_len / page_area * 10000 + table_count = _table_count(page) + is_blank_like = text_len < 20 and image_coverage < 0.02 and drawings_count < 5 + + return { + "page": page_index + 1, # 1-based + "text_length": text_len, + "text_density": round(text_density, 4), + "image_coverage": round(image_coverage, 4), + "image_count": image_count, + "table_count": table_count, + "drawings_count": drawings_count, + "orientation": orientation, + "width": round(float(rect.width), 2), + "height": round(float(rect.height), 2), + "is_blank_like": is_blank_like, + "text_preview": " ".join(text.split())[:300], + } + + +# ── Child-process worker ────────────────────────────────────────────────────── + + +@worker +def _scan_all_pages_worker( + queue, + pdf_path: str, + page_start: int, # 0-based inclusive + page_end: int, # 0-based inclusive (-1 = all) +) -> None: + """Scan all pages (or a subrange) and put feature list onto the queue.""" + import pymupdf # type: ignore[import] + + features: list[dict[str, Any]] = [] + page_count = 0 + + try: + doc = pymupdf.open(pdf_path) + page_count = int(doc.page_count) + end = page_count - 1 if page_end < 0 else min(page_end, page_count - 1) + start = max(0, page_start) + + for idx in range(start, end + 1): + try: + feat = _extract_single_page_feature(doc[idx], idx) + features.append(feat) + except Exception as exc: + # Log and continue; partial feature is better than crash + features.append( + { + "page": idx + 1, + "text_length": 0, + "text_density": 0.0, + "image_coverage": 0.0, + "image_count": 0, + "table_count": 0, + "drawings_count": 0, + "orientation": "portrait", + "width": 0.0, + "height": 0.0, + "is_blank_like": True, + "text_preview": "", + "_error": str(exc), + } + ) + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + + queue.put({"ok": True, "page_count": page_count, "features": features}) + + +# ── Public API ──────────────────────────────────────────────────────────────── + + +def scan_all_page_features( + pdf_path: str, + *, + page_start: int = 0, + page_end: int = -1, + timeout: int = 300, +) -> list[PageFeature]: + """Return a ``PageFeature`` for every page in the PDF (or a subrange). + + Args: + pdf_path: Absolute path to the PDF file. + page_start: 0-based index of first page to scan (default 0 = first). + page_end: 0-based index of last page to scan, inclusive + (default -1 = last page). + timeout: Child process timeout in seconds. Allow 300 s for large docs. + + Returns: + Ordered list of ``PageFeature``, one per page, sorted by page number. + Never raises; returns empty list on subprocess failure. + """ + try: + result = run_in_child_process( + _scan_all_pages_worker, + pdf_path, + page_start, + page_end, + timeout=timeout, + ) + except Exception as exc: + logger.error(f"[scan_all_page_features] subprocess failed: {exc}") + return [] + + raw_features: list[dict] = result.get("features") or [] + page_features: list[PageFeature] = [] + + # Compute median page dimensions for landscape anomaly detection + widths = [f["width"] for f in raw_features if f.get("width", 0) > 0] + heights = [f["height"] for f in raw_features if f.get("height", 0) > 0] + median_w = statistics.median(widths) if widths else 595.0 # A4 width in pt + median_h = statistics.median(heights) if heights else 842.0 + + for raw in raw_features: + page_num = int(raw.get("page") or 0) + if page_num <= 0: + continue + + # Refine orientation: a "portrait" page that is much wider than the + # document median is flagged as a layout anomaly (landscape-rotated page + # that PyMuPDF may not detect via width > height alone). + orientation = raw.get("orientation", "portrait") + w, h = float(raw.get("width") or 0), float(raw.get("height") or 0) + if orientation == "portrait" and median_h > 0 and w > 0: + # If this page's aspect ratio deviates by > 40% vs median, flag it + doc_ratio = median_w / median_h + page_ratio = w / h if h > 0 else 1.0 + if page_ratio > doc_ratio * 1.4: + orientation = "landscape" + + page_features.append( + PageFeature( + page=page_num, + text_length=int(raw.get("text_length") or 0), + text_density=float(raw.get("text_density") or 0.0), + image_coverage=float(raw.get("image_coverage") or 0.0), + image_count=int(raw.get("image_count") or 0), + table_count=int(raw.get("table_count") or 0), + drawings_count=int(raw.get("drawings_count") or 0), + orientation=orientation, + width=w, + height=float(raw.get("height") or 0.0), + is_blank_like=bool(raw.get("is_blank_like")), + text_preview=str(raw.get("text_preview") or "")[:300], + embedding_ref=None, + ) + ) + + page_features.sort(key=lambda pf: pf.page) + logger.info( + f"[scan_all_page_features] scanned {len(page_features)} pages " + f"from '{pdf_path}'" + ) + return page_features diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py new file mode 100644 index 000000000..7b0909dc5 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py @@ -0,0 +1,208 @@ +"""shard_splitter — PDF physical splitting utilities. + +Provides two public functions: + +- ``split_pdf_by_shards()``: write each ``Shard`` as a temporary sub-PDF file. +- ``merge_shard_dataframes()``: merge per-shard DataFrames and correct + ``page_nums`` offsets so they reflect the original document's page numbers. + +Design notes +------------ +- Splitting uses PyMuPDF ``doc.select()`` which preserves all page content + (images, fonts, annotations) but discards cross-shard hyperlinks — acceptable + for a parsing pipeline. +- Each sub-PDF is written to a temp directory under the job's ``output_dir`` + so it lives on the same filesystem as the rest of the job artifacts. +- ``page_offset`` from ``Shard`` is added to every ``page_nums`` value during + merge, keeping page references consistent with the source document. +""" + +from __future__ import annotations + +import os +import gc +from typing import Any + +import pandas as pd +from app.services.document_agent.page_map import Shard +from loguru import logger + + +# ── Split ────────────────────────────────────────────────────────────────────── + + +def split_pdf_by_shards( + pdf_path: str, + shards: list[Shard], + output_dir: str, +) -> list[dict[str, Any]]: + """Write each shard as an independent sub-PDF under ``output_dir``. + + Returns a list of dicts:: + + [ + { + "shard_index": 0, + "shard": Shard(...), + "sub_pdf_path": "/tmp/.../shard_0_p1-p50.pdf", + "sub_output_dir": "/tmp/.../shard_0/", + }, + … + ] + + Raises ``RuntimeError`` if the source PDF cannot be opened. + Never raises for individual shard write failures — they are logged and + skipped (the caller falls back to parsing the full PDF in that case). + """ + try: + import pymupdf # type: ignore[import] + except ImportError as exc: + raise RuntimeError("PyMuPDF (pymupdf) is required for PDF splitting") from exc + + shards_dir = os.path.join(output_dir, "_shards") + os.makedirs(shards_dir, exist_ok=True) + + try: + doc = pymupdf.open(pdf_path) + total_pages = int(doc.page_count) + except Exception as exc: + raise RuntimeError(f"Cannot open PDF '{pdf_path}': {exc}") from exc + + results: list[dict[str, Any]] = [] + + try: + for shard_idx, shard in enumerate(shards): + # Clamp to actual page count + p_start = max(1, int(shard.page_start)) + p_end = min(total_pages, int(shard.page_end)) + + if p_start > total_pages or p_start > p_end: + logger.warning( + f"[shard_splitter] shard {shard_idx} range [{p_start},{p_end}] " + f"is out of bounds (total={total_pages}) — skipping" + ) + continue + + # PyMuPDF uses 0-based indices + page_indices = list(range(p_start - 1, p_end)) + + sub_pdf_name = f"shard_{shard_idx}_p{p_start}-p{p_end}.pdf" + sub_pdf_path = os.path.join(shards_dir, sub_pdf_name) + sub_output_dir = os.path.join(shards_dir, f"shard_{shard_idx}") + os.makedirs(sub_output_dir, exist_ok=True) + + try: + sub_doc = pymupdf.open() # empty document + sub_doc.insert_pdf(doc, from_page=p_start - 1, to_page=p_end - 1) + sub_doc.save(sub_pdf_path, garbage=4, deflate=True) + sub_doc.close() + + logger.info( + f"[shard_splitter] shard {shard_idx}: pages {p_start}-{p_end} " + f"({len(page_indices)} pages) → {sub_pdf_path}" + ) + results.append( + { + "shard_index": shard_idx, + "shard": shard, + "sub_pdf_path": sub_pdf_path, + "sub_output_dir": sub_output_dir, + } + ) + except Exception as exc: + logger.error( + f"[shard_splitter] failed to write shard {shard_idx} " + f"(pages {p_start}-{p_end}): {exc}" + ) + finally: + doc.close() + gc.collect() + + return results + + +# ── Merge ────────────────────────────────────────────────────────────────────── + + +def merge_shard_dataframes( + shard_dfs: list[tuple[Shard, pd.DataFrame]], +) -> pd.DataFrame: + """Merge per-shard DataFrames into a single DataFrame. + + For each shard, adds ``shard.page_offset`` to every value in the + ``page_nums`` column so that page references reflect the original + document rather than the sub-PDF's local page numbers. + + Args: + shard_dfs: Ordered list of ``(Shard, DataFrame)`` pairs. + DataFrames must have the columns produced by the parser + (``content``, ``path``, ``type``, ``page_nums``, …). + + Returns: + A single DataFrame with all rows, sorted by ``page_nums`` (ascending). + Returns an empty DataFrame if ``shard_dfs`` is empty. + """ + if not shard_dfs: + return pd.DataFrame() + + adjusted: list[pd.DataFrame] = [] + + for shard, df in shard_dfs: + if df is None or df.empty: + logger.warning( + f"[shard_splitter] shard pages {shard.page_start}-{shard.page_end} " + f"produced an empty DataFrame — skipping" + ) + continue + + df_copy = df.copy() + + # Correct page_nums: add page_offset to each page number in the list + if "page_nums" in df_copy.columns: + offset = int(shard.page_offset) + + def _shift_page_nums(val: Any, off: int = offset) -> Any: + if isinstance(val, list): + return [p + off for p in val if isinstance(p, int)] + if isinstance(val, str): + # Stored as comma-separated string in some paths + try: + nums = [int(x.strip()) for x in val.split(",") if x.strip()] + return ",".join(str(p + off) for p in nums) + except ValueError: + return val + return val + + df_copy["page_nums"] = df_copy["page_nums"].apply(_shift_page_nums) + + adjusted.append(df_copy) + logger.info( + f"[shard_splitter] merged shard pages " + f"{shard.page_start}-{shard.page_end} ({len(df_copy)} rows, " + f"offset={shard.page_offset})" + ) + + if not adjusted: + return pd.DataFrame() + + merged = pd.concat(adjusted, ignore_index=True) + + # Sort by the first page number of each row's page_nums for document order + if "page_nums" in merged.columns: + def _first_page(val: Any) -> int: + if isinstance(val, list) and val: + return int(val[0]) + if isinstance(val, str): + try: + parts = [int(x.strip()) for x in val.split(",") if x.strip()] + return parts[0] if parts else 0 + except ValueError: + return 0 + return 0 + + merged["_sort_page"] = merged["page_nums"].apply(_first_page) + merged = merged.sort_values("_sort_page").drop(columns=["_sort_page"]) + merged = merged.reset_index(drop=True) + + logger.info(f"[shard_splitter] merge complete: {len(merged)} total rows") + return merged From 9f8a8807b3d022123957de77058d062f9b079ce8 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 21 May 2026 21:43:24 +0800 Subject: [PATCH 02/11] refactor: remove Phase 1 shard planning legacy Delete unused Phase 1 modules that were never integrated into the production parser pipeline: - manifest.py: ShardManifest / ShardSignal / GlobalSignals / SpecialPage - tools/propose_shard_plan.py: LLM + fallback shard planning - tools/probe_sample_pages.py: stratified page sampling - tools/probe_vlm_inspect.py: VLM page screenshot inspection Adapt surviving modules: - classify_special_pages.py: inline SpecialKind Literal type (was imported from deleted manifest.py) - __init__.py / tools/__init__.py: remove all Phase 1 exports - page_map.py / scan_all_page_features.py: clean stale docstring refs The Phase 0 DocumentAnatomyAgent (agent.py, page_map.py, scan_all_page_features, find_h1_boundaries, shard_splitter) is the sole path forward. No production impact: neither Phase 0 nor Phase 1 was ever called from parse_pdfs() or any other production entry point. --- .../app/services/document_agent/__init__.py | 19 +- .../app/services/document_agent/manifest.py | 137 ------- .../app/services/document_agent/page_map.py | 3 - .../services/document_agent/tools/__init__.py | 14 +- .../tools/classify_special_pages.py | 16 +- .../tools/probe_sample_pages.py | 213 ----------- .../document_agent/tools/probe_vlm_inspect.py | 161 -------- .../tools/propose_shard_plan.py | 361 ------------------ .../tools/scan_all_page_features.py | 8 +- 9 files changed, 19 insertions(+), 913 deletions(-) delete mode 100644 apps/worker/app/services/document_agent/manifest.py delete mode 100644 apps/worker/app/services/document_agent/tools/probe_sample_pages.py delete mode 100644 apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py delete mode 100644 apps/worker/app/services/document_agent/tools/propose_shard_plan.py diff --git a/apps/worker/app/services/document_agent/__init__.py b/apps/worker/app/services/document_agent/__init__.py index 9c8475d9a..96081fc35 100644 --- a/apps/worker/app/services/document_agent/__init__.py +++ b/apps/worker/app/services/document_agent/__init__.py @@ -1,19 +1,10 @@ """Document Anatomy Agent package. -Phase 0 (Anatomy): ``DocumentAnatomyAgent`` produces a ``PageMap`` before -any PDF parsing begins, enabling semantically-correct shard decisions. - -Phase 1 (Shard planning): ``ShardManifest`` / ``ShardSignal`` / … are the -original shard-planning primitives retained for backward compatibility. +``DocumentAnatomyAgent`` produces a ``PageMap`` before any PDF parsing begins, +enabling semantically-correct shard decisions for large documents. """ from app.services.document_agent.agent import DocumentAnatomyAgent -from app.services.document_agent.manifest import ( - GlobalSignals, - ShardManifest, - ShardSignal, - SpecialPage, -) from app.services.document_agent.page_map import ( CutPoint, H1BoundaryResult, @@ -24,7 +15,6 @@ ) __all__ = [ - # Phase 0 — Anatomy Agent "DocumentAnatomyAgent", "CutPoint", "H1BoundaryResult", @@ -32,9 +22,4 @@ "PageFeature", "PageMap", "Shard", - # Phase 1 — Shard planning primitives - "GlobalSignals", - "ShardManifest", - "ShardSignal", - "SpecialPage", ] diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py deleted file mode 100644 index 7ed3f9cba..000000000 --- a/apps/worker/app/services/document_agent/manifest.py +++ /dev/null @@ -1,137 +0,0 @@ -"""Stable Phase 1 shard manifest contract. - -The manifest is intentionally independent from existing parser internals so it -can be produced and inspected before parser integration work begins. -""" - -from __future__ import annotations - -from dataclasses import asdict, dataclass, field -from datetime import datetime, timezone -from typing import Any, Literal - - -SpecialKind = Literal[ - "toc", - "blank", - "sparse", - "table_heavy", - "image_heavy", - "landscape", - "single_image", - "normal", -] - -PredominantKind = Literal[ - "text_dense", - "table_heavy", - "image_heavy", - "mixed", - "landscape_block", - "toc", - "sparse", -] - - -@dataclass -class SpecialPage: - page: int - kind: SpecialKind - confidence: float - note: str = "" - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -@dataclass -class ShardSignal: - page_start: int - page_end: int - page_offset: int - predominant_kind: PredominantKind - special_pages: list[SpecialPage] = field(default_factory=list) - estimated_difficulty: str | None = None - parser_hint: str | None = None - cut_rationale: str = "" - - def to_dict(self) -> dict[str, Any]: - data = asdict(self) - data["special_pages"] = [page.to_dict() for page in self.special_pages] - return data - - -@dataclass -class GlobalSignals: - has_toc: bool - toc_pages: list[int] - landscape_ratio: float - table_page_ratio: float - image_page_ratio: float - sample_size: int - notes: str = "" - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -@dataclass -class ShardManifest: - job_id: str - file_uri: str - file_sha: str - page_count: int - shard_count: int - shards: list[ShardSignal] - global_signals: GlobalSignals - decision_log_ref: str - created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) - version: str = "1.0" - - def validate(self) -> None: - """Enforce the downstream coverage contract.""" - if self.page_count < 0: - raise ValueError("page_count must be non-negative") - if self.shard_count != len(self.shards): - raise ValueError("shard_count must match shards length") - if self.page_count == 0: - if self.shards: - raise ValueError("empty documents cannot contain shards") - return - - expected_start = 1 - for shard in self.shards: - if shard.page_start != expected_start: - raise ValueError( - f"non-contiguous shard coverage at page {expected_start}: " - f"got start={shard.page_start}" - ) - if shard.page_end < shard.page_start: - raise ValueError( - f"invalid shard range {shard.page_start}-{shard.page_end}" - ) - if shard.page_offset != shard.page_start - 1: - raise ValueError( - f"invalid page_offset for shard {shard.page_start}-{shard.page_end}" - ) - expected_start = shard.page_end + 1 - - if expected_start != self.page_count + 1: - raise ValueError( - f"shards must cover 1..{self.page_count}, stopped at {expected_start - 1}" - ) - - def to_dict(self) -> dict[str, Any]: - self.validate() - return { - "version": self.version, - "job_id": self.job_id, - "file_uri": self.file_uri, - "file_sha": self.file_sha, - "page_count": self.page_count, - "shard_count": self.shard_count, - "shards": [shard.to_dict() for shard in self.shards], - "global_signals": self.global_signals.to_dict(), - "decision_log_ref": self.decision_log_ref, - "created_at": self.created_at.isoformat(), - } diff --git a/apps/worker/app/services/document_agent/page_map.py b/apps/worker/app/services/document_agent/page_map.py index cff00cb32..fc59ac2ce 100644 --- a/apps/worker/app/services/document_agent/page_map.py +++ b/apps/worker/app/services/document_agent/page_map.py @@ -1,7 +1,4 @@ """PageMap — the output contract of the Document Anatomy Agent. - -Keeps anatomy-specific models separate from the shard-planning primitives in -``manifest.py`` so they can evolve independently. """ from __future__ import annotations diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index 7f0d81d23..5056d7d49 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -1,13 +1,9 @@ """Document Anatomy Agent tools. -Minimal set — only tools with genuine value are exported: - -- ``scan_all_page_features``: full-page structural feature extraction (required). -- ``find_h1_boundaries``: locates level-1 headings via text search (required). +- ``scan_all_page_features``: full-page structural feature extraction. +- ``find_h1_boundaries``: locates level-1 headings via text search. - ``classify_special_pages`` / ``heuristic_classify_special_pages``: page classification used internally by the agent (not called by LLM as a tool). -- ``propose_shard_plan`` / ``sample_pages`` / ``vlm_inspect_pages``: retained - for the existing shard-planning path and optional VLM inspection. """ from app.services.document_agent.tools.classify_special_pages import ( @@ -15,9 +11,6 @@ heuristic_classify_special_pages, ) from app.services.document_agent.tools.find_h1_boundaries import find_h1_boundaries -from app.services.document_agent.tools.probe_sample_pages import sample_pages -from app.services.document_agent.tools.probe_vlm_inspect import vlm_inspect_pages -from app.services.document_agent.tools.propose_shard_plan import propose_shard_plan from app.services.document_agent.tools.scan_all_page_features import ( scan_all_page_features, ) @@ -26,8 +19,5 @@ "classify_special_pages", "find_h1_boundaries", "heuristic_classify_special_pages", - "propose_shard_plan", - "sample_pages", "scan_all_page_features", - "vlm_inspect_pages", ] diff --git a/apps/worker/app/services/document_agent/tools/classify_special_pages.py b/apps/worker/app/services/document_agent/tools/classify_special_pages.py index 0b8efa6c8..83072ccaf 100644 --- a/apps/worker/app/services/document_agent/tools/classify_special_pages.py +++ b/apps/worker/app/services/document_agent/tools/classify_special_pages.py @@ -1,14 +1,24 @@ -"""Classify sampled pages into Phase 1 special page kinds.""" +"""Classify pages into special-page kinds (heuristic + optional LLM).""" from __future__ import annotations import json -from typing import Any +from typing import Any, Literal -from app.services.document_agent.manifest import SpecialKind from app.services.document_agent.tools.llm_json import extract_json_object from loguru import logger +SpecialKind = Literal[ + "toc", + "blank", + "sparse", + "table_heavy", + "image_heavy", + "landscape", + "single_image", + "normal", +] + ALLOWED_KINDS: set[str] = { "toc", "blank", diff --git a/apps/worker/app/services/document_agent/tools/probe_sample_pages.py b/apps/worker/app/services/document_agent/tools/probe_sample_pages.py deleted file mode 100644 index b5b7e1122..000000000 --- a/apps/worker/app/services/document_agent/tools/probe_sample_pages.py +++ /dev/null @@ -1,213 +0,0 @@ -"""Feature-page sampling for the Phase 1 split agent.""" - -from __future__ import annotations - -import gc -import statistics -from typing import Any, Literal - -from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker - -SampleStrategy = Literal["stratified", "uniform", "key_pages"] - - -def _choose_sample_indices( - page_count: int, - *, - strategy: SampleStrategy = "stratified", - max_samples: int = 25, -) -> list[int]: - if page_count <= 0: - return [] - if max_samples <= 0: - return [] - if page_count <= max_samples: - return list(range(page_count)) - - if strategy == "key_pages": - candidates = [0, 1, 2, 3, 4, page_count - 5, page_count - 4, page_count - 3, page_count - 2, page_count - 1] - return sorted({idx for idx in candidates if 0 <= idx < page_count})[:max_samples] - - if strategy == "uniform": - if max_samples == 1: - return [0] - return sorted( - { - round(i * (page_count - 1) / (max_samples - 1)) - for i in range(max_samples) - } - ) - - edge_each = min(5, max_samples // 3) - edge_indices = list(range(edge_each)) + list(range(page_count - edge_each, page_count)) - remaining = max_samples - len(set(edge_indices)) - middle_start = edge_each - middle_end = page_count - edge_each - 1 - middle_indices: list[int] = [] - if remaining > 0 and middle_start <= middle_end: - if remaining == 1: - middle_indices = [(middle_start + middle_end) // 2] - else: - middle_indices = [ - round(middle_start + i * (middle_end - middle_start) / (remaining - 1)) - for i in range(remaining) - ] - return sorted({idx for idx in edge_indices + middle_indices if 0 <= idx < page_count}) - - -def _rect_area(rect: Any) -> float: - return max(float(getattr(rect, "width", 0.0) or 0.0), 0.0) * max( - float(getattr(rect, "height", 0.0) or 0.0), - 0.0, - ) - - -def _measure_image_coverage(page: Any, page_area: float) -> tuple[float, int]: - if page_area <= 0: - return 0.0, 0 - image_area = 0.0 - images = page.get_images(full=True) or [] - seen_rects: set[tuple[float, float, float, float]] = set() - for image in images: - if not image: - continue - xref = image[0] - try: - rects = page.get_image_rects(xref) or [] - except Exception: - rects = [] - for rect in rects: - key = ( - round(float(getattr(rect, "x0", 0.0) or 0.0), 2), - round(float(getattr(rect, "y0", 0.0) or 0.0), 2), - round(float(getattr(rect, "x1", 0.0) or 0.0), 2), - round(float(getattr(rect, "y1", 0.0) or 0.0), 2), - ) - if key in seen_rects: - continue - seen_rects.add(key) - image_area += _rect_area(rect) - return min(image_area / page_area, 1.0), len(images) - - -def _font_stats(page: Any) -> dict[str, float | int]: - sizes: list[float] = [] - try: - text_dict = page.get_text("dict") or {} - except Exception: - text_dict = {} - for block in text_dict.get("blocks", []) or []: - for line in block.get("lines", []) or []: - for span in line.get("spans", []) or []: - size = float(span.get("size") or 0.0) - if size > 0: - sizes.append(size) - if not sizes: - return {"min": 0.0, "max": 0.0, "mean": 0.0, "median": 0.0, "count": 0} - return { - "min": min(sizes), - "max": max(sizes), - "mean": statistics.fmean(sizes), - "median": statistics.median(sizes), - "count": len(sizes), - } - - -def _table_count(page: Any) -> int: - try: - finder = page.find_tables() - return len(getattr(finder, "tables", []) or []) - except Exception: - return 0 - - -def _extract_page_features(page: Any, page_index: int) -> dict[str, Any]: - rect = page.rect - page_area = max(_rect_area(rect), 1.0) - text = page.get_text() or "" - text_len = len(text.strip()) - image_coverage, image_count = _measure_image_coverage(page, page_area) - try: - drawings_count = len(page.get_drawings() or []) - except Exception: - drawings_count = 0 - orientation = "landscape" if float(rect.width) > float(rect.height) else "portrait" - text_density = text_len / page_area * 10000 - table_count = _table_count(page) - is_blank_like = text_len < 20 and image_coverage < 0.02 and drawings_count < 5 - - return { - "page_index": page_index, - "page_number": page_index + 1, - "width": float(rect.width), - "height": float(rect.height), - "orientation": orientation, - "text_length": text_len, - "text_density": round(text_density, 4), - "image_count": image_count, - "image_coverage": round(image_coverage, 4), - "table_count": table_count, - "drawings_count": drawings_count, - "font_size_stats": _font_stats(page), - "is_blank_like": is_blank_like, - "text_preview": " ".join(text.split())[:500], - } - - -@worker -def _sample_pages_worker( - queue, - pdf_path: str, - strategy: str, - max_samples: int, -) -> None: - import pymupdf - - doc = pymupdf.open(pdf_path) - try: - page_count = int(doc.page_count) - safe_strategy: SampleStrategy = "stratified" - if strategy == "uniform": - safe_strategy = "uniform" - elif strategy == "key_pages": - safe_strategy = "key_pages" - indices = _choose_sample_indices( - page_count, - strategy=safe_strategy, - max_samples=max_samples, - ) - sampled_pages = [_extract_page_features(doc[idx], idx) for idx in indices] - finally: - doc.close() - gc.collect() - - queue.put( - { - "ok": True, - "page_count": page_count, - "sample_indices": indices, - "sampled_pages": sampled_pages, - } - ) - - -def sample_pages( - pdf_path: str, - *, - strategy: SampleStrategy = "stratified", - max_samples: int = 25, - timeout: int = 120, -) -> dict[str, Any]: - """Sample structural page features in an isolated PyMuPDF child process.""" - result = run_in_child_process( - _sample_pages_worker, - pdf_path, - strategy, - max_samples, - timeout=timeout, - ) - return { - "page_count": int(result.get("page_count") or 0), - "sample_indices": list(result.get("sample_indices") or []), - "sampled_pages": list(result.get("sampled_pages") or []), - } diff --git a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py deleted file mode 100644 index b073236a5..000000000 --- a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Selective VLM inspection for ambiguous sampled PDF pages.""" - -from __future__ import annotations - -import base64 -import os -import tempfile -from typing import Any - -from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker -from loguru import logger -from openai.types.chat import ( - ChatCompletionContentPartImageParam, - ChatCompletionContentPartParam, - ChatCompletionContentPartTextParam, - ChatCompletionMessageParam, -) - -DEFAULT_QUESTION = ( - "Inspect these PDF page screenshots. For each page, decide whether it is a " - "table-heavy page, image-heavy page, table of contents, blank/sparse page, " - "landscape page, single-image page, or normal page. Return compact JSON with " - "items: [{page, judgement, confidence, note}]." -) - - -@worker -def _render_vlm_pages_worker( - queue, - pdf_path: str, - page_indices: list[int], - dpi: int, - out_dir: str, -) -> None: - import pymupdf - - doc = pymupdf.open(pdf_path) - rendered: list[dict[str, Any]] = [] - try: - mat = pymupdf.Matrix(dpi / 72, dpi / 72) - for idx in page_indices: - if idx < 0 or idx >= doc.page_count: - continue - page = doc[idx] - pix = page.get_pixmap(matrix=mat, alpha=False) - out_path = os.path.join(out_dir, f"vlm_probe_p{idx + 1}.png") - pix.save(out_path) - rendered.append({"page_index": idx, "page_number": idx + 1, "path": out_path}) - pix = None - page = None - finally: - doc.close() - queue.put({"ok": True, "rendered": rendered}) - - -def _png_to_data_url(path: str) -> str | None: - try: - with open(path, "rb") as file: - data = base64.b64encode(file.read()).decode("utf-8") - return f"data:image/png;base64,{data}" - except Exception as exc: - logger.warning(f"[document_agent.vlm_inspect] failed to encode {path}: {exc}") - return None - - -def _call_vlm( - *, - image_items: list[dict[str, Any]], - question: str, - model: str | None = None, - max_tokens: int = 900, -) -> tuple[str, dict[str, int]]: - from shared.core.config import settings - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - effective_model = model or settings.IMAGE_MODEL or "qwen3.5-flash" - client = get_openai_client(model=effective_model) - content: list[ChatCompletionContentPartParam] = [ - ChatCompletionContentPartTextParam( - type="text", - text=question, - ) - ] - - for item in image_items: - url = item.get("data_url") - if not url: - continue - content.append( - ChatCompletionContentPartTextParam( - type="text", - text=f"Page {item['page_number']}:", - ) - ) - content.append( - ChatCompletionContentPartImageParam( - type="image_url", - image_url={"url": url}, - ) - ) - - messages: list[ChatCompletionMessageParam] = [{"role": "user", "content": content}] - return client.chat_completion_with_usage( - messages=messages, - model=effective_model, - temperature=0.0, - max_tokens=max_tokens, - ) - - -def vlm_inspect_pages( - pdf_path: str, - *, - page_indices: list[int], - question: str = DEFAULT_QUESTION, - dpi: int = 120, - model: str | None = None, - max_tokens: int = 900, - timeout: int = 60, -) -> dict[str, Any]: - """Render selected 0-based pages and ask the configured VLM to inspect them.""" - if not page_indices: - return {"observations": [], "raw_response": "", "usage": {}} - - with tempfile.TemporaryDirectory(prefix="doc_agent_vlm_") as tmp_dir: - result = run_in_child_process( - _render_vlm_pages_worker, - pdf_path, - sorted(set(page_indices)), - dpi, - tmp_dir, - timeout=timeout, - ) - image_items = [] - for item in result.get("rendered", []) or []: - data_url = _png_to_data_url(item["path"]) - if data_url is not None: - image_items.append({**item, "data_url": data_url}) - - if not image_items: - return {"observations": [], "raw_response": "", "usage": {}} - - response, usage = _call_vlm( - image_items=image_items, - question=question, - model=model, - max_tokens=max_tokens, - ) - return { - "observations": [ - { - "page": item["page_number"], - "page_index": item["page_index"], - "vlm_judgement": response, - "confidence": None, - } - for item in image_items - ], - "raw_response": response, - "usage": usage, - } 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 deleted file mode 100644 index 1db72592b..000000000 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ /dev/null @@ -1,361 +0,0 @@ -"""Shard planning for the Phase 1 split agent.""" - -from __future__ import annotations - -import json -from collections import Counter -from hashlib import sha256 -from typing import Any - -from app.services.document_agent.manifest import ( - GlobalSignals, - ShardManifest, - ShardSignal, - SpecialPage, -) -from app.services.document_agent.tools.llm_json import extract_json_object -from loguru import logger - -PROMPT = """You are planning PDF shards for a downstream parser. -Return ONLY valid json. - -Goal: -- Cover every page from 1 to page_count exactly once. -- Prefer shards as close to max_pages_per_shard pages as possible without exceeding it. -- Do not cut through obvious table-heavy ranges, continuous image/landscape blocks, or likely TOC pages. -- Align cuts near safer normal/sparse pages when possible. - -JSON schema: -{ - "cuts": [ - {"start": 1, "end": 199, "predominant_kind": "text_dense", "rationale": "short reason"} - ], - "global_notes": "short summary" -} - -Allowed predominant_kind values: text_dense, table_heavy, image_heavy, mixed, landscape_block, toc, sparse. -""" - -ALLOWED_PREDOMINANT = { - "text_dense", - "table_heavy", - "image_heavy", - "mixed", - "landscape_block", - "toc", - "sparse", -} - - -def _special_by_page(classifications: dict[str, Any]) -> dict[int, dict[str, Any]]: - by_page = {} - for item in classifications.get("pages", []) or []: - if not isinstance(item, dict): - continue - page = int(item.get("page") or 0) - if page > 0: - by_page[page] = item - return by_page - - -def _kind_for_range(start: int, end: int, by_page: dict[int, dict[str, Any]]) -> str: - kinds = [ - str(item.get("special_kind") or item.get("kind") or "normal") - for page, item in by_page.items() - if start <= page <= end - ] - if not kinds: - return "text_dense" - counts = Counter(kinds) - special_total = sum(count for kind, count in counts.items() if kind != "normal") - if special_total == 0: - return "text_dense" - top_kind, top_count = counts.most_common(1)[0] - if top_kind in {"table_heavy"}: - return "table_heavy" - if top_kind in {"image_heavy", "single_image"}: - return "image_heavy" - if top_kind == "landscape": - return "landscape_block" - if top_kind == "toc": - return "toc" - if top_kind in {"blank", "sparse"} and top_count >= max(1, (end - start + 1) // 2): - return "sparse" - return "mixed" - - -def _fallback_cuts( - page_count: int, - *, - max_pages_per_shard: int, - classifications: dict[str, Any], -) -> list[dict[str, Any]]: - by_page = _special_by_page(classifications) - cuts = [] - start = 1 - while start <= page_count: - target_end = min(start + max_pages_per_shard - 1, page_count) - end = target_end - if target_end < page_count: - # Prefer a safe boundary close to the shard limit without exceeding it. - target_kind = str(by_page.get(target_end, {}).get("special_kind") or "normal") - if target_kind not in {"blank", "sparse", "normal"}: - window_start = max(start, target_end - 5) - candidates = [] - for page in range(window_start, target_end): - kind = str(by_page.get(page, {}).get("special_kind") or "normal") - priority = {"blank": 0, "sparse": 1, "normal": 2}.get(kind) - if priority is not None: - candidates.append((abs(page - target_end), priority, page)) - if candidates: - end = min(candidates)[2] - cuts.append( - { - "start": start, - "end": end, - "predominant_kind": _kind_for_range(start, end, by_page), - "rationale": "deterministic fallback cut near max shard size", - } - ) - start = end + 1 - return cuts - - -def _normalize_cuts( - cuts: list[dict[str, Any]], - *, - page_count: int, - max_pages_per_shard: int, - classifications: dict[str, Any], -) -> list[dict[str, Any]]: - normalized = [] - expected = 1 - for raw in cuts: - if not isinstance(raw, dict): - continue - start = int(raw.get("start") or raw.get("page_start") or 0) - end = int(raw.get("end") or raw.get("page_end") or 0) - if start != expected or end < start or end > page_count: - raise ValueError("LLM shard cuts are not contiguous") - kind = str(raw.get("predominant_kind") or "mixed") - if kind not in ALLOWED_PREDOMINANT: - kind = "mixed" - normalized.append( - { - "start": start, - "end": end, - "predominant_kind": kind, - "rationale": str(raw.get("rationale") or "")[:500], - } - ) - expected = end + 1 - if expected != page_count + 1: - raise ValueError("LLM shard cuts do not cover all pages") - if not normalized: - raise ValueError("empty LLM shard cuts") - # Avoid accepting pathological single giant cuts except naturally small docs. - if page_count > max_pages_per_shard * 2 and any( - cut["end"] - cut["start"] + 1 > max_pages_per_shard * 2 - for cut in normalized - ): - raise ValueError("LLM shard cut exceeds hard tolerance") - return normalized - - -def _build_global_signals( - *, - sampled_pages: list[dict[str, Any]], - classifications: dict[str, Any], -) -> GlobalSignals: - pages = classifications.get("pages", []) or [] - toc_pages = [ - int(page.get("page") or 0) - for page in pages - if str(page.get("special_kind") or page.get("kind")) == "toc" - ] - sample_size = len(sampled_pages) - if sample_size <= 0: - return GlobalSignals( - has_toc=bool(toc_pages), - toc_pages=toc_pages, - landscape_ratio=0.0, - table_page_ratio=0.0, - image_page_ratio=0.0, - sample_size=0, - notes=str(classifications.get("global_notes") or ""), - ) - landscape_count = sum(1 for page in sampled_pages if page.get("orientation") == "landscape") - table_count = sum( - 1 - for page in pages - if str(page.get("special_kind") or page.get("kind")) == "table_heavy" - ) - image_count = sum( - 1 - for page in pages - if str(page.get("special_kind") or page.get("kind")) in {"image_heavy", "single_image"} - ) - return GlobalSignals( - has_toc=bool(toc_pages), - toc_pages=toc_pages, - landscape_ratio=landscape_count / sample_size, - table_page_ratio=table_count / sample_size, - image_page_ratio=image_count / sample_size, - sample_size=sample_size, - notes=str(classifications.get("global_notes") or ""), - ) - - -def build_manifest_from_cuts( - *, - file_uri: str, - job_id: str, - page_count: int, - sampled_pages: list[dict[str, Any]], - classifications: dict[str, Any], - cuts: list[dict[str, Any]], - decision_log_ref: str = "local-debug", -) -> ShardManifest: - by_page = _special_by_page(classifications) - shards = [] - for cut in cuts: - start = int(cut["start"]) - end = int(cut["end"]) - special_pages = [] - for page in range(start, end + 1): - item = by_page.get(page) - if not item: - continue - kind = str(item.get("special_kind") or item.get("kind") or "normal") - if kind == "normal": - continue - special_pages.append( - SpecialPage( - page=page, - kind=kind, # type: ignore[arg-type] - confidence=float(item.get("confidence") or 0.0), - note=str(item.get("note") or ""), - ) - ) - shards.append( - ShardSignal( - page_start=start, - page_end=end, - page_offset=start - 1, - predominant_kind=cut["predominant_kind"], - special_pages=special_pages, - cut_rationale=str(cut.get("rationale") or ""), - ) - ) - - manifest = ShardManifest( - job_id=job_id, - file_uri=file_uri, - file_sha=_hash_file(file_uri), - page_count=page_count, - shard_count=len(shards), - shards=shards, - global_signals=_build_global_signals( - sampled_pages=sampled_pages, - classifications=classifications, - ), - decision_log_ref=decision_log_ref, - ) - manifest.validate() - return manifest - - -def _hash_file(file_uri: str) -> str: - digest = sha256() - try: - with open(file_uri, "rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - except OSError: - return sha256(file_uri.encode("utf-8")).hexdigest() - - -def propose_shard_plan( - *, - file_uri: str, - job_id: str, - page_count: int, - sampled_pages: list[dict[str, Any]], - classifications: dict[str, Any], - max_pages_per_shard: int = 199, - model: str | None = None, - use_llm: bool = True, -) -> dict[str, Any]: - """Produce a validated shard proposal and manifest.""" - max_pages_per_shard = max(1, int(max_pages_per_shard)) - - cuts: list[dict[str, Any]] - raw_response = "" - if use_llm and page_count > max_pages_per_shard: - try: - from shared.core.config import settings - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL - client = get_openai_client(model=effective_model) - payload = { - "page_count": page_count, - "max_pages_per_shard": max_pages_per_shard, - "sampled_pages": sampled_pages, - "page_classifications": classifications, - } - raw_response = client.chat_completion( - messages=[ - {"role": "system", "content": PROMPT}, - { - "role": "user", - "content": "Propose a shard plan as json:\n" - + json.dumps(payload, ensure_ascii=False), - }, - ], - model=effective_model, - temperature=0.0, - max_tokens=2200, - response_format={"type": "json_object"}, - ) - data = extract_json_object(raw_response) - cuts = _normalize_cuts( - list(data.get("cuts") or []), - page_count=page_count, - max_pages_per_shard=max_pages_per_shard, - classifications=classifications, - ) - if data.get("global_notes") and not classifications.get("global_notes"): - classifications = {**classifications, "global_notes": data.get("global_notes")} - except Exception as exc: - logger.warning( - f"[document_agent.propose_shard_plan] LLM planning failed, " - f"using fallback: {exc}" - ) - cuts = _fallback_cuts( - page_count, - max_pages_per_shard=max_pages_per_shard, - classifications=classifications, - ) - else: - cuts = _fallback_cuts( - page_count, - max_pages_per_shard=max_pages_per_shard, - classifications=classifications, - ) - - manifest = build_manifest_from_cuts( - file_uri=file_uri, - job_id=job_id, - page_count=page_count, - sampled_pages=sampled_pages, - classifications=classifications, - cuts=cuts, - ) - return { - "cuts": cuts, - "manifest": manifest, - "manifest_dict": manifest.to_dict(), - "raw_response": raw_response, - } diff --git a/apps/worker/app/services/document_agent/tools/scan_all_page_features.py b/apps/worker/app/services/document_agent/tools/scan_all_page_features.py index a5a97846c..bd3799704 100644 --- a/apps/worker/app/services/document_agent/tools/scan_all_page_features.py +++ b/apps/worker/app/services/document_agent/tools/scan_all_page_features.py @@ -1,13 +1,9 @@ -"""scan_all_page_features — Tool #2 for DocumentAnatomyAgent. +"""scan_all_page_features — full-page structural feature extraction. Performs a **full-page traversal** (not sampling) of a PDF, extracting structural features for every page. Runs inside an isolated PyMuPDF child process to ensure memory is freed after extraction. -Contrast with ``probe_sample_pages.py`` (which samples up to 25 pages for a -quick diagnostic): this tool is used when the agent needs precise per-page -labels to find semantic cut points. - For a 200-page A4 PDF the child process typically completes in < 8 s. """ @@ -25,7 +21,7 @@ from loguru import logger -# ── Low-level helpers (same logic as probe_sample_pages, kept local) ────────── +# ── Low-level helpers ───────────────────────────────────────────────────────── def _rect_area(rect: Any) -> float: From a487bc73fa886cf872e7450eea462d0e05922e31 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 22 May 2026 11:00:00 +0800 Subject: [PATCH 03/11] refactor: streamline DocumentAnatomyAgent by removing legacy components This commit removes the following unused files and classes related to the DocumentAnatomyAgent: - agent.py: Deleted the DocumentAnatomyAgent class, which was not integrated into the production pipeline. - page_map.py: Removed PageMap, PageFeature, and related classes that were part of the legacy structure. - tools: Deleted all tools related to page feature extraction and heading boundary detection, including scan_all_page_features.py and find_h1_boundaries.py. Updated the __init__.py files to reflect these changes and cleaned up imports accordingly. This refactor simplifies the codebase and focuses on the current implementation of the DocumentAnatomyAgent. --- .../f8a9b0c1d2e3_add_parse_agent_tables.py | 85 +++ .../app/services/document_agent/__init__.py | 30 +- .../app/services/document_agent/agent.py | 674 ------------------ .../app/services/document_agent/budget.py | 65 ++ .../services/document_agent/coordinator.py | 156 ++++ .../services/document_agent/heading_text.py | 83 +++ .../app/services/document_agent/manifest.py | 271 +++++++ .../app/services/document_agent/page_map.py | 149 ---- .../app/services/document_agent/pdf_text.py | 60 ++ .../services/document_agent/profile_agent.py | 40 ++ .../prompts/coordinator_system.md | 12 + .../prompts/propose_hierarchy_assist.md | 11 + .../app/services/document_agent/registry.py | 99 +++ .../app/services/document_agent/state.py | 47 ++ .../services/document_agent/tools/__init__.py | 31 +- .../tools/classify_page_kinds.py | 102 +++ .../tools/classify_special_pages.py | 182 ----- .../tools/find_h1_boundaries.py | 435 +++-------- .../document_agent/tools/find_toc_pages.py | 82 +++ .../services/document_agent/tools/llm_json.py | 28 - .../tools/persist_anatomy_map.py | 72 ++ .../tools/probe_page_features.py | 147 ++++ .../tools/propose_hierarchy_assist.py | 150 ++++ .../tools/propose_shard_plan.py | 176 +++++ .../tools/scan_all_page_features.py | 240 ------- .../tools/validate_anatomy_map.py | 67 ++ .../app/services/document_agent/trace.py | 117 +++ .../app/services/document_agent/validators.py | 121 ++++ .../formats/pdf/shard_splitter.py | 208 ------ apps/worker/build_manifest_sjsyj.py | 159 +++++ apps/worker/run_hierarchy_sjsyj.py | 117 +++ .../shared/models/database/__init__.py | 5 + .../models/database/document_page_plan.py | 37 + .../shared/models/database/parse_agent.py | 63 ++ .../ai/openai_compatible_client_sync.py | 115 ++- 35 files changed, 2578 insertions(+), 1858 deletions(-) create mode 100644 apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py delete mode 100644 apps/worker/app/services/document_agent/agent.py create mode 100644 apps/worker/app/services/document_agent/budget.py create mode 100644 apps/worker/app/services/document_agent/coordinator.py create mode 100644 apps/worker/app/services/document_agent/heading_text.py create mode 100644 apps/worker/app/services/document_agent/manifest.py delete mode 100644 apps/worker/app/services/document_agent/page_map.py create mode 100644 apps/worker/app/services/document_agent/pdf_text.py create mode 100644 apps/worker/app/services/document_agent/profile_agent.py create mode 100644 apps/worker/app/services/document_agent/prompts/coordinator_system.md create mode 100644 apps/worker/app/services/document_agent/prompts/propose_hierarchy_assist.md create mode 100644 apps/worker/app/services/document_agent/registry.py create mode 100644 apps/worker/app/services/document_agent/state.py create mode 100644 apps/worker/app/services/document_agent/tools/classify_page_kinds.py delete mode 100644 apps/worker/app/services/document_agent/tools/classify_special_pages.py create mode 100644 apps/worker/app/services/document_agent/tools/find_toc_pages.py delete mode 100644 apps/worker/app/services/document_agent/tools/llm_json.py create mode 100644 apps/worker/app/services/document_agent/tools/persist_anatomy_map.py create mode 100644 apps/worker/app/services/document_agent/tools/probe_page_features.py create mode 100644 apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py create mode 100644 apps/worker/app/services/document_agent/tools/propose_shard_plan.py delete mode 100644 apps/worker/app/services/document_agent/tools/scan_all_page_features.py create mode 100644 apps/worker/app/services/document_agent/tools/validate_anatomy_map.py create mode 100644 apps/worker/app/services/document_agent/trace.py create mode 100644 apps/worker/app/services/document_agent/validators.py delete mode 100644 apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py create mode 100644 apps/worker/build_manifest_sjsyj.py create mode 100644 apps/worker/run_hierarchy_sjsyj.py create mode 100644 packages/shared-python/shared/models/database/document_page_plan.py create mode 100644 packages/shared-python/shared/models/database/parse_agent.py diff --git a/apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py b/apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py new file mode 100644 index 000000000..8ba86f35e --- /dev/null +++ b/apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py @@ -0,0 +1,85 @@ +"""add parse agent tables + +Revision ID: f8a9b0c1d2e3 +Revises: f7a8b9c0d1e2 +Create Date: 2026-05-22 10:45:00.000000 + +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "f8a9b0c1d2e3" +down_revision: Union[str, Sequence[str], None] = "f7a8b9c0d1e2" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.create_table( + "parse_runs", + sa.Column("run_id", sa.String(length=36), nullable=False), + sa.Column("job_id", sa.String(length=36), nullable=False), + sa.Column("kind", sa.String(length=32), nullable=False, server_default="profile"), + sa.Column("final_status", sa.String(length=32), nullable=False), + sa.Column("rounds_count", sa.Integer(), nullable=False, server_default="0"), + sa.Column("total_tokens", sa.Integer(), nullable=False, server_default="0"), + sa.Column("total_latency_ms", sa.Integer(), nullable=False, server_default="0"), + sa.Column("summary", sa.JSON(), nullable=True), + sa.Column("started_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + sa.Column("finished_at", sa.DateTime(), nullable=True), + sa.ForeignKeyConstraint(["job_id"], ["jobs.job_id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("run_id"), + ) + op.create_index("idx_parse_runs_job_kind", "parse_runs", ["job_id", "kind"]) + op.create_index("idx_parse_runs_started", "parse_runs", ["started_at"]) + + op.create_table( + "parse_steps", + sa.Column("step_id", sa.String(length=36), nullable=False), + sa.Column("run_id", sa.String(length=36), nullable=False), + sa.Column("round_index", sa.Integer(), nullable=False, server_default="0"), + sa.Column("actor", sa.String(length=64), nullable=False), + sa.Column("action_type", sa.String(length=64), nullable=False), + sa.Column("tool_name", sa.String(length=64), nullable=True), + sa.Column("tool_args", sa.JSON(), nullable=True), + sa.Column("observation", sa.JSON(), nullable=True), + sa.Column("tokens_used", sa.Integer(), nullable=False, server_default="0"), + sa.Column("latency_ms", sa.Integer(), nullable=False, server_default="0"), + sa.Column("created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + sa.ForeignKeyConstraint(["run_id"], ["parse_runs.run_id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("step_id"), + ) + op.create_index("idx_parse_steps_run_round", "parse_steps", ["run_id", "round_index"]) + op.create_index("idx_parse_steps_tool", "parse_steps", ["tool_name"]) + + op.create_table( + "document_page_plan", + sa.Column("page_plan_id", sa.String(length=36), nullable=False), + sa.Column("job_id", sa.String(length=36), nullable=False), + sa.Column("page_count", sa.Integer(), nullable=False, server_default="0"), + sa.Column("hierarchy_assist", sa.JSON(), nullable=True), + sa.Column("shard_plan", sa.JSON(), nullable=True), + sa.Column("page_processing_plan", sa.JSON(), nullable=True), + sa.Column("global_signals", sa.JSON(), nullable=True), + sa.Column("created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + sa.ForeignKeyConstraint(["job_id"], ["jobs.job_id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("page_plan_id"), + ) + op.create_index("idx_document_page_plan_job", "document_page_plan", ["job_id"]) + op.create_index("idx_document_page_plan_created", "document_page_plan", ["created_at"]) + + +def downgrade() -> None: + op.drop_index("idx_document_page_plan_created", table_name="document_page_plan") + op.drop_index("idx_document_page_plan_job", table_name="document_page_plan") + op.drop_table("document_page_plan") + op.drop_index("idx_parse_steps_tool", table_name="parse_steps") + op.drop_index("idx_parse_steps_run_round", table_name="parse_steps") + op.drop_table("parse_steps") + op.drop_index("idx_parse_runs_started", table_name="parse_runs") + op.drop_index("idx_parse_runs_job_kind", table_name="parse_runs") + op.drop_table("parse_runs") diff --git a/apps/worker/app/services/document_agent/__init__.py b/apps/worker/app/services/document_agent/__init__.py index 96081fc35..02d729084 100644 --- a/apps/worker/app/services/document_agent/__init__.py +++ b/apps/worker/app/services/document_agent/__init__.py @@ -1,25 +1,19 @@ -"""Document Anatomy Agent package. +"""Page anatomy agent for hierarchy-first PDF profiling.""" -``DocumentAnatomyAgent`` produces a ``PageMap`` before any PDF parsing begins, -enabling semantically-correct shard decisions for large documents. -""" - -from app.services.document_agent.agent import DocumentAnatomyAgent -from app.services.document_agent.page_map import ( - CutPoint, - H1BoundaryResult, - H1Match, +from app.services.document_agent.manifest import ( + HierarchyAssistPlan, + PageAnatomyMap, PageFeature, - PageMap, - Shard, + PageLabel, + ShardPlan, ) +from app.services.document_agent.profile_agent import ProfileAgent __all__ = [ - "DocumentAnatomyAgent", - "CutPoint", - "H1BoundaryResult", - "H1Match", + "HierarchyAssistPlan", + "PageAnatomyMap", "PageFeature", - "PageMap", - "Shard", + "PageLabel", + "ProfileAgent", + "ShardPlan", ] diff --git a/apps/worker/app/services/document_agent/agent.py b/apps/worker/app/services/document_agent/agent.py deleted file mode 100644 index abc0b3931..000000000 --- a/apps/worker/app/services/document_agent/agent.py +++ /dev/null @@ -1,674 +0,0 @@ -"""DocumentAnatomyAgent — produce a PageMap for a PDF via LLM tool-calling. - -Architecture ------------- -The agent holds a **minimal tool set** — only operations that genuinely require -structured data from the PDF are exposed as tools. Classification and -heuristic logic run deterministically inside the agent; only the final -cut-point decision is delegated to the LLM (which has full context by then). - -Two-phase tool calling: -1. ``scan_all_page_features`` — collects per-page structural signals. -2. ``find_h1_boundaries`` — locates level-1 headings via text search. - -The LLM then reasons over the collected evidence to call -``propose_cut_points``, producing the final shard plan. - -Design constraints ------------------- -- No hardcoded page counts, thresholds, or prompt examples. -- Every parameter that influences splitting comes from ``settings`` or is - passed explicitly by the caller. -- Graceful degradation at every step: tools return empty/safe results rather - than raising; the agent falls back to deterministic cuts if the LLM fails. -""" - -from __future__ import annotations - -import json -import os -import traceback -from datetime import datetime, timezone -from typing import Any - -from app.services.document_agent.page_map import ( - CutPoint, - H1BoundaryResult, - H1Match, - PageFeature, - PageMap, - Shard, -) -from app.services.document_agent.tools.classify_special_pages import ( - heuristic_classify_special_pages, -) -from app.services.document_agent.tools.find_h1_boundaries import find_h1_boundaries -from app.services.document_agent.tools.scan_all_page_features import ( - scan_all_page_features, -) -from loguru import logger - - -# ── Tool schemas ─────────────────────────────────────────────────────────────── - -_TOOL_SCHEMAS: list[dict[str, Any]] = [ - { - "type": "function", - "function": { - "name": "scan_all_page_features", - "description": ( - "Perform a full structural scan of every page in the PDF. " - "Returns per-page signals: text length, image coverage, table count, " - "orientation, blank-page flag, and a text preview. " - "Always call this first." - ), - "parameters": {"type": "object", "properties": {}, "required": []}, - }, - }, - { - "type": "function", - "function": { - "name": "find_h1_boundaries", - "description": ( - "Locate level-1 headings in the document body by grepping page " - "texts against TOC entries. Returns the page numbers where each " - "level-1 heading physically starts. Call after scan_all_page_features." - ), - "parameters": {"type": "object", "properties": {}, "required": []}, - }, - }, - { - "type": "function", - "function": { - "name": "propose_cut_points", - "description": ( - "Propose a list of shard boundaries for the document. " - "Each cut_point marks where one shard ends (the next starts at " - "cut_after_page + 1). " - "Constraints you must respect:\n" - "- Shards may not exceed max_pages_per_shard pages.\n" - "- Shards should be at least min_pages_per_shard pages.\n" - "- Do not cut through landscape blocks or table-heavy pages.\n" - "- Prefer h1_heading boundaries; fall back to blank/sparse pages; " - "use forced cuts only as a last resort.\n" - "- If the document is short enough that no split is needed, " - "return an empty cut_points list." - ), - "parameters": { - "type": "object", - "properties": { - "cut_points": { - "type": "array", - "items": { - "type": "object", - "properties": { - "cut_after_page": {"type": "integer"}, - "anchor_type": { - "type": "string", - "enum": [ - "h1_heading", - "blank", - "sparse", - "forced", - ], - }, - "rationale": {"type": "string"}, - "confidence": {"type": "number"}, - }, - "required": [ - "cut_after_page", - "anchor_type", - "rationale", - ], - }, - } - }, - "required": ["cut_points"], - }, - }, - }, - { - "type": "function", - "function": { - "name": "finalize", - "description": "Signal that cut planning is complete. Call last.", - "parameters": {"type": "object", "properties": {}, "required": []}, - }, - }, -] - - -def _build_system_prompt( - split_threshold: int, - max_pages_per_shard: int, - min_pages_per_shard: int, -) -> str: - return ( - "You are a Document Anatomy Agent. Your task is to analyse a PDF and " - "decide how to split it into semantically coherent shards for downstream " - "hierarchical heading extraction.\n\n" - "Call tools in this order:\n" - "1. scan_all_page_features — always first.\n" - "2. find_h1_boundaries — always second.\n" - "3. propose_cut_points — reason carefully over the evidence:\n" - f" - Only propose cuts when the document exceeds {split_threshold} pages.\n" - f" - Each shard: between {min_pages_per_shard} and {max_pages_per_shard} pages.\n" - " - Prefer h1_heading anchors (level-1 heading start pages).\n" - " - Fall back to blank or sparse pages when headings are ambiguous.\n" - " - Use forced cuts only when no semantic boundary is available; " - "avoid cutting through landscape blocks or dense table regions.\n" - "4. finalize — always last.\n\n" - "Be concise. Do not repeat data already returned by tools." - ) - - -# ── Agent ────────────────────────────────────────────────────────────────────── - - -class DocumentAnatomyAgent: - """LLM-orchestrated agent producing a ``PageMap`` for a PDF. - - Parameters - ---------- - model_name: - Override the LLM model. Falls back to ``settings.HIERARCHY_LLM_MODEL`` - then ``settings.NORMOL_MODEL``. - split_threshold: - Minimum page count before physical splitting is considered. If - ``None``, read from ``settings.PDF_ANATOMY_SPLIT_THRESHOLD`` (default - behaviour); the caller can pass an explicit value for testing. - max_pages_per_shard: - Hard cap on shard size. Same settings fallback pattern. - min_pages_per_shard: - Prevent micro-shards that would cause overhead without benefit. - max_iterations: - Hard limit on the LLM tool-calling loop to prevent runaway calls. - """ - - def __init__( - self, - model_name: str | None = None, - split_threshold: int | None = None, - max_pages_per_shard: int | None = None, - min_pages_per_shard: int | None = None, - max_iterations: int = 10, - ) -> None: - self._model_name = model_name - self._split_threshold = split_threshold - self._max_pages_per_shard = max_pages_per_shard - self._min_pages_per_shard = min_pages_per_shard - self._max_iterations = max_iterations - - # Per-run state (reset by run()) - self._pdf_path: str = "" - self._page_features: list[PageFeature] = [] - self._page_labels: list[dict[str, Any]] = [] - self._h1_result: H1BoundaryResult | None = None - self._cut_points: list[CutPoint] = [] - self._decision_log: list[str] = [] - self._finalized: bool = False - - # ── Threshold resolution (defer to settings to avoid hardcoding) ─────────── - - def _resolve_thresholds(self) -> tuple[int, int, int]: - """Return (split_threshold, max_per_shard, min_per_shard) from settings.""" - try: - from shared.core.config import settings - - split = self._split_threshold or getattr( - settings, "PDF_ANATOMY_SPLIT_THRESHOLD", 200 - ) - max_s = self._max_pages_per_shard or getattr( - settings, "PDF_ANATOMY_MAX_PAGES_PER_SHARD", 200 - ) - min_s = self._min_pages_per_shard or getattr( - settings, "PDF_ANATOMY_MIN_PAGES_PER_SHARD", 20 - ) - except Exception: - split = self._split_threshold or 200 - max_s = self._max_pages_per_shard or 200 - min_s = self._min_pages_per_shard or 20 - return int(split), int(max_s), int(min_s) - - def _resolve_model(self) -> str: - try: - from shared.core.config import settings - - return ( - self._model_name - or getattr(settings, "HIERARCHY_LLM_MODEL", None) - or getattr(settings, "NORMOL_MODEL", None) - or "deepseek-chat" - ) - except Exception: - return self._model_name or "deepseek-chat" - - # ── Public entry point ───────────────────────────────────────────────────── - - def run(self, pdf_path: str, job_id: str) -> PageMap: - """Run the anatomy agent; always returns a valid ``PageMap``.""" - self._reset(pdf_path) - split_threshold, max_per_shard, min_per_shard = self._resolve_thresholds() - logger.info( - f"[DocumentAnatomyAgent] start '{os.path.basename(pdf_path)}' " - f"split_threshold={split_threshold} max_shard={max_per_shard}" - ) - try: - page_map = self._run_loop( - job_id, split_threshold, max_per_shard, min_per_shard - ) - except Exception as exc: - logger.error( - f"[DocumentAnatomyAgent] unrecoverable error: {exc}\n" - + traceback.format_exc() - ) - page_map = self._fallback_page_map(job_id, reason=str(exc)) - - logger.info( - f"[DocumentAnatomyAgent] done: {page_map.page_count} pages, " - f"needs_split={page_map.needs_split}, shards={len(page_map.shards)}" - ) - return page_map - - # ── Internal state ───────────────────────────────────────────────────────── - - def _reset(self, pdf_path: str) -> None: - self._pdf_path = pdf_path - self._page_features = [] - self._page_labels = [] - self._h1_result = None - self._cut_points = [] - self._decision_log = [] - self._finalized = False - - def _log(self, msg: str) -> None: - logger.info(f"[DocumentAnatomyAgent] {msg}") - self._decision_log.append(msg) - - # ── Tool dispatch ────────────────────────────────────────────────────────── - - def _tool_scan_all_page_features(self, _args: dict) -> dict[str, Any]: - self._log("→ scan_all_page_features") - features = scan_all_page_features(self._pdf_path) - self._page_features = features - - # Run heuristic classification immediately (deterministic, no LLM cost) - probe_fmt = [ - { - "page_number": f.page, - "text_length": f.text_length, - "image_coverage": f.image_coverage, - "table_count": f.table_count, - "drawings_count": f.drawings_count, - "orientation": f.orientation, - "is_blank_like": f.is_blank_like, - "text_preview": f.text_preview, - } - for f in features - ] - labels = heuristic_classify_special_pages(probe_fmt) - self._page_labels = labels.get("pages") or [] - - # Return a compact summary for the LLM context (not the full feature list) - counts: dict[str, int] = {} - for lbl in self._page_labels: - kind = str(lbl.get("special_kind") or "normal") - counts[kind] = counts.get(kind, 0) + 1 - self._log( - f" {len(features)} pages | " - + ", ".join(f"{k}={v}" for k, v in sorted(counts.items())) - ) - return { - "page_count": len(features), - "page_type_counts": counts, - # Pass sparse structural summary so LLM can reason without wall of text - "page_summary": [ - { - "page": f.page, - "kind": next( - ( - lbl.get("special_kind", "normal") - for lbl in self._page_labels - if lbl.get("page") == f.page - ), - "normal", - ), - "orientation": f.orientation, - "is_blank": f.is_blank_like, - } - for f in features - ], - } - - def _tool_find_h1_boundaries(self, _args: dict) -> dict[str, Any]: - self._log("→ find_h1_boundaries") - result = find_h1_boundaries(self._pdf_path, self._page_features) - self._h1_result = result - self._log( - f" method={result.method}, " - f"toc_pages={result.toc_pages}, " - f"h1_matches={len(result.h1_matches)}" - ) - return result.to_dict() - - def _tool_propose_cut_points(self, args: dict) -> dict[str, Any]: - self._log("→ propose_cut_points") - raw_cuts = args.get("cut_points") or [] - parsed: list[CutPoint] = [] - for raw in raw_cuts: - try: - parsed.append( - CutPoint( - cut_after_page=int(raw["cut_after_page"]), - anchor_type=raw.get("anchor_type", "forced"), - rationale=str(raw.get("rationale", ""))[:400], - confidence=float(raw.get("confidence", 1.0)), - ) - ) - except (KeyError, TypeError, ValueError): - continue - self._cut_points = sorted(parsed, key=lambda c: c.cut_after_page) - self._log(f" {len(self._cut_points)} cuts: {[c.cut_after_page for c in self._cut_points]}") - return {"accepted": len(self._cut_points)} - - def _tool_finalize(self, _args: dict) -> dict[str, Any]: - self._log("→ finalize") - self._finalized = True - return {"status": "ok"} - - _DISPATCH: dict[str, Any] = { - "scan_all_page_features": _tool_scan_all_page_features, - "find_h1_boundaries": _tool_find_h1_boundaries, - "propose_cut_points": _tool_propose_cut_points, - "finalize": _tool_finalize, - } - - def _dispatch(self, name: str, args: dict) -> dict[str, Any]: - handler = self._DISPATCH.get(name) - if handler is None: - return {"error": f"unknown tool: {name}"} - try: - return handler(self, args) - except Exception as exc: - logger.warning(f"[DocumentAnatomyAgent] tool '{name}' error: {exc}") - return {"error": str(exc)} - - # ── LLM loop ─────────────────────────────────────────────────────────────── - - def _run_loop( - self, - job_id: str, - split_threshold: int, - max_per_shard: int, - min_per_shard: int, - ) -> PageMap: - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - model = self._resolve_model() - client = get_openai_client(model=model) - system_prompt = _build_system_prompt(split_threshold, max_per_shard, min_per_shard) - - messages: list[dict[str, Any]] = [ - {"role": "system", "content": system_prompt}, - { - "role": "user", - "content": ( - f"Analyse and plan shards for: {os.path.basename(self._pdf_path)}" - ), - }, - ] - - for iteration in range(self._max_iterations): - self._log(f"iteration {iteration + 1}/{self._max_iterations}") - try: - response = client.chat_completion( - messages=messages, - model=model, - temperature=0.0, - max_tokens=2000, - tools=_TOOL_SCHEMAS, - tool_choice="auto", - ) - except Exception as exc: - logger.error(f"[DocumentAnatomyAgent] LLM call failed: {exc}") - break - - tool_calls = _parse_tool_calls(response) - if not tool_calls: - self._log("no tool calls — exiting loop") - break - - assistant_msg: dict[str, Any] = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": tc["id"], - "type": "function", - "function": { - "name": tc["name"], - "arguments": json.dumps(tc["args"], ensure_ascii=False), - }, - } - for tc in tool_calls - ], - } - messages.append(assistant_msg) - - for tc in tool_calls: - result = self._dispatch(tc["name"], tc["args"]) - messages.append( - { - "role": "tool", - "tool_call_id": tc["id"], - "content": json.dumps(result, ensure_ascii=False, default=str), - } - ) - - if self._finalized: - self._log("finalized — exiting loop") - break - else: - self._log(f"reached max_iterations={self._max_iterations}") - - # Ensure data collection happened even if LLM skipped a step - if not self._page_features: - self._log("fallback: direct scan_all_page_features") - self._tool_scan_all_page_features({}) - - if self._h1_result is None: - self._log("fallback: direct find_h1_boundaries") - self._tool_find_h1_boundaries({}) - - page_count = len(self._page_features) - - # Deterministic fallback cuts if LLM proposed nothing and doc is long - if not self._cut_points and page_count > split_threshold: - self._log("no cuts proposed — applying deterministic fallback") - self._cut_points = _deterministic_cuts( - page_count=page_count, - page_labels=self._page_labels, - h1_result=self._h1_result, - max_per_shard=max_per_shard, - min_per_shard=min_per_shard, - ) - - return self._assemble_page_map(job_id, page_count) - - # ── PageMap assembly ─────────────────────────────────────────────────────── - - def _assemble_page_map(self, job_id: str, page_count: int) -> PageMap: - shards = _cuts_to_shards(self._cut_points, page_count) - h1_result = self._h1_result or H1BoundaryResult( - toc_pages=[], h1_matches=[], method="none" - ) - return PageMap( - job_id=job_id, - file_path=self._pdf_path, - page_count=page_count, - h1_result=h1_result, - page_features=self._page_features, - shards=shards, - needs_split=len(shards) > 1, - global_signals=_global_signals(self._page_features, self._page_labels, h1_result), - agent_decision_log=list(self._decision_log), - created_at=datetime.now(timezone.utc), - ) - - def _fallback_page_map(self, job_id: str, reason: str) -> PageMap: - page_count = len(self._page_features) - shards = [Shard(page_start=1, page_end=max(page_count, 1), page_offset=0)] - return PageMap( - job_id=job_id, - file_path=self._pdf_path, - page_count=page_count, - h1_result=H1BoundaryResult(toc_pages=[], h1_matches=[], method="none"), - page_features=self._page_features, - shards=shards, - needs_split=False, - global_signals={}, - agent_decision_log=self._decision_log + [f"FALLBACK: {reason}"], - created_at=datetime.now(timezone.utc), - ) - - -# ── Module helpers ───────────────────────────────────────────────────────────── - - -def _parse_tool_calls(response: Any) -> list[dict[str, Any]]: - """Normalise an LLM response into [{id, name, args}] dicts.""" - if isinstance(response, str): - return [] - choices = getattr(response, "choices", None) - if not choices: - return [] - tc_list = getattr(choices[0].message, "tool_calls", None) or [] - result = [] - for tc in tc_list: - fn = getattr(tc, "function", None) - if not fn: - continue - try: - args = json.loads(getattr(fn, "arguments", "{}") or "{}") - except (json.JSONDecodeError, TypeError): - args = {} - result.append({"id": getattr(tc, "id", ""), "name": getattr(fn, "name", ""), "args": args}) - return result - - -def _deterministic_cuts( - page_count: int, - page_labels: list[dict[str, Any]], - h1_result: H1BoundaryResult | None, - max_per_shard: int, - min_per_shard: int, -) -> list[CutPoint]: - """Produce cuts without LLM: h1 pages → blank/sparse pages → forced.""" - by_page: dict[int, str] = { - int(p.get("page") or 0): str(p.get("special_kind") or "normal") - for p in page_labels - if p.get("page") - } - avoid = {"table_heavy", "landscape"} - cuts: list[CutPoint] = [] - - # Option A: h1 heading boundaries - h1_pages = sorted({m.page for m in (h1_result.h1_matches if h1_result else [])}) - if h1_pages: - prev = 0 - for h1_page in h1_pages: - cut_page = h1_page - 1 - if cut_page <= 0 or cut_page <= prev: - continue - shard_len = cut_page - prev - if shard_len < min_per_shard: - continue - if shard_len > max_per_shard: - # Need an intermediate forced cut - forced = prev + max_per_shard - while forced < cut_page: - actual = forced - for offset in range(5): - cand = forced - offset - if cand > prev and by_page.get(cand, "normal") not in avoid: - actual = cand - break - cuts.append(CutPoint(cut_after_page=actual, anchor_type="forced", - rationale="intermediate forced cut before h1", confidence=0.5)) - prev = actual - forced = prev + max_per_shard - cuts.append(CutPoint(cut_after_page=cut_page, anchor_type="h1_heading", - rationale=f"h1 heading starts at page {h1_page}", confidence=0.9)) - prev = cut_page - if cuts: - return cuts - - # Option B: blank/sparse pages near shard boundaries - prev = 0 - while prev + max_per_shard < page_count: - target = prev + max_per_shard - actual = target - for offset in range(min(10, max_per_shard // 2)): - for cand in [target - offset, target + offset]: - if prev < cand <= page_count: - kind = by_page.get(cand, "normal") - if kind in {"blank", "sparse"} and cand - prev >= min_per_shard: - actual = cand - break - else: - continue - break - anchor = "blank" if by_page.get(actual, "normal") == "blank" else ( - "sparse" if by_page.get(actual, "normal") == "sparse" else "forced" - ) - cuts.append(CutPoint(cut_after_page=actual, anchor_type=anchor, - rationale=f"deterministic cut near shard boundary", confidence=0.6)) - prev = actual - - return cuts - - -def _cuts_to_shards(cuts: list[CutPoint], page_count: int) -> list[Shard]: - if page_count <= 0: - return [] - if not cuts: - return [Shard(page_start=1, page_end=page_count, page_offset=0)] - shards: list[Shard] = [] - prev = 0 - for cut in sorted(cuts, key=lambda c: c.cut_after_page): - cap = min(int(cut.cut_after_page), page_count) - if cap <= prev: - continue - shards.append(Shard(page_start=prev + 1, page_end=cap, page_offset=prev)) - prev = cap - if prev < page_count: - shards.append(Shard(page_start=prev + 1, page_end=page_count, page_offset=prev)) - return shards - - -def _global_signals( - features: list[PageFeature], - labels: list[dict[str, Any]], - h1_result: H1BoundaryResult, -) -> dict[str, Any]: - total = len(features) - if not total: - return {} - label_counts: dict[str, int] = {} - for lbl in labels: - k = str(lbl.get("special_kind") or "normal") - label_counts[k] = label_counts.get(k, 0) + 1 - return { - "total_pages": total, - "h1_method": h1_result.method, - "toc_pages": h1_result.toc_pages, - "h1_match_count": len(h1_result.h1_matches), - "h1_cut_candidates": h1_result.cut_candidate_pages(), - "landscape_ratio": round( - sum(1 for f in features if f.orientation == "landscape") / total, 3 - ), - "blank_ratio": round( - sum(1 for f in features if f.is_blank_like) / total, 3 - ), - "page_type_counts": label_counts, - } diff --git a/apps/worker/app/services/document_agent/budget.py b/apps/worker/app/services/document_agent/budget.py new file mode 100644 index 000000000..d5ee9e80d --- /dev/null +++ b/apps/worker/app/services/document_agent/budget.py @@ -0,0 +1,65 @@ +"""Small synchronous budget tracker for parse-side agent planning.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class BudgetPool: + capacity: int + used: int = 0 + reserved: int = 0 + + @property + def remaining(self) -> int: + return max(self.capacity - self.used - self.reserved, 0) + + +class BudgetTracker: + """A minimal counter with the same public shape as retrieval's ledger.""" + + def __init__(self, *, plan_budget: int = 5000, max_tool_calls: int = 12) -> None: + self._plan = BudgetPool(capacity=max(int(plan_budget), 0)) + self._max_tool_calls = max(int(max_tool_calls), 1) + self._tool_calls = 0 + + def increment_tool_call(self) -> bool: + if self._tool_calls >= self._max_tool_calls: + return False + self._tool_calls += 1 + return True + + def try_reserve(self, pool: str, est: int) -> bool: + if pool != "plan": + return True + est = max(int(est), 0) + if self._plan.remaining < est: + return False + self._plan.reserved += est + return True + + def commit(self, pool: str, *, actual: int, est: int) -> None: + if pool != "plan": + return + est = max(int(est), 0) + actual = max(int(actual), 0) + self._plan.reserved = max(self._plan.reserved - est, 0) + self._plan.used = min(self._plan.capacity, self._plan.used + actual) + + def refund(self, pool: str, *, est: int) -> None: + if pool != "plan": + return + self._plan.reserved = max(self._plan.reserved - max(int(est), 0), 0) + + def snapshot(self) -> dict[str, object]: + return { + "plan": { + "capacity": self._plan.capacity, + "used": self._plan.used, + "reserved": self._plan.reserved, + "remaining": self._plan.remaining, + }, + "tool_calls": self._tool_calls, + "max_tool_calls": self._max_tool_calls, + } diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py new file mode 100644 index 000000000..8bbf52699 --- /dev/null +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -0,0 +1,156 @@ +"""Synchronous ReAct-style coordinator for the document profile agent.""" + +from __future__ import annotations + +import json +import os +from typing import Any + +from loguru import logger + +from app.services.document_agent.budget import BudgetTracker +from app.services.document_agent.manifest import PageAnatomyMap, ToolContext +from app.services.document_agent.registry import REGISTRY +from app.services.document_agent.state import AgentBlackboard, DocumentAgentState +from app.services.document_agent.tools.persist_anatomy_map import build_anatomy_map +from app.services.document_agent.trace import ParseRunRecorder + + +TRANSITIONS: dict[str, DocumentAgentState] = { + "probe.page_features": DocumentAgentState.PROBED, + "classify.page_kinds": DocumentAgentState.PROBED, + "find.toc_pages": DocumentAgentState.PROBED, + "find.h1_boundaries": DocumentAgentState.H1_FOUND, + "propose.hierarchy_assist": DocumentAgentState.H1_FOUND, + "propose.shard_plan": DocumentAgentState.H1_FOUND, + "validate.anatomy_map": DocumentAgentState.VALIDATED, + "persist.anatomy_map": DocumentAgentState.PERSISTED, +} + +REQUIRED_TOOLS = [ + "probe.page_features", + "classify.page_kinds", + "find.toc_pages", + "find.h1_boundaries", + "propose.hierarchy_assist", + "propose.shard_plan", + "validate.anatomy_map", + "persist.anatomy_map", +] + + +def _parse_tool_calls(response: Any) -> list[dict[str, Any]]: + choices = getattr(response, "choices", None) + if not choices: + return [] + message = choices[0].message + calls = getattr(message, "tool_calls", None) or [] + parsed: list[dict[str, Any]] = [] + for call in calls: + function = getattr(call, "function", None) + if function is None: + continue + try: + args = json.loads(getattr(function, "arguments", "{}") or "{}") + except json.JSONDecodeError: + args = {} + parsed.append( + { + "id": getattr(call, "id", ""), + "name": getattr(function, "name", ""), + "args": args, + } + ) + return parsed + + +class ProfileCoordinator: + def __init__( + self, + *, + pdf_path: str, + job_id: str, + output_dir: str | None = None, + db: Any | None = None, + model: str | None = None, + settings: dict[str, Any] | None = None, + ) -> None: + self.state = DocumentAgentState.INIT + self.blackboard = AgentBlackboard() + self.budget = BudgetTracker( + plan_budget=int(os.environ.get("PARSE_AGENT_PLAN_BUDGET", "5000")), + max_tool_calls=int(os.environ.get("PARSE_AGENT_MAX_TOOL_CALLS", "12")), + ) + effective_settings = settings or {} + if model: + effective_settings["model"] = model + self.ctx = ToolContext( + pdf_path=pdf_path, + job_id=job_id, + blackboard=self.blackboard, + budget=self.budget, + trace=None, + output_dir=output_dir, + settings=effective_settings, + ) + self.trace = ParseRunRecorder(job_id=job_id, db=db) + self.ctx.trace = self.trace + self.round_index = 0 + + def run(self) -> PageAnatomyMap: + # The deterministic tool chain is the primary execution contract. LLM is + # used inside proposal tools where it adds judgement, not to own ordering. + final_status = "ready" + try: + for tool_name in REQUIRED_TOOLS: + if not self.budget.increment_tool_call(): + final_status = "fallback" + break + result = REGISTRY.dispatch(tool_name, self.ctx, {}, self.state) + self.trace.record_step( + round_index=self.round_index, + actor=f"tool:{tool_name}", + action_type="tool_call", + result=result, + tool_name=tool_name, + tool_args={}, + ) + if result.status not in {"ok", "invalid"}: + raise RuntimeError(result.error or f"{tool_name} failed") + self._advance(tool_name) + self._maybe_advance_composite_state() + self.round_index += 1 + if self.state == DocumentAgentState.PERSISTED: + self.state = DocumentAgentState.READY + anatomy = build_anatomy_map(self.ctx) + self.trace.flush( + final_status=final_status, + summary=anatomy.trace_summary | self.trace.summary(), + ) + return anatomy + except Exception as exc: + logger.error(f"[document_agent] profile failed: {exc}") + self.state = DocumentAgentState.FAILED + self.trace.flush(final_status="failed", summary={"error": str(exc)}) + raise + + def _advance(self, tool_name: str) -> None: + next_state = TRANSITIONS.get(tool_name, self.state) + self.state = next_state + self.blackboard.mark(self.state) + + def _maybe_advance_composite_state(self) -> None: + if ( + self.state == DocumentAgentState.PROBED + and self.blackboard.page_labels + and self.blackboard.toc_result is not None + ): + self.state = DocumentAgentState.CLASSIFIED + self.blackboard.mark(self.state) + if ( + self.state == DocumentAgentState.H1_FOUND + and self.blackboard.hierarchy_assist is not None + and self.blackboard.shard_plan is not None + ): + self.state = DocumentAgentState.PLANNED + self.blackboard.mark(self.state) diff --git a/apps/worker/app/services/document_agent/heading_text.py b/apps/worker/app/services/document_agent/heading_text.py new file mode 100644 index 000000000..cae0f0d3b --- /dev/null +++ b/apps/worker/app/services/document_agent/heading_text.py @@ -0,0 +1,83 @@ +"""Heading extraction and matching helpers.""" + +from __future__ import annotations + +import re +import unicodedata + + +LEADING_NUMBER_RE = re.compile( + r"""^ + (?: + 第\s*[零一二三四五六七八九十百千\d]+\s*[章节篇部分] + | [零一二三四五六七八九十百千]+\s*[、。，,] + | [（(]\s*[零一二三四五六七八九十百千\d]+\s*[）)] + | \d+(?:\.\d+)*\.?\s* + | [IVXLCDM]+\.?\s+ + | [A-Za-z]\.\s+ + | Chapter\s+\w+\s* + ) + """, + re.VERBOSE | re.IGNORECASE, +) + +H1_LINE_RE = re.compile( + r"""^\s* + (?: + 第\s*[零一二三四五六七八九十百千\d]+\s*[章篇部] + | [零一二三四五六七八九十百千]+\s*[、。] + | \d+\s*[\.\s] + | Chapter\s+\w+ + | [IVXLCDM]+\.?\s+\w + ) + """, + re.VERBOSE | re.IGNORECASE, +) + +SUBHEADING_RE = re.compile(r"^\s*\d+\.\d+", re.IGNORECASE) +PAGE_SUFFIX_RE = re.compile(r"[\s\.\-·…]+\d+\s*$") + + +def normalize_heading(text: str) -> str: + text = unicodedata.normalize("NFKC", text or "") + text = re.sub(r"\s+", " ", text).strip() + stripped = LEADING_NUMBER_RE.sub("", text).strip() + return stripped if stripped else text + + +def has_numbering(text: str) -> bool: + return bool(LEADING_NUMBER_RE.match(text or "")) + + +def clean_toc_line(line: str) -> str: + return PAGE_SUFFIX_RE.sub("", line or "").strip() + + +def looks_like_h1_line(line: str) -> bool: + stripped = (line or "").strip() + if not stripped or SUBHEADING_RE.match(stripped): + return False + return bool(H1_LINE_RE.match(stripped)) + + +def candidate_allowed(title: str) -> bool: + normalized = normalize_heading(title) + if not normalized: + return False + if len(normalized) < 4 and not has_numbering(title): + return False + return True + + +def fuzzy_match(needle: str, haystack: str) -> bool: + if not needle: + return False + normalized_haystack = re.sub( + r"\s+", + " ", + unicodedata.normalize("NFKC", haystack or ""), + ) + if needle in normalized_haystack: + return True + stripped = normalize_heading(needle) + return bool(stripped and len(stripped) >= 4 and stripped in normalized_haystack) diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py new file mode 100644 index 000000000..16515fafd --- /dev/null +++ b/apps/worker/app/services/document_agent/manifest.py @@ -0,0 +1,271 @@ +"""Contracts for the hierarchy-first document profile agent.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from typing import Any, Literal + + +PageKind = Literal[ + "cover", + "toc", + "preface", + "normal", + "chapter_start", + "section_start", + "table_heavy", + "image_heavy", + "single_image", + "blank", + "separator", + "appendix", + "scan_like", + "landscape", + "sparse", +] + + +@dataclass +class PageFeature: + page: int + raw_text_length: int + text_density: float + image_coverage: float + image_count: int + table_count: int + drawings_count: int + orientation: Literal["portrait", "landscape"] + width: float + height: float + is_blank_like: bool + text_lines_preview: list[str] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class PageLabel: + page: int + kind: PageKind + confidence: float + evidence: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class TocCandidate: + title: str + normalized_title: str + source_page: int + line_index: int + numbering: str = "" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class TocResult: + toc_pages: list[int] = field(default_factory=list) + candidates: list[TocCandidate] = field(default_factory=list) + method: Literal["toc_marker", "none"] = "none" + notes: str = "" + + def to_dict(self) -> dict[str, Any]: + data = asdict(self) + data["candidates"] = [candidate.to_dict() for candidate in self.candidates] + return data + + +@dataclass +class H1Candidate: + title: str + page: int + confidence: float + matched_line: str + source: Literal["toc_exact_top", "toc_fuzzy_top", "heading_grep", "none"] + evidence: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class H1BoundaryResult: + h1_candidates: list[H1Candidate] = field(default_factory=list) + method: Literal["toc_grep", "heading_grep", "none"] = "none" + notes: str = "" + + def to_dict(self) -> dict[str, Any]: + data = asdict(self) + data["h1_candidates"] = [candidate.to_dict() for candidate in self.h1_candidates] + return data + + +@dataclass +class BoundaryHint: + page: int + anchor_type: Literal["h1_boundary", "blank_separator", "separator", "forced_max_size"] + confidence: float + evidence: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class HierarchyAssistPlan: + exclude_pages_from_title_candidates: list[int] = field(default_factory=list) + prefer_h1_start_pages: list[H1Candidate] = field(default_factory=list) + suppress_title_pages: list[int] = field(default_factory=list) + section_boundary_hints: list[BoundaryHint] = field(default_factory=list) + smart_parse_recommendation: Literal["off", "normal", "aggressive"] = "normal" + rationale: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "exclude_pages_from_title_candidates": list( + self.exclude_pages_from_title_candidates + ), + "prefer_h1_start_pages": [ + candidate.to_dict() for candidate in self.prefer_h1_start_pages + ], + "suppress_title_pages": list(self.suppress_title_pages), + "section_boundary_hints": [ + hint.to_dict() for hint in self.section_boundary_hints + ], + "smart_parse_recommendation": self.smart_parse_recommendation, + "rationale": self.rationale, + } + + +@dataclass +class ValidationReport: + valid: bool + errors: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class Shard: + shard_index: int + page_start: int + page_end: int + page_offset: int + anchor_type: Literal["h1_boundary", "blank_separator", "forced_max_size"] + anchor_evidence: str + confidence: float + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class ShardPlan: + enabled: bool + reason: Literal["too_large", "not_needed", "parser_stability", "hierarchy_isolation"] + shards: list[Shard] = field(default_factory=list) + validation: ValidationReport = field( + default_factory=lambda: ValidationReport(valid=True) + ) + + def to_dict(self) -> dict[str, Any]: + return { + "enabled": self.enabled, + "reason": self.reason, + "shards": [shard.to_dict() for shard in self.shards], + "validation": self.validation.to_dict(), + } + + +@dataclass +class PagePlanEntry: + page_index: int + strategy: Literal["vlm_detail", "vlm_lite", "text_only", "skip_tagging"] + expected_kind: str + rationale: str + estimated_cost_tokens: int = 0 + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class PageProcessingPlan: + entries: list[PagePlanEntry] = field(default_factory=list) + global_strategy_summary: dict[str, int] = field(default_factory=dict) + rationale: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "entries": [entry.to_dict() for entry in self.entries], + "global_strategy_summary": dict(self.global_strategy_summary), + "rationale": self.rationale, + } + + +@dataclass +class PageAnatomyMap: + job_id: str + file_path: str + page_count: int + page_features: list[PageFeature] + page_labels: list[PageLabel] + toc_result: TocResult + h1_result: H1BoundaryResult + hierarchy_assist: HierarchyAssistPlan + shard_plan: ShardPlan + page_processing_plan: PageProcessingPlan | None = None + global_signals: dict[str, Any] = field(default_factory=dict) + trace_summary: dict[str, Any] = field(default_factory=dict) + created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + version: str = "1.0" + + def to_dict(self) -> dict[str, Any]: + return { + "version": self.version, + "job_id": self.job_id, + "file_path": self.file_path, + "page_count": self.page_count, + "page_features": [feature.to_dict() for feature in self.page_features], + "page_labels": [label.to_dict() for label in self.page_labels], + "toc_result": self.toc_result.to_dict(), + "h1_result": self.h1_result.to_dict(), + "hierarchy_assist": self.hierarchy_assist.to_dict(), + "shard_plan": self.shard_plan.to_dict(), + "page_processing_plan": ( + self.page_processing_plan.to_dict() + if self.page_processing_plan is not None + else None + ), + "global_signals": dict(self.global_signals), + "trace_summary": dict(self.trace_summary), + "created_at": self.created_at.isoformat(), + } + + +@dataclass +class ToolResult: + status: str + payload: dict[str, Any] = field(default_factory=dict) + latency_ms: int = 0 + error: str | None = None + tokens_used: int = 0 + + +@dataclass +class ToolContext: + pdf_path: str + job_id: str + blackboard: Any + budget: Any + trace: Any + output_dir: str | None = None + settings: dict[str, Any] = field(default_factory=dict) diff --git a/apps/worker/app/services/document_agent/page_map.py b/apps/worker/app/services/document_agent/page_map.py deleted file mode 100644 index fc59ac2ce..000000000 --- a/apps/worker/app/services/document_agent/page_map.py +++ /dev/null @@ -1,149 +0,0 @@ -"""PageMap — the output contract of the Document Anatomy Agent. -""" - -from __future__ import annotations - -from dataclasses import asdict, dataclass, field -from datetime import datetime, timezone -from typing import Any, Literal - - -# ── Page-level feature snapshot ─────────────────────────────────────────────── - - -@dataclass -class PageFeature: - """Raw structural measurements for a single PDF page (1-based numbering).""" - - page: int - text_length: int - text_density: float # chars per 10k pt² of page area - image_coverage: float # fraction [0, 1] - image_count: int - table_count: int - drawings_count: int - orientation: Literal["portrait", "landscape"] - width: float # points - height: float # points - is_blank_like: bool - text_preview: str # first N chars of page text (configurable) - - # Reserved for future Page Memory integration - embedding_ref: str | None = None - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -# ── H1 boundary detection result ───────────────────────────────────────────── - - -@dataclass -class H1Match: - """A single level-1 heading found via text search.""" - - title: str # normalized heading text from TOC - page: int # 1-based page where heading was found in the body - confidence: float # 1.0 = exact match, < 1.0 = fuzzy - match_text: str # the actual text snippet that matched on the page - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -@dataclass -class H1BoundaryResult: - """Output of the ``find_h1_boundaries`` tool.""" - - toc_pages: list[int] # pages identified as TOC (by text content) - h1_matches: list[H1Match] # h1 headings and their body pages - method: Literal["toc_grep", "heading_grep", "none"] - notes: str = "" - - def cut_candidate_pages(self) -> list[int]: - """Pages just before each h1 heading page — natural cut points.""" - pages = sorted({m.page for m in self.h1_matches if m.page > 1}) - # Cut before the chapter starts (i.e., end the previous shard at page-1) - return [p - 1 for p in pages if p > 1] - - def to_dict(self) -> dict[str, Any]: - data = asdict(self) - data["h1_matches"] = [m.to_dict() for m in self.h1_matches] - return data - - -# ── Cut-point ───────────────────────────────────────────────────────────────── - - -@dataclass -class CutPoint: - """A proposed shard boundary produced by the agent.""" - - cut_after_page: int - rationale: str - anchor_type: Literal["h1_heading", "blank", "sparse", "forced"] - confidence: float = 1.0 - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -# ── Shard ────────────────────────────────────────────────────────────────────── - - -@dataclass -class Shard: - """Lightweight shard descriptor. ``page_offset`` is used to correct - ``page_nums`` in sub-PDF DataFrames back to the original document's - page numbering.""" - - page_start: int # 1-based, inclusive - page_end: int # 1-based, inclusive - page_offset: int # = page_start - 1 - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -# ── PageMap ─────────────────────────────────────────────────────────────────── - - -@dataclass -class PageMap: - """Complete anatomy report produced by ``DocumentAnatomyAgent``. - - Downstream consumers: - - ``formats/pdf/parser.py`` — decides whether to physically split. - - ``structure/layout_parser.py`` (future) — page labels annotate headings. - - Page Memory (future) — ``page_features`` become the structural component. - """ - - job_id: str - file_path: str - page_count: int - h1_result: H1BoundaryResult - page_features: list[PageFeature] - shards: list[Shard] - needs_split: bool - global_signals: dict[str, Any] - agent_decision_log: list[str] = field(default_factory=list) - created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) - version: str = "1.0" - - def page_feature_map(self) -> dict[int, PageFeature]: - return {pf.page: pf for pf in self.page_features} - - def to_dict(self) -> dict[str, Any]: - return { - "version": self.version, - "job_id": self.job_id, - "file_path": self.file_path, - "page_count": self.page_count, - "needs_split": self.needs_split, - "h1_result": self.h1_result.to_dict(), - "shards": [s.to_dict() for s in self.shards], - "global_signals": self.global_signals, - "agent_decision_log": self.agent_decision_log, - "page_features": [pf.to_dict() for pf in self.page_features], - "created_at": self.created_at.isoformat(), - } diff --git a/apps/worker/app/services/document_agent/pdf_text.py b/apps/worker/app/services/document_agent/pdf_text.py new file mode 100644 index 000000000..3fc2dd544 --- /dev/null +++ b/apps/worker/app/services/document_agent/pdf_text.py @@ -0,0 +1,60 @@ +"""PyMuPDF helpers used by document-agent tools.""" + +from __future__ import annotations + +import gc +from typing import Any + +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) + + +def normalize_spaces(text: str) -> str: + return " ".join((text or "").split()) + + +@worker +def _read_page_texts_worker(queue, pdf_path: str, pages: list[int]) -> None: + import pymupdf # type: ignore[import] + + texts: dict[int, str] = {} + try: + doc = pymupdf.open(pdf_path) + for page in pages: + idx = page - 1 + if 0 <= idx < doc.page_count: + texts[page] = doc[idx].get_text() or "" + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "texts": texts}) + + +def read_page_texts( + pdf_path: str, + pages: list[int], + *, + timeout: int = 180, +) -> dict[int, str]: + if not pages: + return {} + result = run_in_child_process(_read_page_texts_worker, pdf_path, pages, timeout=timeout) + return {int(k): str(v) for k, v in (result.get("texts") or {}).items()} + + +def meaningful_lines(text: str) -> list[str]: + return [normalize_spaces(line) for line in text.splitlines() if normalize_spaces(line)] + + +def top_lines(text: str, *, max_lines: int = 20) -> list[str]: + lines = meaningful_lines(text) + return lines[: max(max_lines, 0)] + + +def compact_payload_keys(payload: dict[str, Any]) -> list[str]: + return sorted(str(key) for key in payload.keys()) diff --git a/apps/worker/app/services/document_agent/profile_agent.py b/apps/worker/app/services/document_agent/profile_agent.py new file mode 100644 index 000000000..6a73cf74d --- /dev/null +++ b/apps/worker/app/services/document_agent/profile_agent.py @@ -0,0 +1,40 @@ +"""Public entrypoint for document page anatomy profiling.""" + +from __future__ import annotations + +import os +from typing import Any + +from app.services.document_agent.coordinator import ProfileCoordinator +from app.services.document_agent.manifest import PageAnatomyMap + + +class ProfileAgent: + def __init__( + self, + *, + model: str | None = None, + settings: dict[str, Any] | None = None, + ) -> None: + self._model = model + self._settings = settings or {} + + def run( + self, + file_path: str, + job_id: str, + *, + output_dir: str | None = None, + db: Any | None = None, + ) -> PageAnatomyMap: + if not os.path.exists(file_path): + raise FileNotFoundError(file_path) + coordinator = ProfileCoordinator( + pdf_path=file_path, + job_id=job_id, + output_dir=output_dir, + db=db, + model=self._model, + settings=self._settings, + ) + return coordinator.run() diff --git a/apps/worker/app/services/document_agent/prompts/coordinator_system.md b/apps/worker/app/services/document_agent/prompts/coordinator_system.md new file mode 100644 index 000000000..dcfb075d0 --- /dev/null +++ b/apps/worker/app/services/document_agent/prompts/coordinator_system.md @@ -0,0 +1,12 @@ +You are coordinating page anatomy profiling for a PDF. + +Goal: +- identify special pages that affect hierarchy extraction +- identify reliable H1 starts and pages that should not become title candidates +- produce a shard plan for long PDF-to-Markdown execution + +Rules: +- use only the available tools for the current state +- do not invent page numbers or section titles +- prefer evidence from TOC pages and page-start heading matches +- treat tool outputs as structured evidence, not prose diff --git a/apps/worker/app/services/document_agent/prompts/propose_hierarchy_assist.md b/apps/worker/app/services/document_agent/prompts/propose_hierarchy_assist.md new file mode 100644 index 000000000..be6803aef --- /dev/null +++ b/apps/worker/app/services/document_agent/prompts/propose_hierarchy_assist.md @@ -0,0 +1,11 @@ +Return strict JSON for hierarchy assistance. + +Use only page numbers present in the input payload. +Do not invent headings. +Keep reasons brief and evidence-based. + +Required fields: +- exclude_pages_from_title_candidates: integer array +- suppress_title_pages: integer array +- smart_parse_recommendation: one of off, normal, aggressive +- rationale: short string diff --git a/apps/worker/app/services/document_agent/registry.py b/apps/worker/app/services/document_agent/registry.py new file mode 100644 index 000000000..5bf9d4cd2 --- /dev/null +++ b/apps/worker/app/services/document_agent/registry.py @@ -0,0 +1,99 @@ +"""Tool registry with state-based exposure.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable + +from app.services.document_agent.manifest import ToolContext, ToolResult +from app.services.document_agent.state import DocumentAgentState + +ToolHandler = Callable[[ToolContext, dict[str, Any]], ToolResult] + + +@dataclass(frozen=True) +class ToolSpec: + name: str + description: str + parameters: dict[str, Any] + allowed_states: frozenset[DocumentAgentState] + handler: ToolHandler + + def to_openai_schema(self) -> dict[str, Any]: + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": self.parameters, + }, + } + + +class ToolRegistry: + def __init__(self) -> None: + self._tools: dict[str, ToolSpec] = {} + + def register(self, spec: ToolSpec) -> None: + self._tools[spec.name] = spec + + def get(self, name: str) -> ToolSpec | None: + return self._tools.get(name) + + def catalogue_for_state(self, state: DocumentAgentState) -> list[dict[str, Any]]: + return [ + tool.to_openai_schema() + for tool in self._tools.values() + if state in tool.allowed_states + ] + + def allowed_names(self, state: DocumentAgentState) -> list[str]: + return [ + name + for name, tool in self._tools.items() + if state in tool.allowed_states + ] + + def dispatch( + self, + name: str, + ctx: ToolContext, + args: dict[str, Any], + state: DocumentAgentState, + ) -> ToolResult: + tool = self.get(name) + if tool is None: + return ToolResult(status="error", error=f"unknown tool: {name}") + if state not in tool.allowed_states: + return ToolResult( + status="state_error", + payload={"allowed_tools": self.allowed_names(state), "state": state.value}, + error=f"tool {name} is not allowed in state {state.value}", + ) + return tool.handler(ctx, args) + + +REGISTRY = ToolRegistry() + + +def register_tool( + *, + name: str, + description: str, + parameters: dict[str, Any] | None = None, + allowed_states: set[DocumentAgentState], +) -> Callable[[ToolHandler], ToolHandler]: + def _decorator(handler: ToolHandler) -> ToolHandler: + REGISTRY.register( + ToolSpec( + name=name, + description=description, + parameters=parameters + or {"type": "object", "properties": {}, "required": []}, + allowed_states=frozenset(allowed_states), + handler=handler, + ) + ) + return handler + + return _decorator diff --git a/apps/worker/app/services/document_agent/state.py b/apps/worker/app/services/document_agent/state.py new file mode 100644 index 000000000..47f3bc0b4 --- /dev/null +++ b/apps/worker/app/services/document_agent/state.py @@ -0,0 +1,47 @@ +"""State carried by the document profile agent.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +from app.services.document_agent.manifest import ( + H1BoundaryResult, + HierarchyAssistPlan, + PageFeature, + PageLabel, + ShardPlan, + TocResult, +) + + +class DocumentAgentState(str, Enum): + INIT = "init" + PROBED = "probed" + CLASSIFIED = "classified" + H1_FOUND = "h1_found" + PLANNED = "planned" + VALIDATED = "validated" + PERSISTED = "persisted" + READY = "ready" + FAILED = "failed" + PROCESSING_PLAN_PROPOSED = "processing_plan_proposed" + + +@dataclass +class AgentBlackboard: + page_count: int = 0 + page_features: list[PageFeature] = field(default_factory=list) + page_labels: list[PageLabel] = field(default_factory=list) + toc_result: TocResult | None = None + h1_result: H1BoundaryResult | None = None + hierarchy_assist: HierarchyAssistPlan | None = None + shard_plan: ShardPlan | None = None + validation_report: dict[str, Any] | None = None + global_signals: dict[str, Any] = field(default_factory=dict) + errors: list[str] = field(default_factory=list) + state_trace: list[str] = field(default_factory=list) + + def mark(self, state: DocumentAgentState) -> None: + self.state_trace.append(state.value) diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index 5056d7d49..7206a6dce 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -1,23 +1,14 @@ -"""Document Anatomy Agent tools. +"""Import tool modules so decorators register handlers.""" -- ``scan_all_page_features``: full-page structural feature extraction. -- ``find_h1_boundaries``: locates level-1 headings via text search. -- ``classify_special_pages`` / ``heuristic_classify_special_pages``: page - classification used internally by the agent (not called by LLM as a tool). -""" +from app.services.document_agent.registry import REGISTRY -from app.services.document_agent.tools.classify_special_pages import ( - classify_special_pages, - heuristic_classify_special_pages, -) -from app.services.document_agent.tools.find_h1_boundaries import find_h1_boundaries -from app.services.document_agent.tools.scan_all_page_features import ( - scan_all_page_features, -) +from . import classify_page_kinds as classify_page_kinds # noqa: F401 +from . import find_h1_boundaries as find_h1_boundaries # noqa: F401 +from . import find_toc_pages as find_toc_pages # noqa: F401 +from . import persist_anatomy_map as persist_anatomy_map # noqa: F401 +from . import probe_page_features as probe_page_features # noqa: F401 +from . import propose_hierarchy_assist as propose_hierarchy_assist # noqa: F401 +from . import propose_shard_plan as propose_shard_plan # noqa: F401 +from . import validate_anatomy_map as validate_anatomy_map # noqa: F401 -__all__ = [ - "classify_special_pages", - "find_h1_boundaries", - "heuristic_classify_special_pages", - "scan_all_page_features", -] +__all__ = ["REGISTRY"] diff --git a/apps/worker/app/services/document_agent/tools/classify_page_kinds.py b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py new file mode 100644 index 000000000..399741243 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py @@ -0,0 +1,102 @@ +"""Rule-based page kind classification.""" + +from __future__ import annotations + +import time +from typing import Any + +from app.services.document_agent.manifest import PageFeature, PageLabel, ToolContext, ToolResult +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState + + +def _joined_preview(feature: PageFeature) -> str: + return "\n".join(feature.text_lines_preview).lower() + + +def _label_feature(feature: PageFeature) -> PageLabel: + preview = _joined_preview(feature) + page = feature.page + if any(marker in preview.replace(" ", "") for marker in ("目录", "目次", "contents")): + return PageLabel( + page=page, + kind="toc", + confidence=0.86, + evidence={"signal": "toc_marker"}, + ) + if feature.is_blank_like: + return PageLabel( + page=page, + kind="blank", + confidence=0.92, + evidence={"signal": "low_text_image_drawings"}, + ) + if feature.orientation == "landscape": + return PageLabel( + page=page, + kind="landscape", + confidence=0.78, + evidence={"width": feature.width, "height": feature.height}, + ) + if feature.image_coverage >= 0.72 and feature.raw_text_length < 250: + return PageLabel( + page=page, + kind="single_image", + confidence=0.84, + evidence={"image_coverage": feature.image_coverage}, + ) + if feature.raw_text_length < 50 and feature.image_coverage >= 0.35: + return PageLabel( + page=page, + kind="scan_like", + confidence=0.76, + evidence={ + "raw_text_length": feature.raw_text_length, + "image_coverage": feature.image_coverage, + }, + ) + if feature.table_count > 0 or feature.drawings_count >= 80: + return PageLabel( + page=page, + kind="table_heavy", + confidence=0.72, + evidence={ + "table_count": feature.table_count, + "drawings_count": feature.drawings_count, + }, + ) + if feature.image_coverage >= 0.35: + return PageLabel( + page=page, + kind="image_heavy", + confidence=0.72, + evidence={"image_coverage": feature.image_coverage}, + ) + if feature.raw_text_length < 80: + return PageLabel( + page=page, + kind="sparse", + confidence=0.67, + evidence={"raw_text_length": feature.raw_text_length}, + ) + return PageLabel(page=page, kind="normal", confidence=0.65, evidence={}) + + +@register_tool( + name="classify.page_kinds", + description="Classify every probed page into structural page kinds using deterministic signals.", + allowed_states={DocumentAgentState.PROBED}, +) +def classify_page_kinds(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + labels = [_label_feature(feature) for feature in ctx.blackboard.page_features] + ctx.blackboard.page_labels = labels + counts: dict[str, int] = {} + for label in labels: + counts[label.kind] = counts.get(label.kind, 0) + 1 + ctx.blackboard.global_signals["page_kind_counts"] = counts + return ToolResult( + status="ok", + payload={"page_kind_counts": counts}, + latency_ms=int((time.monotonic() - start) * 1000), + ) diff --git a/apps/worker/app/services/document_agent/tools/classify_special_pages.py b/apps/worker/app/services/document_agent/tools/classify_special_pages.py deleted file mode 100644 index 83072ccaf..000000000 --- a/apps/worker/app/services/document_agent/tools/classify_special_pages.py +++ /dev/null @@ -1,182 +0,0 @@ -"""Classify pages into special-page kinds (heuristic + optional LLM).""" - -from __future__ import annotations - -import json -from typing import Any, Literal - -from app.services.document_agent.tools.llm_json import extract_json_object -from loguru import logger - -SpecialKind = Literal[ - "toc", - "blank", - "sparse", - "table_heavy", - "image_heavy", - "landscape", - "single_image", - "normal", -] - -ALLOWED_KINDS: set[str] = { - "toc", - "blank", - "sparse", - "table_heavy", - "image_heavy", - "landscape", - "single_image", - "normal", -} - -PROMPT = """You are classifying sampled PDF pages for a document split planner. -Return ONLY valid json. - -Allowed special_kind values: -- toc: table of contents / contents pages -- blank: blank page -- sparse: very little useful content -- table_heavy: mostly tables or dense tabular rules -- image_heavy: many charts/photos/figures -- landscape: landscape page that should not be split through a landscape block -- single_image: one large screenshot/scanned image dominates the page -- normal: ordinary text page - -Use structural features first, and use VLM observations when present. Be conservative: -only mark toc/table/image/landscape/single_image when there is clear evidence. - -JSON schema: -{ - "pages": [ - {"page": 1, "special_kind": "normal", "confidence": 0.75, "note": "short reason"} - ], - "global_notes": "short summary" -} -""" - - -def _heuristic_kind(page: dict[str, Any]) -> tuple[SpecialKind, float, str]: - text = str(page.get("text_preview") or "") - text_len = int(page.get("text_length") or 0) - image_coverage = float(page.get("image_coverage") or 0.0) - table_count = int(page.get("table_count") or 0) - drawings_count = int(page.get("drawings_count") or 0) - orientation = str(page.get("orientation") or "") - is_blank_like = bool(page.get("is_blank_like")) - - toc_markers = ["目录", "contents", "table of contents"] - if any(marker.lower() in text.lower() for marker in toc_markers): - return "toc", 0.82, "text preview contains TOC marker" - if is_blank_like: - return "blank", 0.9, "very low text/image/drawing signal" - if image_coverage >= 0.72 and text_len < 250: - return "single_image", 0.82, "one or more images dominate the page" - if table_count > 0 or drawings_count >= 80: - return "table_heavy", 0.72, "table detector or dense ruled drawings fired" - if image_coverage >= 0.35: - return "image_heavy", 0.72, "high image coverage" - if orientation == "landscape": - return "landscape", 0.75, "page is landscape" - if text_len < 80: - return "sparse", 0.68, "short text and no stronger special signal" - return "normal", 0.65, "no special signal" - - -def heuristic_classify_special_pages( - sampled_pages: list[dict[str, Any]], -) -> dict[str, Any]: - pages = [] - for page in sampled_pages: - kind, confidence, note = _heuristic_kind(page) - pages.append( - { - "page": int(page.get("page_number") or 0), - "special_kind": kind, - "confidence": confidence, - "note": note, - } - ) - return {"pages": pages, "global_notes": "heuristic classification"} - - -def _normalize_llm_pages( - data: dict[str, Any], - sampled_pages: list[dict[str, Any]], -) -> dict[str, Any]: - by_page = {int(page.get("page_number") or 0): page for page in sampled_pages} - pages = [] - for item in data.get("pages", []) or []: - if not isinstance(item, dict): - continue - page_number = int(item.get("page") or item.get("page_number") or 0) - if page_number not in by_page: - continue - kind = str(item.get("special_kind") or item.get("kind") or "normal") - if kind not in ALLOWED_KINDS: - kind = "normal" - confidence = max(0.0, min(float(item.get("confidence") or 0.0), 1.0)) - if confidence <= 0: - confidence = 0.5 - pages.append( - { - "page": page_number, - "special_kind": kind, - "confidence": confidence, - "note": str(item.get("note") or "")[:300], - } - ) - - seen = {item["page"] for item in pages} - for fallback in heuristic_classify_special_pages(sampled_pages)["pages"]: - if fallback["page"] not in seen: - pages.append(fallback) - pages.sort(key=lambda item: item["page"]) - return { - "pages": pages, - "global_notes": str(data.get("global_notes") or data.get("notes") or "")[:1000], - } - - -def classify_special_pages( - sampled_pages: list[dict[str, Any]], - *, - vlm_observations: list[dict[str, Any]] | None = None, - model: str | None = None, - use_llm: bool = True, -) -> dict[str, Any]: - """Classify special pages with LLM, falling back to deterministic heuristics.""" - if not use_llm: - return heuristic_classify_special_pages(sampled_pages) - - try: - from shared.core.config import settings - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL - client = get_openai_client(model=effective_model) - payload = { - "sampled_pages": sampled_pages, - "vlm_observations": vlm_observations or [], - } - response = client.chat_completion( - messages=[ - {"role": "system", "content": PROMPT}, - { - "role": "user", - "content": "Classify these pages as json:\n" - + json.dumps(payload, ensure_ascii=False), - }, - ], - model=effective_model, - temperature=0.0, - max_tokens=1800, - response_format={"type": "json_object"}, - ) - return _normalize_llm_pages(extract_json_object(response), sampled_pages) - except Exception as exc: - logger.warning( - f"[document_agent.classify_special_pages] LLM classification failed, " - f"using heuristics: {exc}" - ) - return heuristic_classify_special_pages(sampled_pages) diff --git a/apps/worker/app/services/document_agent/tools/find_h1_boundaries.py b/apps/worker/app/services/document_agent/tools/find_h1_boundaries.py index e730b621d..3098fd4ed 100644 --- a/apps/worker/app/services/document_agent/tools/find_h1_boundaries.py +++ b/apps/worker/app/services/document_agent/tools/find_h1_boundaries.py @@ -1,352 +1,119 @@ -"""find_h1_boundaries — locate level-1 headings in a PDF via text search. - -Why not use PyMuPDF ``doc.get_toc()`` page numbers? ---------------------------------------------------- -PDF bookmark page references encode physical page offsets, while printed page -numbers in a TOC reflect logical numbering that often includes unnumbered front -matter. For Chinese documents and scanned PDFs the offsets frequently disagree -by several pages or are entirely absent. - -This tool instead: -1. Identifies TOC pages by detecting TOC-marker text in page features already - gathered by ``scan_all_page_features``. -2. Reads those TOC pages in full to extract level-1 heading candidate strings. -3. Searches every page's full text for those candidates (fuzzy, after - normalisation) to find where they actually appear in the body. - -The result is a set of ``H1Match`` records linking each heading title to the -page where it physically starts — a reliable basis for shard cut decisions. -""" +"""Find H1 starts from TOC candidates or heading-like body lines.""" from __future__ import annotations -import gc -import re -import unicodedata +import time from typing import Any -from app.services.document_agent.page_map import H1BoundaryResult, H1Match, PageFeature -from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( - run_in_child_process, - worker, +from app.services.document_agent.heading_text import ( + candidate_allowed, + clean_toc_line, + fuzzy_match, + looks_like_h1_line, + normalize_heading, ) -from loguru import logger - - -# ── Text normalisation helpers ───────────────────────────────────────────────── - -# Patterns that prefix chapter/section numbers in various languages -_LEADING_NUMBER_RE = re.compile( - r"""^ - (?: - 第\s*[零一二三四五六七八九十百千\d]+\s*[章节篇部分] # 第X章/节 - | [零一二三四五六七八九十百千]+\s*[、。，,] # 一、 - | [（(]\s*[零一二三四五六七八九十百千\d]+\s*[）)] # （一） - | \d+(?:\.\d+)*\.?\s* # 1. / 1.2 / 1.2.3. - | [IVXLCDM]+\.?\s* # Roman I. II. - | [A-Za-z]\.\s* # A. B. - | Chapter\s+\w+\s* # Chapter N - ) - """, - re.VERBOSE | re.IGNORECASE, -) - - -def _normalise_heading(text: str) -> str: - """Strip leading numbers/labels and normalise whitespace for matching.""" - text = unicodedata.normalize("NFKC", text) - text = re.sub(r"\s+", " ", text).strip() - # Strip leading numbering patterns - stripped = _LEADING_NUMBER_RE.sub("", text).strip() - # Keep the original if stripping removed everything (guard against over-stripping) - return stripped if stripped else text - - -def _fuzzy_contains(needle: str, haystack: str, min_len: int = 4) -> bool: - """Check whether ``needle`` (normalised heading) appears in ``haystack``. - - Matching strategy (in order of strictness): - 1. Exact substring after normalisation. - 2. Stripped-number variant of needle appears in normalised haystack. - """ - if not needle or len(needle) < min_len: - return False - norm_hay = re.sub(r"\s+", " ", unicodedata.normalize("NFKC", haystack)) - if needle in norm_hay: - return True - stripped = _normalise_heading(needle) - if stripped and len(stripped) >= min_len and stripped in norm_hay: - return True - return False - - -# ── TOC page detection ───────────────────────────────────────────────────────── - -_TOC_MARKERS = frozenset(["目录", "目次", "contents", "tableofcontents"]) - - -def _is_toc_page(feature: PageFeature) -> bool: - """Heuristic: is this page a Table of Contents page?""" - text = re.sub(r"\s+", "", feature.text_preview.lower()) - return any(marker in text for marker in _TOC_MARKERS) - - -# ── Child-process worker: read full page texts ───────────────────────────────── - - -@worker -def _read_page_texts_worker( - queue, - pdf_path: str, - page_indices: list[int], # 0-based -) -> None: - """Read full text for the requested pages; runs in an isolated process.""" - import pymupdf # type: ignore[import] - - results: dict[int, str] = {} - try: - doc = pymupdf.open(pdf_path) - for idx in page_indices: - if 0 <= idx < doc.page_count: - try: - results[idx] = doc[idx].get_text() or "" - except Exception: - results[idx] = "" - finally: - try: - doc.close() - except Exception: - pass - gc.collect() +from app.services.document_agent.manifest import H1BoundaryResult, H1Candidate, TocCandidate, ToolContext, ToolResult +from app.services.document_agent.pdf_text import meaningful_lines, read_page_texts +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState + + +def _all_non_toc_pages(ctx: ToolContext) -> list[int]: + toc_pages = set(ctx.blackboard.toc_result.toc_pages if ctx.blackboard.toc_result else []) + return [ + feature.page + for feature in ctx.blackboard.page_features + if feature.page not in toc_pages + ] - queue.put({"ok": True, "texts": results}) +def _fallback_candidates_from_previews(ctx: ToolContext) -> list[TocCandidate]: + candidates: list[TocCandidate] = [] + seen: set[str] = set() + for feature in ctx.blackboard.page_features: + for line_index, line in enumerate(feature.text_lines_preview): + cleaned = clean_toc_line(line) + if not looks_like_h1_line(cleaned) or not candidate_allowed(cleaned): + continue + normalized = normalize_heading(cleaned) + if normalized in seen: + continue + seen.add(normalized) + candidates.append( + TocCandidate( + title=cleaned, + normalized_title=normalized, + source_page=feature.page, + line_index=line_index, + ) + ) + return candidates -def _load_full_page_texts( - pdf_path: str, - page_indices: list[int], - timeout: int = 120, -) -> dict[int, str]: - """Return {0-based-index: full_text} for the requested pages.""" - if not page_indices: - return {} - try: - result = run_in_child_process( - _read_page_texts_worker, pdf_path, page_indices, timeout=timeout - ) - return {int(k): str(v) for k, v in (result.get("texts") or {}).items()} - except Exception as exc: - logger.warning(f"[find_h1_boundaries] full-text load failed: {exc}") - return {} +def _match_candidate( + candidate: TocCandidate, + page_texts: dict[int, str], +) -> H1Candidate | None: + for page, text in sorted(page_texts.items()): + lines = meaningful_lines(text) + for line_index, line in enumerate(lines[:20]): + if candidate.title in line: + return H1Candidate( + title=candidate.title, + page=page, + confidence=1.0, + matched_line=line, + source="toc_exact_top", + evidence={"line_index": line_index, "toc_page": candidate.source_page}, + ) + if fuzzy_match(candidate.normalized_title, line): + return H1Candidate( + title=candidate.title, + page=page, + confidence=0.86, + matched_line=line, + source="toc_fuzzy_top", + evidence={"line_index": line_index, "toc_page": candidate.source_page}, + ) + return None -# ── TOC text → h1 candidate extraction ──────────────────────────────────────── -# Patterns that suggest a TOC line is a level-1 heading -# (numbered first-level or occupies a prominent position in the TOC) -_H1_TOC_LINE_RE = re.compile( - r""" - ^\s* - (?: - 第\s*[零一二三四五六七八九十百千\d]+\s*[章篇部] # 第X章 - | [零一二三四五六七八九十百千]+\s*[、。] # 一、 - | \d+\s*[\.\s] # 1. / 1 (level-1 only) - | Chapter\s+\w+ # Chapter … - | [IVXLCDM]+\.?\s+\w # Roman numeral - ) - """, - re.VERBOSE | re.IGNORECASE, +@register_tool( + name="find.h1_boundaries", + description="Match H1 candidates to body page starts with evidence and confidence.", + allowed_states={DocumentAgentState.CLASSIFIED}, ) - -# A line is clearly a sub-heading if it starts with at least two level numbers -_SUBHEADING_RE = re.compile(r"^\s*\d+\.\d+", re.IGNORECASE) - - -def _extract_h1_candidates_from_toc_text(toc_text: str) -> list[str]: - """Parse raw TOC page text and return likely level-1 heading strings.""" - candidates: list[str] = [] - seen: set[str] = set() - - for raw_line in toc_text.splitlines(): - line = raw_line.strip() - if not line: - continue - # Skip sub-headings (e.g. "1.2 something") - if _SUBHEADING_RE.match(line): - continue - # Skip lines that look like page numbers only - if re.fullmatch(r"[\d\s\.\-·…]+", line): - continue - # Must match a level-1 pattern - if not _H1_TOC_LINE_RE.match(line): +def find_h1_boundaries(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + toc_result = ctx.blackboard.toc_result + candidates = list(toc_result.candidates) if toc_result and toc_result.candidates else [] + method: str = "toc_grep" if candidates else "heading_grep" + if not candidates: + candidates = _fallback_candidates_from_previews(ctx) + pages = _all_non_toc_pages(ctx) + page_texts = read_page_texts(ctx.pdf_path, pages) if pages else {} + matches: list[H1Candidate] = [] + seen_pages_titles: set[tuple[int, str]] = set() + for candidate in candidates: + match = _match_candidate(candidate, page_texts) + if match is None: continue - # Strip trailing page-number suffix common in TOCs: " ……… 12" - cleaned = re.sub(r"[\s\.\-·…]+\d+\s*$", "", line).strip() - if not cleaned: + key = (match.page, normalize_heading(match.title)) + if key in seen_pages_titles: continue - norm = _normalise_heading(cleaned) - if norm and norm not in seen and len(norm) >= 2: - candidates.append(cleaned) # keep original for matching fidelity - seen.add(norm) - - return candidates - - -# ── Public API ───────────────────────────────────────────────────────────────── - - -def find_h1_boundaries( - pdf_path: str, - page_features: list[PageFeature], - *, - timeout: int = 180, -) -> H1BoundaryResult: - """Locate level-1 headings in a PDF by grepping full page texts. - - Algorithm - --------- - 1. Identify TOC pages from already-collected ``page_features``. - 2. Read those TOC pages in full; extract level-1 heading candidates. - 3. If no TOC found, fall back to grepping markdown-style headings - (``# Title`` or lines that match level-1 numbering patterns) from - each page's text_preview across all pages. - 4. For each candidate heading, search all pages' full text for the - heading text (after normalisation). Record the first matching page. - - Returns - ------- - ``H1BoundaryResult`` — always non-raising; ``method="none"`` when - nothing useful was found. - """ - if not page_features: - return H1BoundaryResult(toc_pages=[], h1_matches=[], method="none", - notes="no page features provided") - - # Step 1: find TOC pages - toc_pages = [pf.page for pf in page_features if _is_toc_page(pf)] - logger.info(f"[find_h1_boundaries] TOC pages (heuristic): {toc_pages}") - - # Step 2: extract h1 candidates from TOC page full text - h1_candidates: list[str] = [] - if toc_pages: - toc_indices = [p - 1 for p in toc_pages] # 0-based - toc_texts = _load_full_page_texts(pdf_path, toc_indices, timeout=timeout) - for idx, text in sorted(toc_texts.items()): - page_candidates = _extract_h1_candidates_from_toc_text(text) - logger.info( - f"[find_h1_boundaries] TOC page {idx + 1}: " - f"{len(page_candidates)} h1 candidates" - ) - h1_candidates.extend(page_candidates) - - # Step 3: fallback — grep heading-like lines from all page previews - using_fallback = not h1_candidates - if using_fallback: - logger.info( - "[find_h1_boundaries] no TOC or no candidates — " - "grepping all page text_previews for heading-like lines" - ) - for pf in page_features: - for line in pf.text_preview.splitlines(): - line = line.strip() - if _H1_TOC_LINE_RE.match(line) and not _SUBHEADING_RE.match(line): - cleaned = re.sub(r"[\s\.\-·…]+\d+\s*$", "", line).strip() - if cleaned and len(_normalise_heading(cleaned)) >= 2: - h1_candidates.append(cleaned) - - # Deduplicate candidates preserving order - seen_norms: set[str] = set() - unique_candidates: list[str] = [] - for c in h1_candidates: - norm = _normalise_heading(c) - if norm and norm not in seen_norms: - unique_candidates.append(c) - seen_norms.add(norm) - h1_candidates = unique_candidates - - if not h1_candidates: - return H1BoundaryResult( - toc_pages=toc_pages, - h1_matches=[], - method="none", - notes="no h1 candidates extracted", - ) - - logger.info( - f"[find_h1_boundaries] {len(h1_candidates)} unique h1 candidates to search" - ) - - # Step 4: search all pages for each candidate - # Load full text for all pages (excluding confirmed TOC pages to avoid - # matching the TOC entry itself instead of the body heading) - non_toc_indices = [ - pf.page - 1 for pf in page_features if pf.page not in set(toc_pages) - ] - all_texts = _load_full_page_texts(pdf_path, non_toc_indices, timeout=timeout) - - # Build a sorted list of (page_number, full_text) for ordered search - page_text_pairs: list[tuple[int, str]] = sorted( - ((idx + 1, text) for idx, text in all_texts.items()), - key=lambda x: x[0], - ) - - h1_matches: list[H1Match] = [] - for candidate in h1_candidates: - norm_candidate = _normalise_heading(candidate) - matched_page: int | None = None - match_text: str = "" - confidence: float = 0.0 - - for page_num, full_text in page_text_pairs: - # Try exact match first (higher confidence) - if candidate in full_text: - matched_page = page_num - match_text = candidate - confidence = 1.0 - break - # Try normalised match - if norm_candidate and _fuzzy_contains(norm_candidate, full_text): - matched_page = page_num - match_text = norm_candidate - confidence = 0.85 - break - - if matched_page is not None: - h1_matches.append( - H1Match( - title=candidate, - page=matched_page, - confidence=confidence, - match_text=match_text, - ) - ) - logger.debug( - f"[find_h1_boundaries] '{candidate[:40]}' → page {matched_page} " - f"(conf={confidence})" - ) - else: - logger.debug( - f"[find_h1_boundaries] '{candidate[:40]}' → not found in body" - ) - - method: str - if h1_matches and not using_fallback: - method = "toc_grep" - elif h1_matches: - method = "heading_grep" - else: + seen_pages_titles.add(key) + matches.append(match) + if not matches: method = "none" - - logger.info( - f"[find_h1_boundaries] result: method={method}, " - f"{len(h1_matches)}/{len(h1_candidates)} headings matched" + result = H1BoundaryResult( + h1_candidates=matches, + method=method, # type: ignore[arg-type] + notes=f"{len(matches)} of {len(candidates)} candidates matched", ) - return H1BoundaryResult( - toc_pages=toc_pages, - h1_matches=h1_matches, - method=method, - notes=( - f"{len(h1_matches)} of {len(h1_candidates)} candidates matched; " - f"fallback={using_fallback}" - ), + ctx.blackboard.h1_result = result + ctx.blackboard.global_signals["h1_candidate_count"] = len(matches) + return ToolResult( + status="ok", + payload={"h1_match_count": len(matches), "method": method}, + latency_ms=int((time.monotonic() - start) * 1000), ) diff --git a/apps/worker/app/services/document_agent/tools/find_toc_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_pages.py new file mode 100644 index 000000000..e5a3c8fc3 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/find_toc_pages.py @@ -0,0 +1,82 @@ +"""TOC page and level-1 candidate extraction.""" + +from __future__ import annotations + +import time +from typing import Any + +from app.services.document_agent.heading_text import ( + candidate_allowed, + clean_toc_line, + normalize_heading, + looks_like_h1_line, +) +from app.services.document_agent.manifest import TocCandidate, TocResult, ToolContext, ToolResult +from app.services.document_agent.pdf_text import read_page_texts +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState + + +def _is_toc_feature_lines(lines: list[str]) -> bool: + compact = "".join(lines).replace(" ", "").lower() + return any(marker in compact for marker in ("目录", "目次", "contents", "tableofcontents")) + + +def _extract_candidates(page: int, text: str) -> list[TocCandidate]: + candidates: list[TocCandidate] = [] + seen: set[str] = set() + for idx, raw_line in enumerate(text.splitlines()): + line = clean_toc_line(raw_line) + if not looks_like_h1_line(line): + continue + if not candidate_allowed(line): + continue + normalized = normalize_heading(line) + if normalized in seen: + continue + seen.add(normalized) + candidates.append( + TocCandidate( + title=line, + normalized_title=normalized, + source_page=page, + line_index=idx, + numbering=line[: max(len(line) - len(normalized), 0)].strip(), + ) + ) + return candidates + + +@register_tool( + name="find.toc_pages", + description="Find table-of-contents pages and extract level-1 heading candidates.", + allowed_states={DocumentAgentState.PROBED}, +) +def find_toc_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + toc_pages = [ + feature.page + for feature in ctx.blackboard.page_features + if _is_toc_feature_lines(feature.text_lines_preview) + ] + texts = read_page_texts(ctx.pdf_path, toc_pages) if toc_pages else {} + candidates: list[TocCandidate] = [] + for page, text in texts.items(): + candidates.extend(_extract_candidates(page, text)) + result = TocResult( + toc_pages=sorted(toc_pages), + candidates=candidates, + method="toc_marker" if toc_pages else "none", + notes=f"{len(candidates)} h1 candidates extracted", + ) + ctx.blackboard.toc_result = result + ctx.blackboard.global_signals["toc_page_count"] = len(toc_pages) + ctx.blackboard.global_signals["toc_candidate_count"] = len(candidates) + return ToolResult( + status="ok", + payload={ + "toc_pages": sorted(toc_pages), + "candidate_count": len(candidates), + }, + latency_ms=int((time.monotonic() - start) * 1000), + ) diff --git a/apps/worker/app/services/document_agent/tools/llm_json.py b/apps/worker/app/services/document_agent/tools/llm_json.py deleted file mode 100644 index 9531e3690..000000000 --- a/apps/worker/app/services/document_agent/tools/llm_json.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Small JSON helpers for split-agent LLM tools.""" - -from __future__ import annotations - -import json -import re -from typing import Any - - -def extract_json_object(text: str) -> dict[str, Any]: - """Parse JSON from a model response that may include light prose/fences.""" - raw = (text or "").strip() - if not raw: - raise ValueError("empty JSON response") - if raw.startswith("```"): - raw = re.sub(r"^```(?:json)?\s*", "", raw, flags=re.IGNORECASE) - raw = re.sub(r"\s*```$", "", raw) - try: - data = json.loads(raw) - except json.JSONDecodeError: - start = raw.find("{") - end = raw.rfind("}") - if start < 0 or end <= start: - raise - data = json.loads(raw[start : end + 1]) - if not isinstance(data, dict): - raise ValueError("expected JSON object") - return data diff --git a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py new file mode 100644 index 000000000..017504d2c --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py @@ -0,0 +1,72 @@ +"""Persist anatomy map artifacts and optional database records.""" + +from __future__ import annotations + +import json +import os +import time +from pathlib import Path +from typing import Any + +from app.services.document_agent.manifest import PageAnatomyMap, ToolContext, ToolResult +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState + + +def _artifact_dir(ctx: ToolContext) -> Path: + if ctx.output_dir: + return Path(ctx.output_dir) + base = Path(os.path.expanduser("~/.knowhere/_debug_profile")) + return base / Path(ctx.pdf_path).stem + + +def build_anatomy_map(ctx: ToolContext) -> PageAnatomyMap: + if not ( + ctx.blackboard.toc_result + and ctx.blackboard.h1_result + and ctx.blackboard.hierarchy_assist + and ctx.blackboard.shard_plan + ): + raise ValueError("cannot build anatomy map from incomplete blackboard") + return PageAnatomyMap( + job_id=ctx.job_id, + file_path=ctx.pdf_path, + page_count=ctx.blackboard.page_count, + page_features=ctx.blackboard.page_features, + page_labels=ctx.blackboard.page_labels, + toc_result=ctx.blackboard.toc_result, + h1_result=ctx.blackboard.h1_result, + hierarchy_assist=ctx.blackboard.hierarchy_assist, + shard_plan=ctx.blackboard.shard_plan, + page_processing_plan=None, + global_signals=ctx.blackboard.global_signals, + trace_summary={ + "budget": ctx.budget.snapshot(), + "state_trace": list(ctx.blackboard.state_trace), + "validation": ctx.blackboard.validation_report, + }, + ) + + +@register_tool( + name="persist.anatomy_map", + description="Persist anatomy_map.json and buffer SQL trace payloads.", + allowed_states={DocumentAgentState.VALIDATED}, +) +def persist_anatomy_map(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + anatomy = build_anatomy_map(ctx) + output_dir = _artifact_dir(ctx) + output_dir.mkdir(parents=True, exist_ok=True) + artifact_path = output_dir / "anatomy_map.json" + artifact_path.write_text( + json.dumps(anatomy.to_dict(), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + if ctx.trace: + ctx.trace.set_anatomy_map(anatomy, str(artifact_path)) + return ToolResult( + status="ok", + payload={"artifact_path": str(artifact_path)}, + latency_ms=int((time.monotonic() - start) * 1000), + ) 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 new file mode 100644 index 000000000..6eac30a0d --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/probe_page_features.py @@ -0,0 +1,147 @@ +"""Full-page structural probing.""" + +from __future__ import annotations + +import gc +import time +from typing import Any + +from app.services.document_agent.manifest import PageFeature, ToolContext, ToolResult +from app.services.document_agent.pdf_text import top_lines +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) +from loguru import logger + + +def _rect_area(rect: Any) -> float: + width = max(float(getattr(rect, "width", 0.0) or 0.0), 0.0) + height = max(float(getattr(rect, "height", 0.0) or 0.0), 0.0) + return width * height + + +def _image_coverage(page: Any, page_area: float) -> tuple[float, int]: + if page_area <= 0: + return 0.0, 0 + area = 0.0 + images = page.get_images(full=True) or [] + seen: set[tuple[float, float, float, float]] = set() + for image in images: + xref = image[0] + try: + rects = page.get_image_rects(xref) or [] + except Exception: + rects = [] + for rect in rects: + key = ( + round(float(getattr(rect, "x0", 0.0) or 0.0), 2), + round(float(getattr(rect, "y0", 0.0) or 0.0), 2), + round(float(getattr(rect, "x1", 0.0) or 0.0), 2), + round(float(getattr(rect, "y1", 0.0) or 0.0), 2), + ) + if key in seen: + continue + seen.add(key) + area += _rect_area(rect) + return min(area / page_area, 1.0), len(images) + + +def _table_count(page: Any) -> int: + try: + finder = page.find_tables() + return len(getattr(finder, "tables", []) or []) + except Exception: + return 0 + + +def _probe_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()) + image_coverage, image_count = _image_coverage(page, area) + try: + drawings_count = len(page.get_drawings() or []) + except Exception: + drawings_count = 0 + orientation = "landscape" if float(rect.width) > float(rect.height) else "portrait" + return { + "page": page_number, + "raw_text_length": raw_text_length, + "text_density": round(raw_text_length / area * 10000, 4), + "image_coverage": round(image_coverage, 4), + "image_count": image_count, + "table_count": _table_count(page), + "drawings_count": drawings_count, + "orientation": orientation, + "width": round(float(rect.width), 2), + "height": round(float(rect.height), 2), + "is_blank_like": raw_text_length < 20 and image_coverage < 0.02 and drawings_count < 5, + "text_lines_preview": top_lines(text, max_lines=30), + } + + +@worker +def _probe_worker(queue, pdf_path: str) -> None: + import pymupdf # type: ignore[import] + + features: list[dict[str, Any]] = [] + page_count = 0 + try: + doc = pymupdf.open(pdf_path) + page_count = int(doc.page_count) + for idx in range(page_count): + features.append(_probe_one(doc[idx], idx + 1)) + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "page_count": page_count, "features": features}) + + +@register_tool( + name="probe.page_features", + description="Probe every PDF page for structural signals without parsing content semantically.", + allowed_states={DocumentAgentState.INIT}, +) +def probe_page_features(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + try: + result = run_in_child_process(_probe_worker, ctx.pdf_path, timeout=300) + features = [ + PageFeature( + page=int(item["page"]), + raw_text_length=int(item.get("raw_text_length") or 0), + text_density=float(item.get("text_density") or 0.0), + image_coverage=float(item.get("image_coverage") or 0.0), + image_count=int(item.get("image_count") or 0), + table_count=int(item.get("table_count") or 0), + drawings_count=int(item.get("drawings_count") or 0), + orientation=str(item.get("orientation") or "portrait"), # type: ignore[arg-type] + width=float(item.get("width") or 0.0), + height=float(item.get("height") or 0.0), + is_blank_like=bool(item.get("is_blank_like")), + text_lines_preview=list(item.get("text_lines_preview") or []), + ) + for item in (result.get("features") or []) + ] + ctx.blackboard.page_features = sorted(features, key=lambda f: f.page) + ctx.blackboard.page_count = int(result.get("page_count") or len(features)) + ctx.blackboard.global_signals["total_pages"] = ctx.blackboard.page_count + logger.info("[document_agent] probed {} pages", ctx.blackboard.page_count) + return ToolResult( + status="ok", + payload={"page_count": ctx.blackboard.page_count}, + latency_ms=int((time.monotonic() - start) * 1000), + ) + except Exception as exc: + return ToolResult( + status="error", + error=str(exc), + latency_ms=int((time.monotonic() - start) * 1000), + ) diff --git a/apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py b/apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py new file mode 100644 index 000000000..7f11331fb --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py @@ -0,0 +1,150 @@ +"""Build hierarchy hints from page anatomy signals.""" + +from __future__ import annotations + +import json +import time +from typing import Any + +from app.services.document_agent.manifest import ( + BoundaryHint, + H1Candidate, + HierarchyAssistPlan, + PageLabel, + ToolContext, + ToolResult, +) +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState +from app.services.document_agent.validators import repair_hierarchy_assist +from shared.utils.token_estimate import estimate_tokens + + +def _rule_based_plan(labels: list[PageLabel], h1_candidates: list[H1Candidate]) -> HierarchyAssistPlan: + exclude = [ + label.page + for label in labels + if label.kind in {"toc", "blank", "separator", "landscape"} + ] + suppress = [ + label.page + for label in labels + if label.kind in {"table_heavy", "single_image", "scan_like", "image_heavy"} + ] + boundary_hints = [ + BoundaryHint( + page=candidate.page, + anchor_type="h1_boundary", + confidence=candidate.confidence, + evidence=candidate.evidence | {"title": candidate.title}, + ) + for candidate in h1_candidates + ] + scan_like_count = sum(1 for label in labels if label.kind in {"scan_like", "single_image"}) + recommendation = "aggressive" if scan_like_count > max(len(labels) // 3, 0) else "normal" + if h1_candidates and scan_like_count == 0: + recommendation = "normal" + if not h1_candidates and scan_like_count == 0: + recommendation = "aggressive" + return HierarchyAssistPlan( + exclude_pages_from_title_candidates=sorted(set(exclude)), + prefer_h1_start_pages=sorted(h1_candidates, key=lambda item: item.page), + suppress_title_pages=sorted(set(suppress)), + section_boundary_hints=boundary_hints, + smart_parse_recommendation=recommendation, # type: ignore[arg-type] + rationale="Derived from page labels and H1 boundary evidence.", + ) + + +def _parse_llm_plan(raw: str, fallback: HierarchyAssistPlan) -> HierarchyAssistPlan: + try: + data = json.loads(raw) + except json.JSONDecodeError: + return fallback + if not isinstance(data, dict): + return fallback + plan = fallback + plan.exclude_pages_from_title_candidates = [ + int(page) for page in data.get("exclude_pages_from_title_candidates", plan.exclude_pages_from_title_candidates) + ] + plan.suppress_title_pages = [ + int(page) for page in data.get("suppress_title_pages", plan.suppress_title_pages) + ] + recommendation = data.get("smart_parse_recommendation") + if recommendation in {"off", "normal", "aggressive"}: + plan.smart_parse_recommendation = recommendation + if isinstance(data.get("rationale"), str): + plan.rationale = data["rationale"] + return plan + + +@register_tool( + name="propose.hierarchy_assist", + description="Produce hierarchy hints used by section skeleton extraction and title parsing.", + allowed_states={DocumentAgentState.H1_FOUND}, +) +def propose_hierarchy_assist(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + h1_result = ctx.blackboard.h1_result + labels = ctx.blackboard.page_labels + fallback = _rule_based_plan(labels, h1_result.h1_candidates if h1_result else []) + toc_pages = set(ctx.blackboard.toc_result.toc_pages if ctx.blackboard.toc_result else []) + + model = ctx.settings.get("model") + if model: + payload = { + "page_count": ctx.blackboard.page_count, + "page_kind_counts": ctx.blackboard.global_signals.get("page_kind_counts", {}), + "toc_pages": sorted(toc_pages), + "h1_candidates": [ + candidate.to_dict() for candidate in (h1_result.h1_candidates if h1_result else []) + ], + "candidate_exclude_pages": fallback.exclude_pages_from_title_candidates, + "candidate_suppress_pages": fallback.suppress_title_pages, + } + prompt = ( + "Return strict JSON for hierarchy assistance. " + "Use only page numbers present in the payload. " + "Do not invent headings. " + "Fields: exclude_pages_from_title_candidates, suppress_title_pages, " + "smart_parse_recommendation, rationale.\n" + + json.dumps(payload, ensure_ascii=False) + ) + est = estimate_tokens(prompt) + if ctx.budget.try_reserve("plan", est): + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=prompt, + model=model, + temperature=0.0, + max_tokens=1200, + response_format={"type": "json_object"}, + ) + ctx.budget.commit("plan", actual=usage.get("total_tokens", est), est=est) + plan = _parse_llm_plan(raw, fallback) + except Exception: + ctx.budget.refund("plan", est=est) + plan = fallback + else: + plan = fallback + else: + plan = fallback + + repaired = repair_hierarchy_assist( + plan, + page_count=ctx.blackboard.page_count, + toc_pages=toc_pages, + ) + ctx.blackboard.hierarchy_assist = repaired + return ToolResult( + status="ok", + payload={ + "exclude_pages": repaired.exclude_pages_from_title_candidates, + "suppress_pages": repaired.suppress_title_pages, + "h1_hint_count": len(repaired.prefer_h1_start_pages), + }, + latency_ms=int((time.monotonic() - start) * 1000), + ) 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 new file mode 100644 index 000000000..b3a158405 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -0,0 +1,176 @@ +"""Rule-based long-PDF shard planning.""" + +from __future__ import annotations + +import os +import time +from typing import Any + +from app.services.document_agent.manifest import Shard, ShardPlan, ToolContext, ToolResult +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState +from app.services.document_agent.validators import single_shard_plan, validate_shard_plan + + +def _thresholds(ctx: ToolContext) -> tuple[int, int, int]: + threshold = int( + ctx.settings.get("shard_threshold") + or os.environ.get("PARSE_AGENT_SHARD_THRESHOLD", "200") + ) + min_pages = int( + ctx.settings.get("min_pages_per_shard") + or os.environ.get("PARSE_AGENT_MIN_PAGES_PER_SHARD", "20") + ) + max_pages = int( + ctx.settings.get("max_pages_per_shard") + or os.environ.get("PARSE_AGENT_MAX_PAGES_PER_SHARD", "200") + ) + return threshold, min_pages, max_pages + + +def _avoid_pages(ctx: ToolContext) -> set[int]: + return { + label.page + for label in ctx.blackboard.page_labels + if label.kind in {"table_heavy", "landscape"} + } + + +def _separator_pages(ctx: ToolContext) -> set[int]: + return { + label.page + for label in ctx.blackboard.page_labels + if label.kind in {"blank", "separator", "sparse"} + } + + +def _nearest_safe_cut(target: int, previous: int, page_count: int, avoid: set[int], separators: set[int]) -> tuple[int, str]: + window = range(target, max(previous, target - 12), -1) + for page in window: + if previous < page < page_count and page in separators and page not in avoid: + return page, "blank_separator" + for page in range(target, max(previous, target - 5), -1): + if previous < page < page_count and page not in avoid: + return page, "forced_max_size" + return max(previous + 1, min(target, page_count - 1)), "forced_max_size" + + +def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> list[Shard]: + shards: list[Shard] = [] + previous = 0 + for cut_page, anchor_type, evidence, confidence in cuts: + if cut_page <= previous: + continue + shards.append( + Shard( + shard_index=len(shards), + page_start=previous + 1, + page_end=cut_page, + page_offset=previous, + anchor_type=anchor_type, # type: ignore[arg-type] + anchor_evidence=evidence, + confidence=confidence, + ) + ) + previous = cut_page + if previous < page_count: + shards.append( + Shard( + shard_index=len(shards), + page_start=previous + 1, + page_end=page_count, + page_offset=previous, + anchor_type="forced_max_size", + anchor_evidence="final shard", + confidence=1.0, + ) + ) + return shards + + +@register_tool( + name="propose.shard_plan", + description="Create a rule-based PDF segment plan for long-document PDF-to-MD execution.", + allowed_states={DocumentAgentState.H1_FOUND}, +) +def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + page_count = ctx.blackboard.page_count + threshold, min_pages, max_pages = _thresholds(ctx) + if page_count <= threshold: + plan = single_shard_plan(page_count) + ctx.blackboard.shard_plan = plan + return ToolResult( + status="ok", + payload={"enabled": False, "shard_count": len(plan.shards)}, + latency_ms=int((time.monotonic() - start) * 1000), + ) + + avoid = _avoid_pages(ctx) + separators = _separator_pages(ctx) + h1_candidates = [] + if ctx.blackboard.hierarchy_assist: + h1_candidates = ctx.blackboard.hierarchy_assist.prefer_h1_start_pages + elif ctx.blackboard.h1_result: + h1_candidates = ctx.blackboard.h1_result.h1_candidates + + cuts: list[tuple[int, str, str, float]] = [] + previous = 0 + for candidate in sorted(h1_candidates, key=lambda item: item.page): + cut_page = candidate.page - 1 + if cut_page <= previous: + continue + if cut_page - previous < min_pages: + continue + while cut_page - previous > max_pages: + forced_target = previous + max_pages + safe_cut, anchor_type = _nearest_safe_cut( + forced_target, + previous, + page_count, + avoid, + separators, + ) + cuts.append((safe_cut, anchor_type, "max shard size guard", 0.58)) + previous = safe_cut + if cut_page > previous and cut_page - previous >= min_pages: + cuts.append( + ( + cut_page, + "h1_boundary", + f"h1 starts at page {candidate.page}: {candidate.title}", + candidate.confidence, + ) + ) + previous = cut_page + + while page_count - previous > max_pages: + target = previous + max_pages + safe_cut, anchor_type = _nearest_safe_cut(target, previous, page_count, avoid, separators) + cuts.append((safe_cut, anchor_type, "tail max shard size guard", 0.58)) + previous = safe_cut + + shards = _cuts_to_shards(cuts, page_count) + reason = "hierarchy_isolation" if any(cut[1] == "h1_boundary" for cut in cuts) else "too_large" + plan = ShardPlan( + enabled=True, + reason=reason, # type: ignore[arg-type] + shards=shards, + validation=validate_shard_plan( + ShardPlan(enabled=True, reason=reason, shards=shards), # type: ignore[arg-type] + page_count=page_count, + min_pages=min_pages, + max_pages=max_pages, + ), + ) + ctx.blackboard.shard_plan = plan + return ToolResult( + status="ok", + payload={ + "enabled": plan.enabled, + "reason": plan.reason, + "shard_count": len(plan.shards), + "valid": plan.validation.valid, + }, + latency_ms=int((time.monotonic() - start) * 1000), + ) diff --git a/apps/worker/app/services/document_agent/tools/scan_all_page_features.py b/apps/worker/app/services/document_agent/tools/scan_all_page_features.py deleted file mode 100644 index bd3799704..000000000 --- a/apps/worker/app/services/document_agent/tools/scan_all_page_features.py +++ /dev/null @@ -1,240 +0,0 @@ -"""scan_all_page_features — full-page structural feature extraction. - -Performs a **full-page traversal** (not sampling) of a PDF, extracting -structural features for every page. Runs inside an isolated PyMuPDF child -process to ensure memory is freed after extraction. - -For a 200-page A4 PDF the child process typically completes in < 8 s. -""" - -from __future__ import annotations - -import gc -import statistics -from typing import Any - -from app.services.document_agent.page_map import PageFeature -from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( - run_in_child_process, - worker, -) -from loguru import logger - - -# ── Low-level helpers ───────────────────────────────────────────────────────── - - -def _rect_area(rect: Any) -> float: - w = max(float(getattr(rect, "width", 0.0) or 0.0), 0.0) - h = max(float(getattr(rect, "height", 0.0) or 0.0), 0.0) - return w * h - - -def _measure_image_coverage(page: Any, page_area: float) -> tuple[float, int]: - if page_area <= 0: - return 0.0, 0 - image_area = 0.0 - images = page.get_images(full=True) or [] - seen_rects: set[tuple] = set() - for image in images: - if not image: - continue - xref = image[0] - try: - rects = page.get_image_rects(xref) or [] - except Exception: - rects = [] - for rect in rects: - key = ( - round(float(getattr(rect, "x0", 0.0) or 0.0), 2), - round(float(getattr(rect, "y0", 0.0) or 0.0), 2), - round(float(getattr(rect, "x1", 0.0) or 0.0), 2), - round(float(getattr(rect, "y1", 0.0) or 0.0), 2), - ) - if key in seen_rects: - continue - seen_rects.add(key) - image_area += _rect_area(rect) - return min(image_area / page_area, 1.0), len(images) - - -def _table_count(page: Any) -> int: - try: - finder = page.find_tables() - return len(getattr(finder, "tables", []) or []) - except Exception: - return 0 - - -def _extract_single_page_feature(page: Any, page_index: int) -> dict[str, Any]: - """Extract features for one page; returns a plain dict for IPC.""" - rect = page.rect - page_area = max(_rect_area(rect), 1.0) - text = page.get_text() or "" - text_len = len(text.strip()) - image_coverage, image_count = _measure_image_coverage(page, page_area) - try: - drawings_count = len(page.get_drawings() or []) - except Exception: - drawings_count = 0 - orientation = "landscape" if float(rect.width) > float(rect.height) else "portrait" - text_density = text_len / page_area * 10000 - table_count = _table_count(page) - is_blank_like = text_len < 20 and image_coverage < 0.02 and drawings_count < 5 - - return { - "page": page_index + 1, # 1-based - "text_length": text_len, - "text_density": round(text_density, 4), - "image_coverage": round(image_coverage, 4), - "image_count": image_count, - "table_count": table_count, - "drawings_count": drawings_count, - "orientation": orientation, - "width": round(float(rect.width), 2), - "height": round(float(rect.height), 2), - "is_blank_like": is_blank_like, - "text_preview": " ".join(text.split())[:300], - } - - -# ── Child-process worker ────────────────────────────────────────────────────── - - -@worker -def _scan_all_pages_worker( - queue, - pdf_path: str, - page_start: int, # 0-based inclusive - page_end: int, # 0-based inclusive (-1 = all) -) -> None: - """Scan all pages (or a subrange) and put feature list onto the queue.""" - import pymupdf # type: ignore[import] - - features: list[dict[str, Any]] = [] - page_count = 0 - - try: - doc = pymupdf.open(pdf_path) - page_count = int(doc.page_count) - end = page_count - 1 if page_end < 0 else min(page_end, page_count - 1) - start = max(0, page_start) - - for idx in range(start, end + 1): - try: - feat = _extract_single_page_feature(doc[idx], idx) - features.append(feat) - except Exception as exc: - # Log and continue; partial feature is better than crash - features.append( - { - "page": idx + 1, - "text_length": 0, - "text_density": 0.0, - "image_coverage": 0.0, - "image_count": 0, - "table_count": 0, - "drawings_count": 0, - "orientation": "portrait", - "width": 0.0, - "height": 0.0, - "is_blank_like": True, - "text_preview": "", - "_error": str(exc), - } - ) - finally: - try: - doc.close() - except Exception: - pass - gc.collect() - - queue.put({"ok": True, "page_count": page_count, "features": features}) - - -# ── Public API ──────────────────────────────────────────────────────────────── - - -def scan_all_page_features( - pdf_path: str, - *, - page_start: int = 0, - page_end: int = -1, - timeout: int = 300, -) -> list[PageFeature]: - """Return a ``PageFeature`` for every page in the PDF (or a subrange). - - Args: - pdf_path: Absolute path to the PDF file. - page_start: 0-based index of first page to scan (default 0 = first). - page_end: 0-based index of last page to scan, inclusive - (default -1 = last page). - timeout: Child process timeout in seconds. Allow 300 s for large docs. - - Returns: - Ordered list of ``PageFeature``, one per page, sorted by page number. - Never raises; returns empty list on subprocess failure. - """ - try: - result = run_in_child_process( - _scan_all_pages_worker, - pdf_path, - page_start, - page_end, - timeout=timeout, - ) - except Exception as exc: - logger.error(f"[scan_all_page_features] subprocess failed: {exc}") - return [] - - raw_features: list[dict] = result.get("features") or [] - page_features: list[PageFeature] = [] - - # Compute median page dimensions for landscape anomaly detection - widths = [f["width"] for f in raw_features if f.get("width", 0) > 0] - heights = [f["height"] for f in raw_features if f.get("height", 0) > 0] - median_w = statistics.median(widths) if widths else 595.0 # A4 width in pt - median_h = statistics.median(heights) if heights else 842.0 - - for raw in raw_features: - page_num = int(raw.get("page") or 0) - if page_num <= 0: - continue - - # Refine orientation: a "portrait" page that is much wider than the - # document median is flagged as a layout anomaly (landscape-rotated page - # that PyMuPDF may not detect via width > height alone). - orientation = raw.get("orientation", "portrait") - w, h = float(raw.get("width") or 0), float(raw.get("height") or 0) - if orientation == "portrait" and median_h > 0 and w > 0: - # If this page's aspect ratio deviates by > 40% vs median, flag it - doc_ratio = median_w / median_h - page_ratio = w / h if h > 0 else 1.0 - if page_ratio > doc_ratio * 1.4: - orientation = "landscape" - - page_features.append( - PageFeature( - page=page_num, - text_length=int(raw.get("text_length") or 0), - text_density=float(raw.get("text_density") or 0.0), - image_coverage=float(raw.get("image_coverage") or 0.0), - image_count=int(raw.get("image_count") or 0), - table_count=int(raw.get("table_count") or 0), - drawings_count=int(raw.get("drawings_count") or 0), - orientation=orientation, - width=w, - height=float(raw.get("height") or 0.0), - is_blank_like=bool(raw.get("is_blank_like")), - text_preview=str(raw.get("text_preview") or "")[:300], - embedding_ref=None, - ) - ) - - page_features.sort(key=lambda pf: pf.page) - logger.info( - f"[scan_all_page_features] scanned {len(page_features)} pages " - f"from '{pdf_path}'" - ) - return page_features diff --git a/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py new file mode 100644 index 000000000..fe4459a5b --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py @@ -0,0 +1,67 @@ +"""Validate the current anatomy blackboard.""" + +from __future__ import annotations + +import os +import time +from typing import Any + +from app.services.document_agent.manifest import PageAnatomyMap, ToolContext, ToolResult +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState +from app.services.document_agent.validators import validate_anatomy_map + + +def _thresholds(ctx: ToolContext) -> tuple[int, int]: + min_pages = int( + ctx.settings.get("min_pages_per_shard") + or os.environ.get("PARSE_AGENT_MIN_PAGES_PER_SHARD", "20") + ) + max_pages = int( + ctx.settings.get("max_pages_per_shard") + or os.environ.get("PARSE_AGENT_MAX_PAGES_PER_SHARD", "200") + ) + return min_pages, max_pages + + +@register_tool( + name="validate.anatomy_map", + description="Validate page anatomy, hierarchy hints, and shard coverage.", + allowed_states={DocumentAgentState.PLANNED}, +) +def validate_current_anatomy(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + if not ( + ctx.blackboard.toc_result + and ctx.blackboard.h1_result + and ctx.blackboard.hierarchy_assist + and ctx.blackboard.shard_plan + ): + return ToolResult( + status="error", + error="blackboard is missing anatomy outputs", + latency_ms=int((time.monotonic() - start) * 1000), + ) + anatomy = PageAnatomyMap( + job_id=ctx.job_id, + file_path=ctx.pdf_path, + page_count=ctx.blackboard.page_count, + page_features=ctx.blackboard.page_features, + page_labels=ctx.blackboard.page_labels, + toc_result=ctx.blackboard.toc_result, + h1_result=ctx.blackboard.h1_result, + hierarchy_assist=ctx.blackboard.hierarchy_assist, + shard_plan=ctx.blackboard.shard_plan, + global_signals=ctx.blackboard.global_signals, + trace_summary={}, + ) + min_pages, max_pages = _thresholds(ctx) + report = validate_anatomy_map(anatomy, min_pages=min_pages, max_pages=max_pages) + ctx.blackboard.validation_report = report.to_dict() + if ctx.blackboard.shard_plan: + ctx.blackboard.shard_plan.validation = report if not report.valid else ctx.blackboard.shard_plan.validation + return ToolResult( + status="ok" if report.valid else "invalid", + payload=report.to_dict(), + latency_ms=int((time.monotonic() - start) * 1000), + ) diff --git a/apps/worker/app/services/document_agent/trace.py b/apps/worker/app/services/document_agent/trace.py new file mode 100644 index 000000000..02a03e364 --- /dev/null +++ b/apps/worker/app/services/document_agent/trace.py @@ -0,0 +1,117 @@ +"""Best-effort parse-agent trace buffering and database persistence.""" + +from __future__ import annotations + +import time +from datetime import datetime +from typing import Any +from uuid import uuid4 + +from loguru import logger + +from app.services.document_agent.manifest import PageAnatomyMap, ToolResult + + +class ParseRunRecorder: + def __init__(self, *, job_id: str, db: Any | None = None) -> None: + self.run_id = f"prof_{uuid4().hex[:12]}" + self.job_id = job_id + self._db = db + self._started = time.monotonic() + self._steps: list[dict[str, Any]] = [] + self._anatomy: PageAnatomyMap | None = None + self._artifact_path: str | None = None + + def record_step( + self, + *, + round_index: int, + actor: str, + action_type: str, + result: ToolResult, + tool_name: str | None = None, + tool_args: dict[str, Any] | None = None, + ) -> None: + self._steps.append( + { + "round_index": round_index, + "actor": actor, + "action_type": action_type, + "tool_name": tool_name, + "tool_args": tool_args or {}, + "observation": { + "status": result.status, + "payload_keys": sorted(result.payload.keys()), + "error": result.error, + }, + "tokens_used": result.tokens_used, + "latency_ms": result.latency_ms, + "created_at": datetime.utcnow(), + } + ) + + def set_anatomy_map(self, anatomy: PageAnatomyMap, artifact_path: str) -> None: + self._anatomy = anatomy + self._artifact_path = artifact_path + + def summary(self) -> dict[str, Any]: + return { + "run_id": self.run_id, + "step_count": len(self._steps), + "artifact_path": self._artifact_path, + "latency_ms": int((time.monotonic() - self._started) * 1000), + } + + def flush(self, *, final_status: str, summary: dict[str, Any] | None = None) -> None: + if self._db is None: + return + try: + from shared.models.database.parse_agent import ParseRun, ParseStep + from shared.models.database.document_page_plan import DocumentPagePlan + + run = ParseRun( + run_id=self.run_id, + job_id=self.job_id, + kind="profile", + final_status=final_status, + rounds_count=max((step["round_index"] for step in self._steps), default=0) + 1, + total_tokens=sum(int(step.get("tokens_used") or 0) for step in self._steps), + total_latency_ms=int((time.monotonic() - self._started) * 1000), + summary=summary or self.summary(), + ) + self._db.add(run) + for index, step in enumerate(self._steps): + self._db.add( + ParseStep( + step_id=f"prst_{uuid4().hex[:12]}", + run_id=self.run_id, + round_index=int(step["round_index"]), + actor=str(step["actor"]), + action_type=str(step["action_type"]), + tool_name=step.get("tool_name"), + tool_args=step.get("tool_args"), + observation=step.get("observation"), + tokens_used=int(step.get("tokens_used") or 0), + latency_ms=int(step.get("latency_ms") or 0), + created_at=step.get("created_at"), + ) + ) + if self._anatomy is not None: + self._db.add( + DocumentPagePlan( + page_plan_id=f"dpp_{uuid4().hex[:12]}", + job_id=self.job_id, + page_count=self._anatomy.page_count, + hierarchy_assist=self._anatomy.hierarchy_assist.to_dict(), + shard_plan=self._anatomy.shard_plan.to_dict(), + page_processing_plan=None, + global_signals=self._anatomy.global_signals, + ) + ) + self._db.flush() + except Exception as exc: + logger.debug(f"parse agent trace flush failed: {exc}") + try: + self._db.rollback() + except Exception: + pass diff --git a/apps/worker/app/services/document_agent/validators.py b/apps/worker/app/services/document_agent/validators.py new file mode 100644 index 000000000..44e3cf537 --- /dev/null +++ b/apps/worker/app/services/document_agent/validators.py @@ -0,0 +1,121 @@ +"""Validation and repair for page anatomy outputs.""" + +from __future__ import annotations + +from app.services.document_agent.manifest import ( + HierarchyAssistPlan, + PageAnatomyMap, + Shard, + ShardPlan, + ValidationReport, +) + + +def valid_pages(pages: list[int], page_count: int) -> list[int]: + return sorted({page for page in pages if 1 <= int(page) <= page_count}) + + +def repair_hierarchy_assist( + plan: HierarchyAssistPlan, + *, + page_count: int, + toc_pages: set[int], +) -> HierarchyAssistPlan: + plan.exclude_pages_from_title_candidates = valid_pages( + plan.exclude_pages_from_title_candidates, + page_count, + ) + plan.suppress_title_pages = valid_pages(plan.suppress_title_pages, page_count) + plan.prefer_h1_start_pages = [ + candidate + for candidate in plan.prefer_h1_start_pages + if 1 <= candidate.page <= page_count and candidate.page not in toc_pages + ] + plan.section_boundary_hints = [ + hint for hint in plan.section_boundary_hints if 1 <= hint.page <= page_count + ] + return plan + + +def validate_shard_plan( + plan: ShardPlan, + *, + page_count: int, + min_pages: int, + max_pages: int, +) -> ValidationReport: + errors: list[str] = [] + warnings: list[str] = [] + if not plan.shards: + errors.append("shard_plan has no shards") + return ValidationReport(valid=False, errors=errors, warnings=warnings) + expected_start = 1 + for shard in sorted(plan.shards, key=lambda item: item.shard_index): + if shard.page_start != expected_start: + errors.append( + f"shard {shard.shard_index} starts at {shard.page_start}, expected {expected_start}" + ) + if shard.page_end < shard.page_start: + errors.append(f"shard {shard.shard_index} has invalid range") + if shard.page_offset != shard.page_start - 1: + errors.append(f"shard {shard.shard_index} page_offset mismatch") + length = shard.page_end - shard.page_start + 1 + if plan.enabled and length > max_pages: + errors.append(f"shard {shard.shard_index} exceeds max_pages={max_pages}") + if plan.enabled and length < min_pages and shard.page_end != page_count: + warnings.append(f"shard {shard.shard_index} shorter than min_pages={min_pages}") + expected_start = shard.page_end + 1 + if expected_start != page_count + 1: + errors.append("shard_plan does not cover full document") + return ValidationReport(valid=not errors, errors=errors, warnings=warnings) + + +def single_shard_plan(page_count: int) -> ShardPlan: + return ShardPlan( + enabled=False, + reason="not_needed", + shards=[ + Shard( + shard_index=0, + page_start=1, + page_end=max(page_count, 1), + page_offset=0, + anchor_type="forced_max_size", + anchor_evidence="document within shard threshold", + confidence=1.0, + ) + ], + ) + + +def validate_anatomy_map( + anatomy: PageAnatomyMap, + *, + min_pages: int, + max_pages: int, +) -> ValidationReport: + errors: list[str] = [] + warnings: list[str] = [] + page_count = anatomy.page_count + feature_pages = {feature.page for feature in anatomy.page_features} + label_pages = {label.page for label in anatomy.page_labels} + expected_pages = set(range(1, page_count + 1)) + if feature_pages != expected_pages: + errors.append("page_features do not cover every page") + if label_pages != expected_pages: + errors.append("page_labels do not cover every page") + toc_pages = set(anatomy.toc_result.toc_pages) + for candidate in anatomy.h1_result.h1_candidates: + if candidate.page in toc_pages: + errors.append(f"h1 candidate points to toc page {candidate.page}") + if candidate.page < 1 or candidate.page > page_count: + errors.append(f"h1 candidate page {candidate.page} out of range") + shard_report = validate_shard_plan( + anatomy.shard_plan, + page_count=page_count, + min_pages=min_pages, + max_pages=max_pages, + ) + errors.extend(shard_report.errors) + warnings.extend(shard_report.warnings) + return ValidationReport(valid=not errors, errors=errors, warnings=warnings) diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py deleted file mode 100644 index 7b0909dc5..000000000 --- a/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py +++ /dev/null @@ -1,208 +0,0 @@ -"""shard_splitter — PDF physical splitting utilities. - -Provides two public functions: - -- ``split_pdf_by_shards()``: write each ``Shard`` as a temporary sub-PDF file. -- ``merge_shard_dataframes()``: merge per-shard DataFrames and correct - ``page_nums`` offsets so they reflect the original document's page numbers. - -Design notes ------------- -- Splitting uses PyMuPDF ``doc.select()`` which preserves all page content - (images, fonts, annotations) but discards cross-shard hyperlinks — acceptable - for a parsing pipeline. -- Each sub-PDF is written to a temp directory under the job's ``output_dir`` - so it lives on the same filesystem as the rest of the job artifacts. -- ``page_offset`` from ``Shard`` is added to every ``page_nums`` value during - merge, keeping page references consistent with the source document. -""" - -from __future__ import annotations - -import os -import gc -from typing import Any - -import pandas as pd -from app.services.document_agent.page_map import Shard -from loguru import logger - - -# ── Split ────────────────────────────────────────────────────────────────────── - - -def split_pdf_by_shards( - pdf_path: str, - shards: list[Shard], - output_dir: str, -) -> list[dict[str, Any]]: - """Write each shard as an independent sub-PDF under ``output_dir``. - - Returns a list of dicts:: - - [ - { - "shard_index": 0, - "shard": Shard(...), - "sub_pdf_path": "/tmp/.../shard_0_p1-p50.pdf", - "sub_output_dir": "/tmp/.../shard_0/", - }, - … - ] - - Raises ``RuntimeError`` if the source PDF cannot be opened. - Never raises for individual shard write failures — they are logged and - skipped (the caller falls back to parsing the full PDF in that case). - """ - try: - import pymupdf # type: ignore[import] - except ImportError as exc: - raise RuntimeError("PyMuPDF (pymupdf) is required for PDF splitting") from exc - - shards_dir = os.path.join(output_dir, "_shards") - os.makedirs(shards_dir, exist_ok=True) - - try: - doc = pymupdf.open(pdf_path) - total_pages = int(doc.page_count) - except Exception as exc: - raise RuntimeError(f"Cannot open PDF '{pdf_path}': {exc}") from exc - - results: list[dict[str, Any]] = [] - - try: - for shard_idx, shard in enumerate(shards): - # Clamp to actual page count - p_start = max(1, int(shard.page_start)) - p_end = min(total_pages, int(shard.page_end)) - - if p_start > total_pages or p_start > p_end: - logger.warning( - f"[shard_splitter] shard {shard_idx} range [{p_start},{p_end}] " - f"is out of bounds (total={total_pages}) — skipping" - ) - continue - - # PyMuPDF uses 0-based indices - page_indices = list(range(p_start - 1, p_end)) - - sub_pdf_name = f"shard_{shard_idx}_p{p_start}-p{p_end}.pdf" - sub_pdf_path = os.path.join(shards_dir, sub_pdf_name) - sub_output_dir = os.path.join(shards_dir, f"shard_{shard_idx}") - os.makedirs(sub_output_dir, exist_ok=True) - - try: - sub_doc = pymupdf.open() # empty document - sub_doc.insert_pdf(doc, from_page=p_start - 1, to_page=p_end - 1) - sub_doc.save(sub_pdf_path, garbage=4, deflate=True) - sub_doc.close() - - logger.info( - f"[shard_splitter] shard {shard_idx}: pages {p_start}-{p_end} " - f"({len(page_indices)} pages) → {sub_pdf_path}" - ) - results.append( - { - "shard_index": shard_idx, - "shard": shard, - "sub_pdf_path": sub_pdf_path, - "sub_output_dir": sub_output_dir, - } - ) - except Exception as exc: - logger.error( - f"[shard_splitter] failed to write shard {shard_idx} " - f"(pages {p_start}-{p_end}): {exc}" - ) - finally: - doc.close() - gc.collect() - - return results - - -# ── Merge ────────────────────────────────────────────────────────────────────── - - -def merge_shard_dataframes( - shard_dfs: list[tuple[Shard, pd.DataFrame]], -) -> pd.DataFrame: - """Merge per-shard DataFrames into a single DataFrame. - - For each shard, adds ``shard.page_offset`` to every value in the - ``page_nums`` column so that page references reflect the original - document rather than the sub-PDF's local page numbers. - - Args: - shard_dfs: Ordered list of ``(Shard, DataFrame)`` pairs. - DataFrames must have the columns produced by the parser - (``content``, ``path``, ``type``, ``page_nums``, …). - - Returns: - A single DataFrame with all rows, sorted by ``page_nums`` (ascending). - Returns an empty DataFrame if ``shard_dfs`` is empty. - """ - if not shard_dfs: - return pd.DataFrame() - - adjusted: list[pd.DataFrame] = [] - - for shard, df in shard_dfs: - if df is None or df.empty: - logger.warning( - f"[shard_splitter] shard pages {shard.page_start}-{shard.page_end} " - f"produced an empty DataFrame — skipping" - ) - continue - - df_copy = df.copy() - - # Correct page_nums: add page_offset to each page number in the list - if "page_nums" in df_copy.columns: - offset = int(shard.page_offset) - - def _shift_page_nums(val: Any, off: int = offset) -> Any: - if isinstance(val, list): - return [p + off for p in val if isinstance(p, int)] - if isinstance(val, str): - # Stored as comma-separated string in some paths - try: - nums = [int(x.strip()) for x in val.split(",") if x.strip()] - return ",".join(str(p + off) for p in nums) - except ValueError: - return val - return val - - df_copy["page_nums"] = df_copy["page_nums"].apply(_shift_page_nums) - - adjusted.append(df_copy) - logger.info( - f"[shard_splitter] merged shard pages " - f"{shard.page_start}-{shard.page_end} ({len(df_copy)} rows, " - f"offset={shard.page_offset})" - ) - - if not adjusted: - return pd.DataFrame() - - merged = pd.concat(adjusted, ignore_index=True) - - # Sort by the first page number of each row's page_nums for document order - if "page_nums" in merged.columns: - def _first_page(val: Any) -> int: - if isinstance(val, list) and val: - return int(val[0]) - if isinstance(val, str): - try: - parts = [int(x.strip()) for x in val.split(",") if x.strip()] - return parts[0] if parts else 0 - except ValueError: - return 0 - return 0 - - merged["_sort_page"] = merged["page_nums"].apply(_first_page) - merged = merged.sort_values("_sort_page").drop(columns=["_sort_page"]) - merged = merged.reset_index(drop=True) - - logger.info(f"[shard_splitter] merge complete: {len(merged)} total rows") - return merged diff --git a/apps/worker/build_manifest_sjsyj.py b/apps/worker/build_manifest_sjsyj.py new file mode 100644 index 000000000..87a0319fb --- /dev/null +++ b/apps/worker/build_manifest_sjsyj.py @@ -0,0 +1,159 @@ +""" +从 preds_5_final_output.csv 构建 HIERARCHY 树， +输出 manifest.json 并与原始 MD 的 # 标题对比 +""" + +import json +import os +import re +import pandas as pd +from collections import OrderedDict + +CSV_PATH = ( + "/Users/wuchengke/Desktop/temp/ontos/parse_comparison/" + "sjsyj_mineru_notable/SJSYJ-SC-2024 企业制度汇编（上册）/txt/" + "preds_5_final_output.csv" +) +MD_PATH = ( + "/Users/wuchengke/Desktop/temp/ontos/parse_comparison/" + "sjsyj_mineru_notable/SJSYJ-SC-2024 企业制度汇编（上册）/txt/" + "SJSYJ-SC-2024 企业制度汇编（上册）.md" +) +OUT_DIR = os.path.dirname(CSV_PATH) +SOURCE_FILE = "SJSYJ-SC-2024 企业制度汇编（上册）.pdf" + + +def clean_heading(h: str) -> str: + """去掉 Markdown # 前缀和前后空格""" + return re.sub(r"^#+\s*", "", str(h)).strip() + + +def build_hierarchy(df: pd.DataFrame) -> dict: + """从有层级的行构建嵌套字典树""" + root = OrderedDict() + root["Root"] = {} + stack = [] # [(level, dict)] + + valid = df[df["level"].astype(str).str.match(r"^[1-9]")].copy() + valid["level"] = valid["level"].astype(int) + + for _, row in valid.iterrows(): + title = clean_heading(str(row["heading"])) + level = int(row["level"]) + + # 弹栈到当前父节点 + while stack and stack[-1][0] >= level: + stack.pop() + + parent = stack[-1][1] if stack else root + # 处理重名 + key, suffix = title, 2 + while key in parent: + key = f"{title} ({suffix})" + suffix += 1 + parent[key] = OrderedDict() + stack.append((level, parent[key])) + + return root + + +def extract_md_headings(md_path: str) -> list[dict]: + """从 MD 文件提取所有 # 标题""" + headings = [] + with open(md_path, encoding="utf-8") as f: + for lineno, line in enumerate(f, 1): + m = re.match(r"^(#{1,6})\s+(.+)", line.rstrip()) + if m: + level = len(m.group(1)) + text = m.group(2).strip() + headings.append({"lineno": lineno, "level": level, "text": text}) + return headings + + +def compare_with_md(hierarchy: dict, md_headings: list[dict]): + """对比 hierarchy 树中的标题 vs MD 中的 # 标题""" + def flatten(d, prefix="", result=None): + if result is None: + result = [] + for k, v in d.items(): + result.append(k) + if isinstance(v, dict) and v: + flatten(v, prefix + k + "/", result) + return result + + tree_titles = set(flatten(hierarchy)) + tree_titles.discard("Root") + + md_texts = set(h["text"] for h in md_headings) + + in_tree_not_md = tree_titles - md_texts + in_md_not_tree = md_texts - tree_titles + both = tree_titles & md_texts + + print(f"\n{'='*60}") + print(f"📊 标题对比") + print(f"{'='*60}") + print(f" MD 中 # 标题数: {len(md_texts)}") + print(f" Hierarchy 树标题数: {len(tree_titles)}") + print(f" 完全匹配: {len(both)}") + print(f" 仅在 Hierarchy 中: {len(in_tree_not_md)}") + print(f" 仅在 MD 中: {len(in_md_not_tree)}") + + if in_tree_not_md: + print(f"\n⚠️ Hierarchy 有但 MD # 里没有（LLM 可能误识别）:") + for t in sorted(in_tree_not_md)[:15]: + print(f" - {t[:80]}") + + if in_md_not_tree: + print(f"\n⚠️ MD # 有但 Hierarchy 没有（可能被降级为正文）:") + for t in sorted(in_md_not_tree)[:15]: + print(f" - {t[:80]}") + + return {"matched": len(both), "only_tree": len(in_tree_not_md), "only_md": len(in_md_not_tree)} + + +def main(): + print(f"📄 读取: {CSV_PATH}") + df = pd.read_csv(CSV_PATH, encoding="utf-8-sig") + print(f" 总行数: {len(df)}") + + # 构建树 + hierarchy = build_hierarchy(df) + top_level = [k for k in hierarchy if k != "Root"] + print(f" 顶层章节数: {len(top_level)}") + + # 输出 manifest.json + manifest = { + "version": "2.0", + "job_id": SOURCE_FILE, + "source_file_name": SOURCE_FILE, + "processing_date": "2026-05-21T04:46:35Z", + "statistics": { + "total_chunks": int((df["level"] == -1).sum()), + "heading_count": int((df["level"].astype(str).str.match(r"^[1-9]")).sum()), + }, + "HIERARCHY": hierarchy, + } + + out_path = os.path.join(OUT_DIR, "manifest.json") + with open(out_path, "w", encoding="utf-8") as f: + json.dump(manifest, f, ensure_ascii=False, indent=2) + size_kb = os.path.getsize(out_path) / 1024 + print(f"\n✅ 已保存: {out_path} ({size_kb:.0f} KB)") + + # 预览顶层结构 + print(f"\n📋 顶层章节（前 20）:") + for i, k in enumerate(top_level[:20], 1): + children = hierarchy[k] + child_count = len(children) if isinstance(children, dict) else 0 + print(f" L1 [{i:2d}] {k[:60]}{'...' if len(k)>60 else ''} → {child_count} 子节点") + + # 与 MD 对比 + print(f"\n📄 读取 MD: {MD_PATH}") + md_headings = extract_md_headings(MD_PATH) + print(f" MD # 标题数: {len(md_headings)}") + compare_with_md(hierarchy, md_headings) + + +if __name__ == "__main__": + main() diff --git a/apps/worker/run_hierarchy_sjsyj.py b/apps/worker/run_hierarchy_sjsyj.py new file mode 100644 index 000000000..a3021b7d8 --- /dev/null +++ b/apps/worker/run_hierarchy_sjsyj.py @@ -0,0 +1,117 @@ +""" +针对 MinerU 解析输出的 MD 文件，运行 pred_titles 完整流程 +输出 preds_3_llm_base.csv 和 preds_5_final_output.csv 到同目录 +""" + +import os +import sys +import json +import shutil + +# ── 配置 ── +MODEL_NAME = "qwen3.5-27b" +ENABLE_THINKING = False + +# 目标 MD 文件 +TARGET_MD = ( + "/Users/wuchengke/Desktop/temp/ontos/parse_comparison/" + "sjsyj_mineru_notable/SJSYJ-SC-2024 企业制度汇编（上册）/txt/" + "SJSYJ-SC-2024 企业制度汇编（上册）.md" +) +OUTPUT_DIR = os.path.dirname(TARGET_MD) + +# ── 工程路径注入 ── +WORKER_DIR = "/Users/wuchengke/Desktop/knowhere/knowhereapi-main/apps/worker" +SHARED_DIR = "/Users/wuchengke/Desktop/knowhere/knowhereapi-main/packages/shared-python" +sys.path.insert(0, SHARED_DIR) +sys.path.insert(0, WORKER_DIR) + +from dotenv import load_dotenv +load_dotenv(os.path.join(WORKER_DIR, ".env")) +os.environ["LOCAL_DEBUG"] = "1" + +import openai +from loguru import logger + +# ── Qwen thinking mode 关闭 patch ── +_original_create = openai.resources.chat.completions.Completions.create + +def patched_create(self, *args, **kwargs): + kwargs.setdefault("extra_body", {}) + kwargs["extra_body"]["enable_thinking"] = False + return _original_create(self, *args, **kwargs) + +openai.resources.chat.completions.Completions.create = patched_create +logger.info("🚫 Thinking mode 已关闭 (enable_thinking=false)") + +from app.services.document_parser.structure.layout_parser import pred_titles +from app.services.document_parser.structure.toc_parser import detect_tocs_in_texts +from app.services.document_parser.formats.html.parser import merge_html_tables + + +def main(): + logger.info("=" * 60) + logger.info(f"🚀 SJSYJ hierarchy 检测 — 模型: {MODEL_NAME}") + logger.info(f" 输入: {TARGET_MD}") + logger.info(f" 输出: {OUTPUT_DIR}") + logger.info("=" * 60) + + # 1. 加载 MD + with open(TARGET_MD, "r", encoding="utf-8") as f: + md_lines = f.read().splitlines() + md_lines = [line.strip() for line in md_lines if line.strip()] + md_lines = merge_html_tables(md_lines) + logger.info(f"📄 MD 加载完毕: {len(md_lines)} 行") + + # 2. TOC 检测 + toc_json_path = os.path.join(OUTPUT_DIR, "toc_hierarchies.json") + toc_hierarchies = None + if os.path.exists(toc_json_path): + os.remove(toc_json_path) + logger.info("🗑️ 已删除旧 toc_hierarchies.json") + + logger.info("🔍 检测 TOC...") + toc_hierarchies, md_lines = detect_tocs_in_texts(md_lines, model_name=MODEL_NAME) + toc_hierarchies = toc_hierarchies or [] + if toc_hierarchies: + with open(toc_json_path, "w", encoding="utf-8") as f: + json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) + logger.info(f"✅ TOC 检测完毕: {len(toc_hierarchies)} 个区域 → {toc_json_path}") + else: + logger.info(" 未检测到 TOC") + + # 3. 运行 pred_titles（完整流程，输出 CSV） + logger.info(f"🧠 运行 pred_titles (model={MODEL_NAME}, smart_parse=True)...") + heading_preds = pred_titles( + infos=md_lines, + doc_type="md", + toc_hierarchies=toc_hierarchies or [], + prompt_limt=4000, + enable_regx=True, + smart_parse=True, + model_name=MODEL_NAME, + output_dir=OUTPUT_DIR, # ← CSV 保存到这里 + layout_json_path=None, + ) + + if heading_preds.empty: + logger.warning("⚠️ 没有检测到任何有效标题") + return + + valid = heading_preds[heading_preds["level"] > 0] + logger.info(f"✅ 完成! 有效标题 {len(valid)} 个 / 总行 {len(heading_preds)}") + logger.info(f" 层级分布:\n{heading_preds['level'].value_counts().sort_index().to_string()}") + + # 汇报生成的 CSV + for csv_name in ["preds_3_llm_base.csv", "preds_4_llm_final.csv", "preds_5_final_output.csv"]: + p = os.path.join(OUTPUT_DIR, csv_name) + if os.path.exists(p): + size_kb = os.path.getsize(p) / 1024 + logger.info(f" 📄 {csv_name} → {p} ({size_kb:.0f} KB)") + + logger.info("=" * 60) + logger.info("🎉 Done!") + + +if __name__ == "__main__": + main() diff --git a/packages/shared-python/shared/models/database/__init__.py b/packages/shared-python/shared/models/database/__init__.py index a0aba9d3d..2d0e7d8a8 100644 --- a/packages/shared-python/shared/models/database/__init__.py +++ b/packages/shared-python/shared/models/database/__init__.py @@ -20,10 +20,12 @@ RetrievalRun, RetrievalStep, ) +from .document_page_plan import DocumentPagePlan from .demo_materialization import DemoMaterialization from .guest_device import GuestDevice from .job import Job from .job_result import JobChunk, JobResult +from .parse_agent import ParseRun, ParseStep # 3. Job-related log models from .job_state_audit_log import JobStateAuditLog @@ -54,12 +56,15 @@ "Document", "DocumentSection", "DocumentChunk", + "DocumentPagePlan", "DemoMaterialization", "GraphNode", "GraphEdge", "RetrievalHitStat", "RetrievalRun", "RetrievalStep", + "ParseRun", + "ParseStep", "StripePriceConfig", "PaymentRecord", "JobStateAuditLog", diff --git a/packages/shared-python/shared/models/database/document_page_plan.py b/packages/shared-python/shared/models/database/document_page_plan.py new file mode 100644 index 000000000..b7f2424b7 --- /dev/null +++ b/packages/shared-python/shared/models/database/document_page_plan.py @@ -0,0 +1,37 @@ +"""Persisted page anatomy and future page-processing plans.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Dict, Optional + +from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String +from sqlalchemy.orm import Mapped, mapped_column + +from shared.core.database import Base +from shared.utils.utc_now import utc_now_naive + + +class DocumentPagePlan(Base): + __tablename__ = "document_page_plan" + + page_plan_id: Mapped[str] = mapped_column(String(36), primary_key=True) + job_id: Mapped[str] = mapped_column( + String(36), ForeignKey("jobs.job_id", ondelete="CASCADE"), nullable=False + ) + page_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + hierarchy_assist: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + shard_plan: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + page_processing_plan: Mapped[Optional[Dict[str, Any]]] = mapped_column( + JSON, + nullable=True, + ) + global_signals: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + Index("idx_document_page_plan_job", "job_id"), + Index("idx_document_page_plan_created", "created_at"), + ) diff --git a/packages/shared-python/shared/models/database/parse_agent.py b/packages/shared-python/shared/models/database/parse_agent.py new file mode 100644 index 000000000..430957c43 --- /dev/null +++ b/packages/shared-python/shared/models/database/parse_agent.py @@ -0,0 +1,63 @@ +"""Parse-side agent trace models.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Dict, Optional + +from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String +from sqlalchemy.orm import Mapped, mapped_column + +from shared.core.database import Base +from shared.utils.utc_now import utc_now_naive + + +class ParseRun(Base): + __tablename__ = "parse_runs" + + run_id: Mapped[str] = mapped_column(String(36), primary_key=True) + job_id: Mapped[str] = mapped_column( + String(36), ForeignKey("jobs.job_id", ondelete="CASCADE"), nullable=False + ) + kind: Mapped[str] = mapped_column(String(32), nullable=False, default="profile") + final_status: Mapped[str] = mapped_column(String(32), nullable=False) + rounds_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + total_tokens: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + total_latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + summary: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + started_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + finished_at: Mapped[Optional[datetime]] = mapped_column( + DateTime, default=utc_now_naive, nullable=True + ) + + __table_args__ = ( + Index("idx_parse_runs_job_kind", "job_id", "kind"), + Index("idx_parse_runs_started", "started_at"), + ) + + +class ParseStep(Base): + __tablename__ = "parse_steps" + + step_id: Mapped[str] = mapped_column(String(36), primary_key=True) + run_id: Mapped[str] = mapped_column( + String(36), ForeignKey("parse_runs.run_id", ondelete="CASCADE"), nullable=False + ) + round_index: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + actor: Mapped[str] = mapped_column(String(64), nullable=False) + action_type: Mapped[str] = mapped_column(String(64), nullable=False) + tool_name: Mapped[Optional[str]] = mapped_column(String(64), nullable=True) + tool_args: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + observation: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + tokens_used: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + Index("idx_parse_steps_run_round", "run_id", "round_index"), + Index("idx_parse_steps_tool", "tool_name"), + ) diff --git a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py index 79a00e97a..e8d5f4135 100644 --- a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py +++ b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py @@ -176,14 +176,14 @@ def _should_use_ali_pool(self) -> bool: return False return _is_ali_model(self.default_model) - def _make_ali_pool_call( + def _make_ali_pool_raw_call( self, model: str, all_messages: List[ChatCompletionMessageParam], temperature: float, max_tokens: int, api_kwargs: Dict[str, Any], - ) -> tuple[str, LLMUsage]: + ) -> tuple[Any, LLMUsage]: """Acquire a token, make the call, and retry inline on 429.""" from shared.services.ai.ali_quota_manager import get_ali_quota_manager @@ -208,13 +208,12 @@ def _make_ali_pool_call( max_tokens=max_tokens, **api_kwargs, ) - choices = response.choices - if not choices: + if not response.choices: raise LLMServiceException( internal_message="AI returned empty result", provider=self.default_model, ) - return choices[0].message.content or "", _extract_usage(response) + return response, _extract_usage(response) except openai.RateLimitError as exc: retry_after = _parse_retry_after(exc) quota_manager.mark_rate_limited(lease.token_id, retry_after) @@ -253,8 +252,114 @@ def _make_ali_pool_call( provider=self.default_model, ) + def _make_ali_pool_call( + self, + model: str, + all_messages: List[ChatCompletionMessageParam], + temperature: float, + max_tokens: int, + api_kwargs: Dict[str, Any], + ) -> tuple[str, LLMUsage]: + response, usage = self._make_ali_pool_raw_call( + model=model, + all_messages=all_messages, + temperature=temperature, + max_tokens=max_tokens, + api_kwargs=api_kwargs, + ) + return response.choices[0].message.content or "", usage + # ------------------------------------------------------------------ + def chat_completion_raw_with_usage( + self, + messages: Union[str, List[ChatCompletionMessageParam]], + model: Optional[str] = None, + temperature: float = 0.1, + max_tokens: int = 4096, + top_p: Optional[float] = None, + timeout: Optional[int] = None, + **kwargs, + ) -> tuple[Any, LLMUsage]: + all_messages: List[ChatCompletionMessageParam] + if isinstance(messages, list): + all_messages = messages # type: ignore[assignment] + else: + all_messages = [{"role": "user", "content": str(messages)}] + + api_kwargs: Dict[str, Any] = {} + if top_p is not None: + api_kwargs["top_p"] = top_p + if timeout is not None: + api_kwargs["timeout"] = timeout + allowed_api_params = { + "n", "stop", "presence_penalty", "frequency_penalty", + "logit_bias", "user", "seed", "tools", "tool_choice", + "response_format", "logprobs", "top_logprobs", + } + for key, value in kwargs.items(): + if key in allowed_api_params: + api_kwargs[key] = value + + extra_body = api_kwargs.get("extra_body", {}) + if isinstance(extra_body, dict): + extra_body.setdefault("enable_thinking", False) + else: + extra_body = {"enable_thinking": False} + api_kwargs["extra_body"] = extra_body + + effective_model = model or self.default_model + if _should_mock_llm_calls(): + content = build_mock_chat_completion_response( + messages=all_messages, + model_name=effective_model, + ) + return {"mock_content": content}, _empty_usage() + + if self._should_use_ali_pool(): + return self._make_ali_pool_raw_call( + model=effective_model, + all_messages=all_messages, + temperature=temperature, + max_tokens=max_tokens, + api_kwargs=api_kwargs, + ) + + client = self._client + if client is None: + raise LLMServiceException( + internal_message="OpenAI client is not initialized for direct provider requests", + provider=self.default_model, + ) + try: + response = client.chat.completions.create( + model=effective_model, + messages=all_messages, + temperature=temperature, + max_tokens=max_tokens, + **api_kwargs, + ) + if not response.choices: + raise LLMServiceException( + internal_message="AI returned empty result", + provider=self.default_model, + ) + return response, _extract_usage(response) + except LLMServiceException: + raise + except Exception as exc: + logger.error( + "LLM raw request failed: model={model}, base_url={base_url}, error_chain={error_chain}", + model=effective_model, + base_url=client.base_url, + error_chain=_summarize_exception_chain(exc), + ) + raise LLMServiceException( + internal_message=f"API request failed: {_summarize_exception_chain(exc)}", + provider=self.default_model, + original_exception=exc, + ) from exc + def chat_completion_with_usage( self, messages: Union[str, List[ChatCompletionMessageParam]], From a5950c36540d7fbe2c9263655d9ba598a614718d Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 22 May 2026 12:06:21 +0800 Subject: [PATCH 04/11] refactor: update DocumentAgent to use boundary candidates and streamline transitions This commit modifies the DocumentAgent's state management and tool integration by replacing the legacy H1 boundary and TOC page detection tools with a new boundary candidate system. Key changes include: - Updated the state transitions to utilize `collect.boundary_candidates` instead of the removed `find.h1_boundaries` and `find.toc_pages`. - Introduced a new `BoundaryCandidate` data structure in the manifest to encapsulate boundary candidate details. - Adjusted the `ProfileCoordinator` and validation logic to accommodate the new boundary candidates. - Enhanced the `ParseRunRecorder` to include detailed trace information for boundary candidates. - Removed obsolete tools related to H1 and TOC detection, simplifying the toolset. These changes aim to improve the clarity and efficiency of the DocumentAgent's processing pipeline. --- .../services/document_agent/coordinator.py | 7 +- .../app/services/document_agent/manifest.py | 36 +- .../app/services/document_agent/state.py | 2 + .../services/document_agent/tools/__init__.py | 3 +- .../tools/classify_page_kinds.py | 30 +- .../tools/collect_boundary_candidates.py | 134 ++++++++ .../tools/find_h1_boundaries.py | 119 ------- .../document_agent/tools/find_toc_pages.py | 82 ----- .../tools/persist_anatomy_map.py | 1 + .../tools/propose_hierarchy_assist.py | 81 +---- .../tools/propose_shard_plan.py | 307 +++++++++++++----- .../tools/validate_anatomy_map.py | 1 + .../app/services/document_agent/trace.py | 35 ++ .../app/services/document_agent/validators.py | 19 ++ 14 files changed, 504 insertions(+), 353 deletions(-) create mode 100644 apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py delete mode 100644 apps/worker/app/services/document_agent/tools/find_h1_boundaries.py delete mode 100644 apps/worker/app/services/document_agent/tools/find_toc_pages.py diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 8bbf52699..1f5fa0968 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -19,8 +19,7 @@ TRANSITIONS: dict[str, DocumentAgentState] = { "probe.page_features": DocumentAgentState.PROBED, "classify.page_kinds": DocumentAgentState.PROBED, - "find.toc_pages": DocumentAgentState.PROBED, - "find.h1_boundaries": DocumentAgentState.H1_FOUND, + "collect.boundary_candidates": DocumentAgentState.H1_FOUND, "propose.hierarchy_assist": DocumentAgentState.H1_FOUND, "propose.shard_plan": DocumentAgentState.H1_FOUND, "validate.anatomy_map": DocumentAgentState.VALIDATED, @@ -30,8 +29,7 @@ REQUIRED_TOOLS = [ "probe.page_features", "classify.page_kinds", - "find.toc_pages", - "find.h1_boundaries", + "collect.boundary_candidates", "propose.hierarchy_assist", "propose.shard_plan", "validate.anatomy_map", @@ -143,7 +141,6 @@ def _maybe_advance_composite_state(self) -> None: if ( self.state == DocumentAgentState.PROBED and self.blackboard.page_labels - and self.blackboard.toc_result is not None ): self.state = DocumentAgentState.CLASSIFIED self.blackboard.mark(self.state) diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index 16515fafd..58aa13c7f 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -25,6 +25,14 @@ "sparse", ] +BoundaryCandidateKind = Literal[ + "h1", + "toc", + "blank", + "sparse", + "separator", +] + @dataclass class PageFeature: @@ -117,6 +125,18 @@ def to_dict(self) -> dict[str, Any]: return asdict(self) +@dataclass +class BoundaryCandidate: + page: int + kind: BoundaryCandidateKind + priority: int + confidence: float + evidence: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + @dataclass class HierarchyAssistPlan: exclude_pages_from_title_candidates: list[int] = field(default_factory=list) @@ -170,7 +190,13 @@ def to_dict(self) -> dict[str, Any]: @dataclass class ShardPlan: enabled: bool - reason: Literal["too_large", "not_needed", "parser_stability", "hierarchy_isolation"] + reason: Literal[ + "too_large", + "not_needed", + "parser_stability", + "hierarchy_isolation", + "llm_boundary_decision", + ] shards: list[Shard] = field(default_factory=list) validation: ValidationReport = field( default_factory=lambda: ValidationReport(valid=True) @@ -222,6 +248,7 @@ class PageAnatomyMap: h1_result: H1BoundaryResult hierarchy_assist: HierarchyAssistPlan shard_plan: ShardPlan + boundary_candidates: list[BoundaryCandidate] = field(default_factory=list) page_processing_plan: PageProcessingPlan | None = None global_signals: dict[str, Any] = field(default_factory=dict) trace_summary: dict[str, Any] = field(default_factory=dict) @@ -240,6 +267,9 @@ def to_dict(self) -> dict[str, Any]: "h1_result": self.h1_result.to_dict(), "hierarchy_assist": self.hierarchy_assist.to_dict(), "shard_plan": self.shard_plan.to_dict(), + "boundary_candidates": [ + candidate.to_dict() for candidate in self.boundary_candidates + ], "page_processing_plan": ( self.page_processing_plan.to_dict() if self.page_processing_plan is not None @@ -258,6 +288,10 @@ class ToolResult: latency_ms: int = 0 error: str | None = None tokens_used: int = 0 + input_summary: dict[str, Any] | None = None + output_summary: dict[str, Any] | None = None + warnings: list[str] = field(default_factory=list) + debug: dict[str, Any] | None = None @dataclass diff --git a/apps/worker/app/services/document_agent/state.py b/apps/worker/app/services/document_agent/state.py index 47f3bc0b4..a752f129b 100644 --- a/apps/worker/app/services/document_agent/state.py +++ b/apps/worker/app/services/document_agent/state.py @@ -7,6 +7,7 @@ from typing import Any from app.services.document_agent.manifest import ( + BoundaryCandidate, H1BoundaryResult, HierarchyAssistPlan, PageFeature, @@ -36,6 +37,7 @@ class AgentBlackboard: page_labels: list[PageLabel] = field(default_factory=list) toc_result: TocResult | None = None h1_result: H1BoundaryResult | None = None + boundary_candidates: list[BoundaryCandidate] = field(default_factory=list) hierarchy_assist: HierarchyAssistPlan | None = None shard_plan: ShardPlan | None = None validation_report: dict[str, Any] | None = None diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index 7206a6dce..52084ae75 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -3,8 +3,7 @@ from app.services.document_agent.registry import REGISTRY from . import classify_page_kinds as classify_page_kinds # noqa: F401 -from . import find_h1_boundaries as find_h1_boundaries # noqa: F401 -from . import find_toc_pages as find_toc_pages # noqa: F401 +from . import collect_boundary_candidates as collect_boundary_candidates # noqa: F401 from . import persist_anatomy_map as persist_anatomy_map # noqa: F401 from . import probe_page_features as probe_page_features # noqa: F401 from . import propose_hierarchy_assist as propose_hierarchy_assist # noqa: F401 diff --git a/apps/worker/app/services/document_agent/tools/classify_page_kinds.py b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py index 399741243..2305652c1 100644 --- a/apps/worker/app/services/document_agent/tools/classify_page_kinds.py +++ b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py @@ -3,6 +3,7 @@ from __future__ import annotations import time +from collections import Counter, defaultdict from typing import Any from app.services.document_agent.manifest import PageFeature, PageLabel, ToolContext, ToolResult @@ -91,12 +92,33 @@ def classify_page_kinds(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() labels = [_label_feature(feature) for feature in ctx.blackboard.page_features] ctx.blackboard.page_labels = labels - counts: dict[str, int] = {} - for label in labels: - counts[label.kind] = counts.get(label.kind, 0) + 1 + counts = Counter(label.kind for label in labels) ctx.blackboard.global_signals["page_kind_counts"] = counts + features_by_page = {feature.page: feature for feature in ctx.blackboard.page_features} + samples: dict[str, list[dict[str, Any]]] = defaultdict(list) + for label in labels: + if len(samples[label.kind]) >= 8: + continue + feature = features_by_page.get(label.page) + samples[label.kind].append( + { + "page": label.page, + "confidence": label.confidence, + "evidence": label.evidence, + "raw_text_length": feature.raw_text_length if feature else None, + "image_coverage": feature.image_coverage if feature else None, + "table_count": feature.table_count if feature else None, + "drawings_count": feature.drawings_count if feature else None, + "text_preview": (feature.text_lines_preview[:4] if feature else []), + } + ) return ToolResult( status="ok", - payload={"page_kind_counts": counts}, + payload={"page_kind_counts": dict(counts)}, latency_ms=int((time.monotonic() - start) * 1000), + input_summary={"page_count": ctx.blackboard.page_count}, + output_summary={ + "page_kind_counts": dict(counts), + "sample_pages_by_kind": dict(samples), + }, ) diff --git a/apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py b/apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py new file mode 100644 index 000000000..aec0cd279 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py @@ -0,0 +1,134 @@ +"""Collect candidate split pages without deciding where to split.""" + +from __future__ import annotations + +import time +from collections import Counter +from typing import Any + +from app.services.document_agent.manifest import ( + BoundaryCandidate, + H1BoundaryResult, + TocResult, + ToolContext, + ToolResult, +) +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState + + +BOUNDARY_PAGE_KINDS = {"blank", "sparse", "separator", "toc"} +PRIORITY_BY_KIND = { + "h1": 100, + "toc": 80, + "separator": 70, + "blank": 45, + "sparse": 40, +} + + +def _feature_by_page(ctx: ToolContext) -> dict[int, Any]: + return {feature.page: feature for feature in ctx.blackboard.page_features} + + +def _candidate_evidence(ctx: ToolContext, page: int, kind: str) -> dict[str, Any]: + feature = _feature_by_page(ctx).get(page) + if feature is None: + return {} + return { + "source": "page_label", + "label_kind": kind, + "position_ratio": round(page / max(ctx.blackboard.page_count, 1), 4), + "raw_text_length": feature.raw_text_length, + "image_coverage": feature.image_coverage, + "table_count": feature.table_count, + "drawings_count": feature.drawings_count, + "text_preview": feature.text_lines_preview[:5], + } + + +@register_tool( + name="collect.boundary_candidates", + description="Collect sparse, blank, TOC, and H1 pages as split candidates without making split decisions.", + allowed_states={DocumentAgentState.CLASSIFIED}, +) +def collect_boundary_candidates(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + candidates: list[BoundaryCandidate] = [] + seen: set[tuple[int, str]] = set() + + # The robust TOC parser will populate these later. Until then, keep the + # anatomy contract explicit without inventing TOC/H1 results here. + ctx.blackboard.toc_result = ctx.blackboard.toc_result or TocResult( + method="none", + notes="TOC parser not run in boundary candidate collection", + ) + ctx.blackboard.h1_result = ctx.blackboard.h1_result or H1BoundaryResult( + method="none", + notes="H1 parser not run in boundary candidate collection", + ) + + for label in ctx.blackboard.page_labels: + if label.kind not in BOUNDARY_PAGE_KINDS: + continue + key = (label.page, label.kind) + if key in seen: + continue + seen.add(key) + candidates.append( + BoundaryCandidate( + page=label.page, + kind=label.kind, # type: ignore[arg-type] + priority=PRIORITY_BY_KIND.get(label.kind, 0), + confidence=label.confidence, + evidence=_candidate_evidence(ctx, label.page, label.kind), + ) + ) + + # Preserve future H1 candidates as highest-priority hints when upstream TOC + # parsing is wired in, but do not derive them with the old fuzzy matcher. + if ctx.blackboard.h1_result: + for h1 in ctx.blackboard.h1_result.h1_candidates: + key = (h1.page, "h1") + if key in seen: + continue + seen.add(key) + candidates.append( + BoundaryCandidate( + page=h1.page, + kind="h1", + priority=PRIORITY_BY_KIND["h1"], + confidence=h1.confidence, + evidence={ + "source": h1.source, + "title": h1.title, + "matched_line": h1.matched_line, + "position_ratio": round( + h1.page / max(ctx.blackboard.page_count, 1), 4 + ), + **h1.evidence, + }, + ) + ) + + candidates.sort(key=lambda item: (item.page, -item.priority)) + ctx.blackboard.boundary_candidates = candidates + counts = Counter(candidate.kind for candidate in candidates) + ctx.blackboard.global_signals["boundary_candidate_counts"] = dict(counts) + + return ToolResult( + status="ok", + payload={ + "candidate_count": len(candidates), + "candidate_counts": dict(counts), + }, + latency_ms=int((time.monotonic() - start) * 1000), + input_summary={ + "page_count": ctx.blackboard.page_count, + "page_kind_counts": ctx.blackboard.global_signals.get("page_kind_counts", {}), + }, + output_summary={ + "candidate_counts": dict(counts), + "sample_candidates": [candidate.to_dict() for candidate in candidates[:20]], + }, + ) diff --git a/apps/worker/app/services/document_agent/tools/find_h1_boundaries.py b/apps/worker/app/services/document_agent/tools/find_h1_boundaries.py deleted file mode 100644 index 3098fd4ed..000000000 --- a/apps/worker/app/services/document_agent/tools/find_h1_boundaries.py +++ /dev/null @@ -1,119 +0,0 @@ -"""Find H1 starts from TOC candidates or heading-like body lines.""" - -from __future__ import annotations - -import time -from typing import Any - -from app.services.document_agent.heading_text import ( - candidate_allowed, - clean_toc_line, - fuzzy_match, - looks_like_h1_line, - normalize_heading, -) -from app.services.document_agent.manifest import H1BoundaryResult, H1Candidate, TocCandidate, ToolContext, ToolResult -from app.services.document_agent.pdf_text import meaningful_lines, read_page_texts -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState - - -def _all_non_toc_pages(ctx: ToolContext) -> list[int]: - toc_pages = set(ctx.blackboard.toc_result.toc_pages if ctx.blackboard.toc_result else []) - return [ - feature.page - for feature in ctx.blackboard.page_features - if feature.page not in toc_pages - ] - - -def _fallback_candidates_from_previews(ctx: ToolContext) -> list[TocCandidate]: - candidates: list[TocCandidate] = [] - seen: set[str] = set() - for feature in ctx.blackboard.page_features: - for line_index, line in enumerate(feature.text_lines_preview): - cleaned = clean_toc_line(line) - if not looks_like_h1_line(cleaned) or not candidate_allowed(cleaned): - continue - normalized = normalize_heading(cleaned) - if normalized in seen: - continue - seen.add(normalized) - candidates.append( - TocCandidate( - title=cleaned, - normalized_title=normalized, - source_page=feature.page, - line_index=line_index, - ) - ) - return candidates - - -def _match_candidate( - candidate: TocCandidate, - page_texts: dict[int, str], -) -> H1Candidate | None: - for page, text in sorted(page_texts.items()): - lines = meaningful_lines(text) - for line_index, line in enumerate(lines[:20]): - if candidate.title in line: - return H1Candidate( - title=candidate.title, - page=page, - confidence=1.0, - matched_line=line, - source="toc_exact_top", - evidence={"line_index": line_index, "toc_page": candidate.source_page}, - ) - if fuzzy_match(candidate.normalized_title, line): - return H1Candidate( - title=candidate.title, - page=page, - confidence=0.86, - matched_line=line, - source="toc_fuzzy_top", - evidence={"line_index": line_index, "toc_page": candidate.source_page}, - ) - return None - - -@register_tool( - name="find.h1_boundaries", - description="Match H1 candidates to body page starts with evidence and confidence.", - allowed_states={DocumentAgentState.CLASSIFIED}, -) -def find_h1_boundaries(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: - start = time.monotonic() - toc_result = ctx.blackboard.toc_result - candidates = list(toc_result.candidates) if toc_result and toc_result.candidates else [] - method: str = "toc_grep" if candidates else "heading_grep" - if not candidates: - candidates = _fallback_candidates_from_previews(ctx) - pages = _all_non_toc_pages(ctx) - page_texts = read_page_texts(ctx.pdf_path, pages) if pages else {} - matches: list[H1Candidate] = [] - seen_pages_titles: set[tuple[int, str]] = set() - for candidate in candidates: - match = _match_candidate(candidate, page_texts) - if match is None: - continue - key = (match.page, normalize_heading(match.title)) - if key in seen_pages_titles: - continue - seen_pages_titles.add(key) - matches.append(match) - if not matches: - method = "none" - result = H1BoundaryResult( - h1_candidates=matches, - method=method, # type: ignore[arg-type] - notes=f"{len(matches)} of {len(candidates)} candidates matched", - ) - ctx.blackboard.h1_result = result - ctx.blackboard.global_signals["h1_candidate_count"] = len(matches) - return ToolResult( - status="ok", - payload={"h1_match_count": len(matches), "method": method}, - latency_ms=int((time.monotonic() - start) * 1000), - ) diff --git a/apps/worker/app/services/document_agent/tools/find_toc_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_pages.py deleted file mode 100644 index e5a3c8fc3..000000000 --- a/apps/worker/app/services/document_agent/tools/find_toc_pages.py +++ /dev/null @@ -1,82 +0,0 @@ -"""TOC page and level-1 candidate extraction.""" - -from __future__ import annotations - -import time -from typing import Any - -from app.services.document_agent.heading_text import ( - candidate_allowed, - clean_toc_line, - normalize_heading, - looks_like_h1_line, -) -from app.services.document_agent.manifest import TocCandidate, TocResult, ToolContext, ToolResult -from app.services.document_agent.pdf_text import read_page_texts -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState - - -def _is_toc_feature_lines(lines: list[str]) -> bool: - compact = "".join(lines).replace(" ", "").lower() - return any(marker in compact for marker in ("目录", "目次", "contents", "tableofcontents")) - - -def _extract_candidates(page: int, text: str) -> list[TocCandidate]: - candidates: list[TocCandidate] = [] - seen: set[str] = set() - for idx, raw_line in enumerate(text.splitlines()): - line = clean_toc_line(raw_line) - if not looks_like_h1_line(line): - continue - if not candidate_allowed(line): - continue - normalized = normalize_heading(line) - if normalized in seen: - continue - seen.add(normalized) - candidates.append( - TocCandidate( - title=line, - normalized_title=normalized, - source_page=page, - line_index=idx, - numbering=line[: max(len(line) - len(normalized), 0)].strip(), - ) - ) - return candidates - - -@register_tool( - name="find.toc_pages", - description="Find table-of-contents pages and extract level-1 heading candidates.", - allowed_states={DocumentAgentState.PROBED}, -) -def find_toc_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: - start = time.monotonic() - toc_pages = [ - feature.page - for feature in ctx.blackboard.page_features - if _is_toc_feature_lines(feature.text_lines_preview) - ] - texts = read_page_texts(ctx.pdf_path, toc_pages) if toc_pages else {} - candidates: list[TocCandidate] = [] - for page, text in texts.items(): - candidates.extend(_extract_candidates(page, text)) - result = TocResult( - toc_pages=sorted(toc_pages), - candidates=candidates, - method="toc_marker" if toc_pages else "none", - notes=f"{len(candidates)} h1 candidates extracted", - ) - ctx.blackboard.toc_result = result - ctx.blackboard.global_signals["toc_page_count"] = len(toc_pages) - ctx.blackboard.global_signals["toc_candidate_count"] = len(candidates) - return ToolResult( - status="ok", - payload={ - "toc_pages": sorted(toc_pages), - "candidate_count": len(candidates), - }, - latency_ms=int((time.monotonic() - start) * 1000), - ) diff --git a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py index 017504d2c..48ed6ea66 100644 --- a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py @@ -38,6 +38,7 @@ def build_anatomy_map(ctx: ToolContext) -> PageAnatomyMap: h1_result=ctx.blackboard.h1_result, hierarchy_assist=ctx.blackboard.hierarchy_assist, shard_plan=ctx.blackboard.shard_plan, + boundary_candidates=ctx.blackboard.boundary_candidates, page_processing_plan=None, global_signals=ctx.blackboard.global_signals, trace_summary={ diff --git a/apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py b/apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py index 7f11331fb..50cd350bf 100644 --- a/apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py +++ b/apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py @@ -2,7 +2,6 @@ from __future__ import annotations -import json import time from typing import Any @@ -17,7 +16,6 @@ from app.services.document_agent.registry import register_tool from app.services.document_agent.state import DocumentAgentState from app.services.document_agent.validators import repair_hierarchy_assist -from shared.utils.token_estimate import estimate_tokens def _rule_based_plan(labels: list[PageLabel], h1_candidates: list[H1Candidate]) -> HierarchyAssistPlan: @@ -56,28 +54,6 @@ def _rule_based_plan(labels: list[PageLabel], h1_candidates: list[H1Candidate]) ) -def _parse_llm_plan(raw: str, fallback: HierarchyAssistPlan) -> HierarchyAssistPlan: - try: - data = json.loads(raw) - except json.JSONDecodeError: - return fallback - if not isinstance(data, dict): - return fallback - plan = fallback - plan.exclude_pages_from_title_candidates = [ - int(page) for page in data.get("exclude_pages_from_title_candidates", plan.exclude_pages_from_title_candidates) - ] - plan.suppress_title_pages = [ - int(page) for page in data.get("suppress_title_pages", plan.suppress_title_pages) - ] - recommendation = data.get("smart_parse_recommendation") - if recommendation in {"off", "normal", "aggressive"}: - plan.smart_parse_recommendation = recommendation - if isinstance(data.get("rationale"), str): - plan.rationale = data["rationale"] - return plan - - @register_tool( name="propose.hierarchy_assist", description="Produce hierarchy hints used by section skeleton extraction and title parsing.", @@ -89,49 +65,7 @@ def propose_hierarchy_assist(ctx: ToolContext, _args: dict[str, Any]) -> ToolRes labels = ctx.blackboard.page_labels fallback = _rule_based_plan(labels, h1_result.h1_candidates if h1_result else []) toc_pages = set(ctx.blackboard.toc_result.toc_pages if ctx.blackboard.toc_result else []) - - model = ctx.settings.get("model") - if model: - payload = { - "page_count": ctx.blackboard.page_count, - "page_kind_counts": ctx.blackboard.global_signals.get("page_kind_counts", {}), - "toc_pages": sorted(toc_pages), - "h1_candidates": [ - candidate.to_dict() for candidate in (h1_result.h1_candidates if h1_result else []) - ], - "candidate_exclude_pages": fallback.exclude_pages_from_title_candidates, - "candidate_suppress_pages": fallback.suppress_title_pages, - } - prompt = ( - "Return strict JSON for hierarchy assistance. " - "Use only page numbers present in the payload. " - "Do not invent headings. " - "Fields: exclude_pages_from_title_candidates, suppress_title_pages, " - "smart_parse_recommendation, rationale.\n" - + json.dumps(payload, ensure_ascii=False) - ) - est = estimate_tokens(prompt) - if ctx.budget.try_reserve("plan", est): - try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - client = get_openai_client(model=model) - raw, usage = client.chat_completion_with_usage( - messages=prompt, - model=model, - temperature=0.0, - max_tokens=1200, - response_format={"type": "json_object"}, - ) - ctx.budget.commit("plan", actual=usage.get("total_tokens", est), est=est) - plan = _parse_llm_plan(raw, fallback) - except Exception: - ctx.budget.refund("plan", est=est) - plan = fallback - else: - plan = fallback - else: - plan = fallback + plan = fallback repaired = repair_hierarchy_assist( plan, @@ -147,4 +81,17 @@ def propose_hierarchy_assist(ctx: ToolContext, _args: dict[str, Any]) -> ToolRes "h1_hint_count": len(repaired.prefer_h1_start_pages), }, latency_ms=int((time.monotonic() - start) * 1000), + input_summary={ + "page_kind_counts": ctx.blackboard.global_signals.get("page_kind_counts", {}), + "boundary_candidate_counts": ctx.blackboard.global_signals.get( + "boundary_candidate_counts", {} + ), + }, + output_summary={ + "exclude_pages": repaired.exclude_pages_from_title_candidates[:50], + "suppress_pages_count": len(repaired.suppress_title_pages), + "h1_hint_count": len(repaired.prefer_h1_start_pages), + "smart_parse_recommendation": repaired.smart_parse_recommendation, + "rationale": repaired.rationale, + }, ) 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 b3a158405..91836936c 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -1,15 +1,24 @@ -"""Rule-based long-PDF shard planning.""" +"""LLM-guided long-PDF shard planning from candidate split pages.""" from __future__ import annotations +import json import os +import re import time from typing import Any -from app.services.document_agent.manifest import Shard, ShardPlan, ToolContext, ToolResult +from app.services.document_agent.manifest import ( + BoundaryCandidate, + Shard, + ShardPlan, + ToolContext, + ToolResult, +) from app.services.document_agent.registry import register_tool from app.services.document_agent.state import DocumentAgentState from app.services.document_agent.validators import single_shard_plan, validate_shard_plan +from shared.utils.token_estimate import estimate_tokens def _thresholds(ctx: ToolContext) -> tuple[int, int, int]: @@ -28,33 +37,6 @@ def _thresholds(ctx: ToolContext) -> tuple[int, int, int]: return threshold, min_pages, max_pages -def _avoid_pages(ctx: ToolContext) -> set[int]: - return { - label.page - for label in ctx.blackboard.page_labels - if label.kind in {"table_heavy", "landscape"} - } - - -def _separator_pages(ctx: ToolContext) -> set[int]: - return { - label.page - for label in ctx.blackboard.page_labels - if label.kind in {"blank", "separator", "sparse"} - } - - -def _nearest_safe_cut(target: int, previous: int, page_count: int, avoid: set[int], separators: set[int]) -> tuple[int, str]: - window = range(target, max(previous, target - 12), -1) - for page in window: - if previous < page < page_count and page in separators and page not in avoid: - return page, "blank_separator" - for page in range(target, max(previous, target - 5), -1): - if previous < page < page_count and page not in avoid: - return page, "forced_max_size" - return max(previous + 1, min(target, page_count - 1)), "forced_max_size" - - def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> list[Shard]: shards: list[Shard] = [] previous = 0 @@ -88,9 +70,158 @@ def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> return shards +def _compact_candidate(candidate: BoundaryCandidate, page_count: int) -> dict[str, Any]: + evidence = candidate.evidence or {} + return { + "page": candidate.page, + "kind": candidate.kind, + "priority": candidate.priority, + "confidence": candidate.confidence, + "position_ratio": round(candidate.page / max(page_count, 1), 4), + "raw_text_length": evidence.get("raw_text_length"), + "image_coverage": evidence.get("image_coverage"), + "table_count": evidence.get("table_count"), + "drawings_count": evidence.get("drawings_count"), + "text_preview": evidence.get("text_preview", [])[:3], + "title": evidence.get("title"), + } + + +def _build_prompt( + *, + page_count: int, + min_pages: int, + max_pages: int, + candidates: list[BoundaryCandidate], + page_kind_counts: dict[str, int], +) -> str: + payload = { + "page_count": page_count, + "min_pages_per_shard": min_pages, + "max_pages_per_shard": max_pages, + "page_kind_counts": page_kind_counts, + "candidate_priority": { + "h1": "highest semantic priority, but still decide using document size and spacing", + "toc": "high priority marker, usually not a cut by itself unless it indicates nearby structure", + "separator": "explicit sparse separator page", + "blank": "sparse/blank structural gap candidate", + "sparse": "low-density candidate, useful when no stronger signal exists", + }, + "candidates": [ + _compact_candidate(candidate, page_count) for candidate in candidates + ], + } + return ( + "You are a senior document parsing architect. Decide whether to split a PDF " + "and where to split it using only the provided candidate pages and document-scale " + "features.\n" + "Rules:\n" + "- Return strict JSON only.\n" + "- Do not invent pages. Every cut_after_page must be one of the candidate pages, " + "or candidate_page - 1 when the candidate is a semantic start page such as h1.\n" + "- H1 candidates have the highest semantic priority, but do not blindly split on " + "every H1. Consider total page_count, candidate spacing, min/max shard sizes, and " + "whether splitting would over-fragment the document.\n" + "- Blank and sparse pages are valid split candidates because they often mark section " + "gaps, especially when TOC/H1 evidence is weak or absent.\n" + "- Prefer fewer, semantically coherent shards over many tiny shards.\n" + "- Do not use domain-specific hardcoded labels or examples; decide from the supplied " + "features and positions only. Do not quote business/category words from text_preview " + "in rationale; refer to them generically as sparse separator text.\n" + "- Every resulting shard length must be <= max_pages_per_shard unless enabled=false. " + "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\" | \"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 _sanitize_rationale(text: str) -> str: + # Keep rationales structural. The model may quote page preview text; those + # literals are useful as input evidence but should not become baked-in rules. + sanitized = re.sub(r"'[^']{1,40}'", "sparse separator text", text or "") + sanitized = re.sub(r'"[^"]{1,40}"', "sparse separator text", sanitized) + return sanitized + + +def _validate_cut_lengths(cuts: list[tuple[int, str, str, float]], page_count: int, max_pages: int) -> None: + previous = 0 + for cut_page, *_ in cuts: + if cut_page - previous > max_pages: + raise ValueError( + f"LLM cut plan creates shard length {cut_page - previous} > max_pages={max_pages}" + ) + previous = cut_page + if page_count - previous > max_pages: + raise ValueError( + f"LLM cut plan creates final shard length {page_count - previous} > max_pages={max_pages}" + ) + + +def _parse_llm_plan(raw: str, page_count: int, max_pages: int) -> tuple[bool, list[tuple[int, str, str, float]], str, str]: + data = json.loads(raw) + if not isinstance(data, dict): + raise ValueError("LLM shard plan is not an object") + enabled = bool(data.get("enabled")) + reason = str(data.get("reason") or ("llm_boundary_decision" if enabled else "not_needed")) + rationale = _sanitize_rationale(str(data.get("rationale") or "")) + cuts: list[tuple[int, str, str, float]] = [] + for item in data.get("cuts") or []: + if not isinstance(item, dict): + continue + cut_page = int(item.get("cut_after_page")) + if not 1 <= cut_page < page_count: + continue + anchor_type = str(item.get("anchor_type") or "separator") + if anchor_type not in {"h1_boundary", "blank_separator", "separator", "forced_max_size"}: + anchor_type = "separator" + confidence = float(item.get("confidence") or 0.5) + cuts.append((cut_page, anchor_type, _sanitize_rationale(str(item.get("rationale") or rationale)), confidence)) + cuts = sorted({cut[0]: cut for cut in cuts}.values(), key=lambda cut: cut[0]) + if enabled: + _validate_cut_lengths(cuts, page_count, max_pages) + return enabled, cuts, reason, rationale + + +def _deterministic_guardrail_plan( + *, + page_count: int, + max_pages: int, + candidates: list[BoundaryCandidate], +) -> tuple[list[tuple[int, str, str, float]], str]: + cuts: list[tuple[int, str, str, float]] = [] + previous = 0 + while page_count - previous > max_pages: + target = previous + max_pages + eligible = [ + candidate for candidate in candidates if previous < candidate.page <= target + ] + if eligible: + chosen = max(eligible, key=lambda item: (item.priority, item.page)) + cut_page = chosen.page - 1 if chosen.kind == "h1" and chosen.page > previous + 1 else chosen.page + anchor_type = "h1_boundary" if chosen.kind == "h1" else "blank_separator" + cuts.append((cut_page, anchor_type, f"guardrail candidate {chosen.kind} at page {chosen.page}", 0.35)) + previous = cut_page + else: + cuts.append((target, "forced_max_size", "guardrail max shard size", 0.25)) + previous = target + return cuts, "too_large" + + @register_tool( name="propose.shard_plan", - description="Create a rule-based PDF segment plan for long-document PDF-to-MD execution.", + description="Ask the LLM to decide whether and where to split using candidate boundary pages.", allowed_states={DocumentAgentState.H1_FOUND}, ) def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: @@ -106,58 +237,69 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) - avoid = _avoid_pages(ctx) - separators = _separator_pages(ctx) - h1_candidates = [] - if ctx.blackboard.hierarchy_assist: - h1_candidates = ctx.blackboard.hierarchy_assist.prefer_h1_start_pages - elif ctx.blackboard.h1_result: - h1_candidates = ctx.blackboard.h1_result.h1_candidates + candidates = list(ctx.blackboard.boundary_candidates) + model = ctx.settings.get("model") + prompt = _build_prompt( + page_count=page_count, + min_pages=min_pages, + max_pages=max_pages, + candidates=candidates, + page_kind_counts=ctx.blackboard.global_signals.get("page_kind_counts", {}), + ) + prompt_tokens_est = estimate_tokens(prompt) + warnings: list[str] = [] + raw_response = "" + rationale = "" + llm_attempted = False + if model and ctx.budget.try_reserve("plan", prompt_tokens_est): + try: + llm_attempted = True + from shared.services.ai.openai_compatible_client_sync import get_openai_client - cuts: list[tuple[int, str, str, float]] = [] - previous = 0 - for candidate in sorted(h1_candidates, key=lambda item: item.page): - cut_page = candidate.page - 1 - if cut_page <= previous: - continue - if cut_page - previous < min_pages: - continue - while cut_page - previous > max_pages: - forced_target = previous + max_pages - safe_cut, anchor_type = _nearest_safe_cut( - forced_target, - previous, - page_count, - avoid, - separators, + client = get_openai_client(model=model) + raw_response, usage = client.chat_completion_with_usage( + messages=prompt, + model=model, + temperature=0.0, + max_tokens=1600, + response_format={"type": "json_object"}, ) - cuts.append((safe_cut, anchor_type, "max shard size guard", 0.58)) - previous = safe_cut - if cut_page > previous and cut_page - previous >= min_pages: - cuts.append( - ( - cut_page, - "h1_boundary", - f"h1 starts at page {candidate.page}: {candidate.title}", - candidate.confidence, - ) + ctx.budget.commit("plan", actual=usage.get("total_tokens", prompt_tokens_est), est=prompt_tokens_est) + enabled, cuts, reason, rationale = _parse_llm_plan(raw_response, page_count, max_pages) + if not enabled: + cuts = [] + reason = "not_needed" + except Exception as exc: + ctx.budget.refund("plan", est=prompt_tokens_est) + warnings.append(f"LLM shard decision rejected, using guardrail plan: {exc}") + cuts, reason = _deterministic_guardrail_plan( + page_count=page_count, + max_pages=max_pages, + candidates=candidates, ) - previous = cut_page - - while page_count - previous > max_pages: - target = previous + max_pages - safe_cut, anchor_type = _nearest_safe_cut(target, previous, page_count, avoid, separators) - cuts.append((safe_cut, anchor_type, "tail max shard size guard", 0.58)) - previous = safe_cut + rationale = "Guardrail plan after LLM shard decision failure." + else: + if not model: + warnings.append("No model configured for shard decision; using guardrail plan.") + else: + warnings.append("Insufficient plan budget for shard decision; using guardrail plan.") + cuts, reason = _deterministic_guardrail_plan( + page_count=page_count, + max_pages=max_pages, + candidates=candidates, + ) + rationale = "Guardrail plan without LLM decision." shards = _cuts_to_shards(cuts, page_count) - reason = "hierarchy_isolation" if any(cut[1] == "h1_boundary" for cut in cuts) else "too_large" + enabled = len(shards) > 1 + if not enabled: + reason = "not_needed" plan = ShardPlan( - enabled=True, + enabled=enabled, reason=reason, # type: ignore[arg-type] shards=shards, validation=validate_shard_plan( - ShardPlan(enabled=True, reason=reason, shards=shards), # type: ignore[arg-type] + ShardPlan(enabled=enabled, reason=reason, shards=shards), # type: ignore[arg-type] page_count=page_count, min_pages=min_pages, max_pages=max_pages, @@ -173,4 +315,23 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: "valid": plan.validation.valid, }, latency_ms=int((time.monotonic() - start) * 1000), + tokens_used=ctx.budget.snapshot()["plan"]["used"], + input_summary={ + "page_count": page_count, + "candidate_count": len(candidates), + "candidate_counts": ctx.blackboard.global_signals.get("boundary_candidate_counts", {}), + "model": model, + }, + output_summary={ + "enabled": plan.enabled, + "reason": plan.reason, + "rationale": rationale, + "shards": [shard.to_dict() for shard in plan.shards], + }, + warnings=warnings, + debug={ + "prompt_excerpt": prompt[:4000], + "raw_response_excerpt": raw_response[:4000], + "llm_attempted": llm_attempted, + }, ) diff --git a/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py index fe4459a5b..d37811853 100644 --- a/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py @@ -52,6 +52,7 @@ def validate_current_anatomy(ctx: ToolContext, _args: dict[str, Any]) -> ToolRes h1_result=ctx.blackboard.h1_result, hierarchy_assist=ctx.blackboard.hierarchy_assist, shard_plan=ctx.blackboard.shard_plan, + boundary_candidates=ctx.blackboard.boundary_candidates, global_signals=ctx.blackboard.global_signals, trace_summary={}, ) diff --git a/apps/worker/app/services/document_agent/trace.py b/apps/worker/app/services/document_agent/trace.py index 02a03e364..30f9b46b6 100644 --- a/apps/worker/app/services/document_agent/trace.py +++ b/apps/worker/app/services/document_agent/trace.py @@ -4,6 +4,7 @@ import time from datetime import datetime +from pathlib import Path from typing import Any from uuid import uuid4 @@ -42,6 +43,11 @@ def record_step( "observation": { "status": result.status, "payload_keys": sorted(result.payload.keys()), + "payload": result.payload, + "input_summary": result.input_summary, + "output_summary": result.output_summary, + "warnings": list(result.warnings), + "debug": result.debug, "error": result.error, }, "tokens_used": result.tokens_used, @@ -53,6 +59,35 @@ def record_step( def set_anatomy_map(self, anatomy: PageAnatomyMap, artifact_path: str) -> None: self._anatomy = anatomy self._artifact_path = artifact_path + self.write_trace_json(str(Path(artifact_path).with_name("trace.json"))) + + def write_trace_json(self, trace_path: str) -> None: + try: + import json + + serializable_steps = [] + for step in self._steps: + item = dict(step) + created_at = item.get("created_at") + if hasattr(created_at, "isoformat"): + item["created_at"] = created_at.isoformat() + serializable_steps.append(item) + Path(trace_path).write_text( + json.dumps( + { + "run_id": self.run_id, + "job_id": self.job_id, + "artifact_path": self._artifact_path, + "steps": serializable_steps, + }, + ensure_ascii=False, + indent=2, + default=str, + ), + encoding="utf-8", + ) + except Exception as exc: + logger.debug(f"parse agent trace json write failed: {exc}") def summary(self) -> dict[str, Any]: return { diff --git a/apps/worker/app/services/document_agent/validators.py b/apps/worker/app/services/document_agent/validators.py index 44e3cf537..1a5da7777 100644 --- a/apps/worker/app/services/document_agent/validators.py +++ b/apps/worker/app/services/document_agent/validators.py @@ -110,6 +110,17 @@ def validate_anatomy_map( errors.append(f"h1 candidate points to toc page {candidate.page}") if candidate.page < 1 or candidate.page > page_count: errors.append(f"h1 candidate page {candidate.page} out of range") + candidate_count = len(anatomy.boundary_candidates) + if page_count > max_pages and candidate_count == 0: + warnings.append("long document has no boundary candidates") + h1_candidate_count = sum(1 for candidate in anatomy.boundary_candidates if candidate.kind == "h1") + sparse_candidate_count = sum( + 1 + for candidate in anatomy.boundary_candidates + if candidate.kind in {"blank", "sparse", "separator"} + ) + if page_count > max_pages and h1_candidate_count == 0 and sparse_candidate_count == 0: + warnings.append("long document has neither H1 nor sparse/blank boundary candidates") shard_report = validate_shard_plan( anatomy.shard_plan, page_count=page_count, @@ -118,4 +129,12 @@ def validate_anatomy_map( ) errors.extend(shard_report.errors) warnings.extend(shard_report.warnings) + if anatomy.shard_plan.enabled: + forced_count = sum( + 1 + for shard in anatomy.shard_plan.shards + if shard.anchor_type == "forced_max_size" + ) + if forced_count == len(anatomy.shard_plan.shards): + warnings.append("all shards are based on forced max-size boundaries") return ValidationReport(valid=not errors, errors=errors, warnings=warnings) From bf519da9f09c0c941131188b6949ec96be862c27 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 25 May 2026 15:30:42 +0800 Subject: [PATCH 05/11] feat: implement H1-to-page matching tools and integrate TOC hierarchies into document agent state. --- .../services/document_agent/coordinator.py | 8 +- .../services/document_agent/heading_text.py | 83 -- .../app/services/document_agent/manifest.py | 18 +- .../app/services/document_agent/state.py | 4 + .../services/document_agent/tools/__init__.py | 3 + .../tools/classify_page_kinds.py | 19 +- .../tools/collect_boundary_candidates.py | 25 +- .../tools/extract_toc_with_boundaries.py | 529 ++++++++++++ .../tools/find_toc_anchor_pages.py | 217 +++++ .../document_agent/tools/match_h1_pages.py | 196 +++++ .../tools/persist_anatomy_map.py | 1 + .../tools/test_vlm_toc_extract.py | 817 ++++++++++++++++++ 12 files changed, 1800 insertions(+), 120 deletions(-) delete mode 100644 apps/worker/app/services/document_agent/heading_text.py create mode 100644 apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py create mode 100644 apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py create mode 100644 apps/worker/app/services/document_agent/tools/match_h1_pages.py create mode 100644 apps/worker/app/services/document_agent/tools/test_vlm_toc_extract.py diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 1f5fa0968..5330b0fa6 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -19,6 +19,9 @@ TRANSITIONS: dict[str, DocumentAgentState] = { "probe.page_features": DocumentAgentState.PROBED, "classify.page_kinds": DocumentAgentState.PROBED, + "find.toc_anchor_pages": DocumentAgentState.CLASSIFIED, + "extract.toc_with_boundaries": DocumentAgentState.CLASSIFIED, + "match.h1_pages": DocumentAgentState.H1_FOUND, "collect.boundary_candidates": DocumentAgentState.H1_FOUND, "propose.hierarchy_assist": DocumentAgentState.H1_FOUND, "propose.shard_plan": DocumentAgentState.H1_FOUND, @@ -29,6 +32,9 @@ REQUIRED_TOOLS = [ "probe.page_features", "classify.page_kinds", + "find.toc_anchor_pages", + "extract.toc_with_boundaries", + "match.h1_pages", "collect.boundary_candidates", "propose.hierarchy_assist", "propose.shard_plan", @@ -77,7 +83,7 @@ def __init__( self.blackboard = AgentBlackboard() self.budget = BudgetTracker( plan_budget=int(os.environ.get("PARSE_AGENT_PLAN_BUDGET", "5000")), - max_tool_calls=int(os.environ.get("PARSE_AGENT_MAX_TOOL_CALLS", "12")), + max_tool_calls=int(os.environ.get("PARSE_AGENT_MAX_TOOL_CALLS", "15")), ) effective_settings = settings or {} if model: diff --git a/apps/worker/app/services/document_agent/heading_text.py b/apps/worker/app/services/document_agent/heading_text.py deleted file mode 100644 index cae0f0d3b..000000000 --- a/apps/worker/app/services/document_agent/heading_text.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Heading extraction and matching helpers.""" - -from __future__ import annotations - -import re -import unicodedata - - -LEADING_NUMBER_RE = re.compile( - r"""^ - (?: - 第\s*[零一二三四五六七八九十百千\d]+\s*[章节篇部分] - | [零一二三四五六七八九十百千]+\s*[、。，,] - | [（(]\s*[零一二三四五六七八九十百千\d]+\s*[）)] - | \d+(?:\.\d+)*\.?\s* - | [IVXLCDM]+\.?\s+ - | [A-Za-z]\.\s+ - | Chapter\s+\w+\s* - ) - """, - re.VERBOSE | re.IGNORECASE, -) - -H1_LINE_RE = re.compile( - r"""^\s* - (?: - 第\s*[零一二三四五六七八九十百千\d]+\s*[章篇部] - | [零一二三四五六七八九十百千]+\s*[、。] - | \d+\s*[\.\s] - | Chapter\s+\w+ - | [IVXLCDM]+\.?\s+\w - ) - """, - re.VERBOSE | re.IGNORECASE, -) - -SUBHEADING_RE = re.compile(r"^\s*\d+\.\d+", re.IGNORECASE) -PAGE_SUFFIX_RE = re.compile(r"[\s\.\-·…]+\d+\s*$") - - -def normalize_heading(text: str) -> str: - text = unicodedata.normalize("NFKC", text or "") - text = re.sub(r"\s+", " ", text).strip() - stripped = LEADING_NUMBER_RE.sub("", text).strip() - return stripped if stripped else text - - -def has_numbering(text: str) -> bool: - return bool(LEADING_NUMBER_RE.match(text or "")) - - -def clean_toc_line(line: str) -> str: - return PAGE_SUFFIX_RE.sub("", line or "").strip() - - -def looks_like_h1_line(line: str) -> bool: - stripped = (line or "").strip() - if not stripped or SUBHEADING_RE.match(stripped): - return False - return bool(H1_LINE_RE.match(stripped)) - - -def candidate_allowed(title: str) -> bool: - normalized = normalize_heading(title) - if not normalized: - return False - if len(normalized) < 4 and not has_numbering(title): - return False - return True - - -def fuzzy_match(needle: str, haystack: str) -> bool: - if not needle: - return False - normalized_haystack = re.sub( - r"\s+", - " ", - unicodedata.normalize("NFKC", haystack or ""), - ) - if needle in normalized_haystack: - return True - stripped = normalize_heading(needle) - return bool(stripped and len(stripped) >= 4 and stripped in normalized_haystack) diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index 58aa13c7f..9bb5ea12a 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -15,12 +15,10 @@ "chapter_start", "section_start", "table_heavy", - "image_heavy", "single_image", "blank", "separator", "appendix", - "scan_like", "landscape", "sparse", ] @@ -76,11 +74,23 @@ def to_dict(self) -> dict[str, Any]: return asdict(self) +@dataclass +class TocAnchorPage: + """A candidate TOC start page identified by keyword scan, pending VLM confirmation.""" + + page: int # 1-based page number + png_path: str # local PNG path for VLM inspection + source: Literal["page_label", "text_scan"] # how this anchor was discovered + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + @dataclass class TocResult: toc_pages: list[int] = field(default_factory=list) candidates: list[TocCandidate] = field(default_factory=list) - method: Literal["toc_marker", "none"] = "none" + method: Literal["toc_marker", "vlm_progressive", "none"] = "none" notes: str = "" def to_dict(self) -> dict[str, Any]: @@ -249,6 +259,7 @@ class PageAnatomyMap: hierarchy_assist: HierarchyAssistPlan shard_plan: ShardPlan boundary_candidates: list[BoundaryCandidate] = field(default_factory=list) + toc_hierarchies: list[dict[str, Any]] | None = None page_processing_plan: PageProcessingPlan | None = None global_signals: dict[str, Any] = field(default_factory=dict) trace_summary: dict[str, Any] = field(default_factory=dict) @@ -270,6 +281,7 @@ def to_dict(self) -> dict[str, Any]: "boundary_candidates": [ candidate.to_dict() for candidate in self.boundary_candidates ], + "toc_hierarchies": self.toc_hierarchies, "page_processing_plan": ( self.page_processing_plan.to_dict() if self.page_processing_plan is not None diff --git a/apps/worker/app/services/document_agent/state.py b/apps/worker/app/services/document_agent/state.py index a752f129b..f54278fa6 100644 --- a/apps/worker/app/services/document_agent/state.py +++ b/apps/worker/app/services/document_agent/state.py @@ -13,6 +13,7 @@ PageFeature, PageLabel, ShardPlan, + TocAnchorPage, TocResult, ) @@ -35,7 +36,9 @@ class AgentBlackboard: page_count: int = 0 page_features: list[PageFeature] = field(default_factory=list) page_labels: list[PageLabel] = field(default_factory=list) + toc_anchor_pages: list[TocAnchorPage] = field(default_factory=list) toc_result: TocResult | None = None + toc_hierarchies: list[dict[str, Any]] | None = None h1_result: H1BoundaryResult | None = None boundary_candidates: list[BoundaryCandidate] = field(default_factory=list) hierarchy_assist: HierarchyAssistPlan | None = None @@ -47,3 +50,4 @@ class AgentBlackboard: def mark(self, state: DocumentAgentState) -> None: self.state_trace.append(state.value) + diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index 52084ae75..d212eaa14 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -4,6 +4,9 @@ from . import classify_page_kinds as classify_page_kinds # noqa: F401 from . import collect_boundary_candidates as collect_boundary_candidates # noqa: F401 +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 match_h1_pages as match_h1_pages # noqa: F401 from . import persist_anatomy_map as persist_anatomy_map # noqa: F401 from . import probe_page_features as probe_page_features # noqa: F401 from . import propose_hierarchy_assist as propose_hierarchy_assist # noqa: F401 diff --git a/apps/worker/app/services/document_agent/tools/classify_page_kinds.py b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py index 2305652c1..44fb98f41 100644 --- a/apps/worker/app/services/document_agent/tools/classify_page_kinds.py +++ b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py @@ -39,23 +39,13 @@ def _label_feature(feature: PageFeature) -> PageLabel: confidence=0.78, evidence={"width": feature.width, "height": feature.height}, ) - if feature.image_coverage >= 0.72 and feature.raw_text_length < 250: + if feature.image_coverage >= 0.35 and feature.raw_text_length < 250: return PageLabel( page=page, kind="single_image", confidence=0.84, evidence={"image_coverage": feature.image_coverage}, ) - if feature.raw_text_length < 50 and feature.image_coverage >= 0.35: - return PageLabel( - page=page, - kind="scan_like", - confidence=0.76, - evidence={ - "raw_text_length": feature.raw_text_length, - "image_coverage": feature.image_coverage, - }, - ) if feature.table_count > 0 or feature.drawings_count >= 80: return PageLabel( page=page, @@ -66,13 +56,6 @@ def _label_feature(feature: PageFeature) -> PageLabel: "drawings_count": feature.drawings_count, }, ) - if feature.image_coverage >= 0.35: - return PageLabel( - page=page, - kind="image_heavy", - confidence=0.72, - evidence={"image_coverage": feature.image_coverage}, - ) if feature.raw_text_length < 80: return PageLabel( page=page, diff --git a/apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py b/apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py index aec0cd279..89e72610f 100644 --- a/apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py +++ b/apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py @@ -17,11 +17,10 @@ from app.services.document_agent.state import DocumentAgentState -BOUNDARY_PAGE_KINDS = {"blank", "sparse", "separator", "toc"} +BOUNDARY_PAGE_KINDS = {"blank", "sparse", "toc"} PRIORITY_BY_KIND = { "h1": 100, "toc": 80, - "separator": 70, "blank": 45, "sparse": 40, } @@ -50,23 +49,19 @@ def _candidate_evidence(ctx: ToolContext, page: int, kind: str) -> dict[str, Any @register_tool( name="collect.boundary_candidates", description="Collect sparse, blank, TOC, and H1 pages as split candidates without making split decisions.", - allowed_states={DocumentAgentState.CLASSIFIED}, + allowed_states={DocumentAgentState.H1_FOUND}, ) def collect_boundary_candidates(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() candidates: list[BoundaryCandidate] = [] seen: set[tuple[int, str]] = set() - # The robust TOC parser will populate these later. Until then, keep the - # anatomy contract explicit without inventing TOC/H1 results here. - ctx.blackboard.toc_result = ctx.blackboard.toc_result or TocResult( - method="none", - notes="TOC parser not run in boundary candidate collection", - ) - ctx.blackboard.h1_result = ctx.blackboard.h1_result or H1BoundaryResult( - method="none", - notes="H1 parser not run in boundary candidate collection", - ) + # Upstream tools (extract.toc_with_boundaries, match.h1_pages) populate + # these before collect runs. Provide safe defaults if they were skipped. + if ctx.blackboard.toc_result is None: + ctx.blackboard.toc_result = TocResult(method="none") + if ctx.blackboard.h1_result is None: + ctx.blackboard.h1_result = H1BoundaryResult(method="none") for label in ctx.blackboard.page_labels: if label.kind not in BOUNDARY_PAGE_KINDS: @@ -85,8 +80,8 @@ def collect_boundary_candidates(ctx: ToolContext, _args: dict[str, Any]) -> Tool ) ) - # Preserve future H1 candidates as highest-priority hints when upstream TOC - # parsing is wired in, but do not derive them with the old fuzzy matcher. + # Add H1 candidates from upstream match.h1_pages as highest-priority + # boundary hints for shard planning. if ctx.blackboard.h1_result: for h1 in ctx.blackboard.h1_result.h1_candidates: key = (h1.page, "h1") diff --git a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py new file mode 100644 index 000000000..fe0dc617c --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py @@ -0,0 +1,529 @@ +"""VLM-driven progressive TOC boundary detection + mineru local MD + toc_parser reuse.""" + +from __future__ import annotations + +import gc +import json +import os +import subprocess +import time +from pathlib import Path +from typing import Any + +from app.services.document_agent.manifest import ( + TocAnchorPage, + TocResult, + ToolContext, + ToolResult, +) +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) +from loguru import logger + +# -- Constants ----------------------------------------------------------------- + +BOUNDARY_STEP_PAGES = 5 +MAX_BOUNDARY_ROUNDS = 6 +MAX_TOC_PAGES = BOUNDARY_STEP_PAGES * MAX_BOUNDARY_ROUNDS # 30 +MINERU_TIMEOUT_SECONDS = 180 + + +# -- PyMuPDF workers (must be top-level for multiprocessing pickle) ------------ + + +@worker +def _render_single_page_worker( + queue, pdf_path: str, page_num: int, output_path: str, dpi: int +) -> None: + import pymupdf # type: ignore[import] + + try: + doc = pymupdf.open(pdf_path) + idx = page_num - 1 + if 0 <= idx < doc.page_count: + page = doc[idx] + mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) + pix = page.get_pixmap(matrix=mat) + pix.save(output_path) + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "png_path": output_path}) + + +# -- VLM helpers --------------------------------------------------------------- + + +def _vlm_confirm_anchors( + anchor_pages: list[TocAnchorPage], + model: str, +) -> list[TocAnchorPage]: + """Phase 1: send all anchor PNGs to VLM, ask which are real TOC starts.""" + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + if not anchor_pages: + return [] + + import base64 + + # Build multi-image message + content_parts: list[dict[str, Any]] = [ + { + "type": "text", + "text": ( + "You are a document structure analysis expert. " + "Below are screenshot(s) of candidate pages extracted from a PDF. " + "These pages contained keywords such as 'Table of Contents' / 'Contents' " + "during a text scan.\n\n" + "For each page, determine whether it is truly the **start page** of a " + "Table of Contents (TOC).\n\n" + "Criteria for a real TOC page:\n" + "- Contains a list of section titles paired with page numbers\n" + "- Titles are connected to page numbers via dots, ellipses, or spaces\n" + "- Titles have a systematic numbering scheme (e.g. 1. / 1.1 / Chapter 1)\n\n" + "NOT a TOC page:\n" + "- Body text that casually mentions 'contents'\n" + "- A page with only a 'Contents' heading but body text below\n\n" + "Return a strict JSON array (no markdown fences):\n" + '[{"page": , "is_toc_start": true/false, "reason": "brief reason"}]' + ), + } + ] + + for anchor in anchor_pages: + with open(anchor.png_path, "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + content_parts.append( + { + "type": "text", + "text": f"\n--- Page {anchor.page} ---", + } + ) + content_parts.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + } + ) + + messages = [{"role": "user", "content": content_parts}] + + try: + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=messages, + model=model, + temperature=0.1, + max_tokens=500, + response_format={"type": "json_object"}, + ) + data = json.loads(raw) + if isinstance(data, dict): + items = data.get("pages") or data.get("results") or data.get("data") or [] + if not items and len(data) == 1: + items = list(data.values())[0] + elif isinstance(data, list): + items = data + else: + items = [] + + confirmed_pages: set[int] = set() + for item in items: + if isinstance(item, dict) and item.get("is_toc_start"): + confirmed_pages.add(int(item["page"])) + + confirmed = [a for a in anchor_pages if a.page in confirmed_pages] + rejected = [a.page for a in anchor_pages if a.page not in confirmed_pages] + logger.info( + "[extract.toc] VLM confirmed {} TOC starts, rejected pages: {}", + len(confirmed), + rejected, + ) + return confirmed + except Exception as exc: + logger.warning( + "[extract.toc] VLM anchor confirmation failed: {}, " + "falling back to no confirmed anchors (safe degradation)", + exc, + ) + return [] + + +def _vlm_check_boundary_page( + png_path: str, + page_num: int, + model: str, +) -> bool: + """Phase 2: check if a single page still contains TOC content.""" + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + import base64 + + with open(png_path, "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + + content_parts: list[dict[str, Any]] = [ + { + "type": "text", + "text": ( + "You are a document structure analysis expert. " + "Below is a screenshot of one page from a PDF. " + "I am determining the boundary of a Table of Contents (TOC) region.\n\n" + "Does this page still contain TOC content?\n\n" + "TOC content characteristics:\n" + "- Entry titles paired with page numbers\n" + "- Dots (...), leader lines (.....), or spaces connecting titles to page numbers\n" + "- Systematic numbering (e.g. 1. / 1.1 / Chapter 1 / (1))\n\n" + "NOT TOC content:\n" + "- Body text paragraphs\n" + "- Data tables\n" + "- Image-heavy pages\n" + "- A single heading with no page-number listing\n\n" + "Return strict JSON (no markdown fences):\n" + '{"still_toc": true/false, "confidence": "high"/"medium"/"low", ' + '"reason": "brief reason"}' + ), + }, + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + }, + ] + + try: + client = get_openai_client(model=model) + raw, _ = client.chat_completion_with_usage( + messages=[{"role": "user", "content": content_parts}], + model=model, + temperature=0.1, + max_tokens=200, + response_format={"type": "json_object"}, + ) + data = json.loads(raw) + return bool(data.get("still_toc")) + except Exception as exc: + logger.warning( + "[extract.toc] boundary VLM check failed for page {}: {}", page_num, exc + ) + # Conservative: stop expansion on failure + return False + + +# -- Progressive boundary detection -------------------------------------------- + + +def _detect_toc_range_for_anchor( + *, + anchor_page: int, + pdf_path: str, + page_count: int, + output_dir: str, + dpi: int, + model: str, +) -> tuple[int, int, list[dict[str, Any]]]: + """Progressively expand from anchor_page to find the TOC end boundary. + + Returns: + (start_page, end_page, trace_rounds) -- all 1-based inclusive. + """ + start_page = anchor_page + current_end = min(anchor_page + BOUNDARY_STEP_PAGES - 1, page_count) + trace_rounds: list[dict[str, Any]] = [] + + for round_idx in range(MAX_BOUNDARY_ROUNDS): + check_page = current_end + png_path = os.path.join(output_dir, f"toc_boundary_p{check_page}.png") + run_in_child_process( + _render_single_page_worker, + pdf_path, + check_page, + png_path, + dpi, + timeout=60, + ) + + still_toc = _vlm_check_boundary_page(png_path, check_page, model) + trace_rounds.append( + { + "round": round_idx, + "check_page": check_page, + "window": [start_page, current_end], + "still_toc": still_toc, + } + ) + logger.info( + "[extract.toc] round {}: page {} still_toc={}", + round_idx, + check_page, + still_toc, + ) + + if not still_toc: + # Boundary page is NOT TOC; TOC ends at previous page + current_end = max(check_page - 1, start_page) + break + + next_end = min(current_end + BOUNDARY_STEP_PAGES, page_count) + if next_end == current_end: + break + current_end = next_end + + return start_page, current_end, trace_rounds + + +# -- MinerU local extraction --------------------------------------------------- + + +def _run_mineru_local( + pdf_path: str, + start_page_0based: int, + end_page_0based: int, + output_dir: str, +) -> list[str]: + """Run mineru CLI on a page range and return the resulting markdown lines.""" + os.makedirs(output_dir, exist_ok=True) + cmd = [ + "mineru", + "-p", + pdf_path, + "-o", + output_dir, + "-s", + str(start_page_0based), + "-e", + str(end_page_0based), + "-t", + "false", # skip table parsing for speed + "-b", + "pipeline", + "-m", + "txt", + ] + logger.info( + "[extract.toc] mineru local: pages {}-{}, cmd: {}", + start_page_0based, + end_page_0based, + " ".join(cmd), + ) + + try: + proc = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=MINERU_TIMEOUT_SECONDS, + ) + if proc.returncode != 0: + logger.error( + "[extract.toc] mineru failed (code={}): {}", + proc.returncode, + proc.stderr[-500:] if proc.stderr else "", + ) + return [] + except subprocess.TimeoutExpired: + logger.error("[extract.toc] mineru timed out after {}s", MINERU_TIMEOUT_SECONDS) + return [] + + md_lines: list[str] = [] + output_path = Path(output_dir) + for md_file in sorted(output_path.rglob("*.md")): + content = md_file.read_text(encoding="utf-8", errors="replace") + md_lines.extend(content.splitlines()) + + logger.info("[extract.toc] mineru produced {} markdown lines", len(md_lines)) + return md_lines + + +# -- Main tool ----------------------------------------------------------------- + + +@register_tool( + name="extract.toc_with_boundaries", + description=( + "VLM-confirms TOC anchor pages, progressively detects TOC boundaries, " + "runs mineru local extraction, then reuses toc_parser for hierarchy." + ), + allowed_states={DocumentAgentState.CLASSIFIED}, +) +def extract_toc_with_boundaries( + ctx: ToolContext, _args: dict[str, Any] +) -> ToolResult: + start = time.monotonic() + anchors = ctx.blackboard.toc_anchor_pages + warnings: list[str] = [] + debug_info: dict[str, Any] = {} + + if not anchors: + logger.info("[extract.toc] no anchor pages, skipping") + ctx.blackboard.toc_result = TocResult( + method="none", + notes="No TOC anchor pages found by find.toc_anchor_pages", + ) + return ToolResult( + status="ok", + payload={"toc_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + ) + + model = ctx.settings.get("vlm_model") or os.environ.get( + "IMAGE_MODEL", "qwen3.5-flash" + ) + dpi = int(ctx.settings.get("toc_png_dpi", "144")) + page_count = ctx.blackboard.page_count + output_dir = str( + Path(ctx.output_dir or os.path.expanduser("~/.knowhere/_debug_profile")) + / "toc_pages" + ) + os.makedirs(output_dir, exist_ok=True) + + # -- Phase 1: VLM confirm anchors ----------------------------------------- + confirmed = _vlm_confirm_anchors(anchors, model) + debug_info["phase1_confirmed"] = [a.page for a in confirmed] + debug_info["phase1_rejected"] = [ + a.page for a in anchors if a not in confirmed + ] + + if not confirmed: + ctx.blackboard.toc_result = TocResult( + method="none", + notes="VLM rejected all TOC anchor candidates", + ) + return ToolResult( + status="ok", + payload={"toc_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + warnings=["VLM rejected all anchor pages"], + debug=debug_info, + ) + + # -- Phase 2: progressive boundary detection ------------------------------- + toc_ranges: list[tuple[int, int]] = [] + all_trace_rounds: list[dict[str, Any]] = [] + + for anchor in confirmed: + toc_start, toc_end, trace_rounds = _detect_toc_range_for_anchor( + anchor_page=anchor.page, + pdf_path=ctx.pdf_path, + page_count=page_count, + output_dir=output_dir, + dpi=dpi, + model=model, + ) + toc_ranges.append((toc_start, toc_end)) + all_trace_rounds.extend(trace_rounds) + logger.info( + "[extract.toc] TOC region: pages {}-{}", toc_start, toc_end + ) + + debug_info["phase2_ranges"] = toc_ranges + debug_info["phase2_trace_rounds"] = all_trace_rounds + + # -- Phase 3: mineru local extraction -------------------------------------- + all_md_lines: list[str] = [] + for i, (toc_start, toc_end) in enumerate(toc_ranges): + region_dir = os.path.join(output_dir, f"mineru_region_{i}") + md_lines = _run_mineru_local( + pdf_path=ctx.pdf_path, + start_page_0based=toc_start - 1, # mineru uses 0-based + end_page_0based=toc_end - 1, + output_dir=region_dir, + ) + if md_lines: + all_md_lines.extend(md_lines) + else: + warnings.append( + f"mineru produced no output for region {i} (pages {toc_start}-{toc_end})" + ) + + debug_info["phase3_md_line_count"] = len(all_md_lines) + + if not all_md_lines: + ctx.blackboard.toc_result = TocResult( + toc_pages=[p for s, e in toc_ranges for p in range(s, e + 1)], + method="vlm_progressive", + notes="VLM detected TOC ranges but mineru produced no markdown", + ) + warnings.append("mineru produced no markdown for any TOC region") + return ToolResult( + status="ok", + payload={"toc_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + warnings=warnings, + debug=debug_info, + ) + + # -- Phase 4: toc_parser reuse --------------------------------------------- + try: + from app.services.document_parser.structure.toc_parser import ( + detect_tocs_in_texts, + ) + + hierarchy_model = ctx.settings.get("model") or os.environ.get( + "HIERARCHY_LLM_MODEL" + ) or os.environ.get("NORMOL_MODEL") + + toc_hierarchies, _filtered = detect_tocs_in_texts( + all_md_lines, + model_name=hierarchy_model, + hierarchy_model_name=hierarchy_model, + branch="normal", + limit_=150, + ) + except Exception as exc: + logger.error("[extract.toc] toc_parser failed: {}", exc) + toc_hierarchies = None + warnings.append(f"toc_parser failed: {exc}") + + # -- Write results to blackboard ------------------------------------------- + all_toc_pages = sorted( + {p for s, e in toc_ranges for p in range(s, e + 1)} + ) + + ctx.blackboard.toc_result = TocResult( + toc_pages=all_toc_pages, + method="vlm_progressive", + notes=( + f"VLM confirmed {len(confirmed)} TOC starts, " + f"expanded to {len(toc_ranges)} ranges: {toc_ranges}" + ), + ) + ctx.blackboard.toc_hierarchies = toc_hierarchies if toc_hierarchies else None + + # Persist toc_hierarchies to disk for inspection / downstream reuse + if toc_hierarchies and ctx.output_dir: + toc_json_path = os.path.join(ctx.output_dir, "toc_hierarchies.json") + try: + with open(toc_json_path, "w", encoding="utf-8") as f: + json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) + logger.info("[extract.toc] wrote toc_hierarchies to {}", toc_json_path) + except Exception as exc: + logger.warning("[extract.toc] failed to write toc_hierarchies: {}", exc) + + toc_summary: dict[str, Any] = { + "toc_ranges": toc_ranges, + "toc_page_count": len(all_toc_pages), + } + if toc_hierarchies: + for i, hier in enumerate(toc_hierarchies): + tree = hier.get("toc_tree", {}) + toc_summary[f"region_{i}_level1_count"] = len(tree) + toc_summary[f"region_{i}_level1_titles"] = list(tree.keys())[:10] + + return ToolResult( + status="ok", + payload={ + "toc_count": len(toc_hierarchies) if toc_hierarchies else 0, + "toc_page_count": len(all_toc_pages), + }, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary=toc_summary, + warnings=warnings, + debug=debug_info, + ) diff --git a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py new file mode 100644 index 000000000..6d9bc4437 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py @@ -0,0 +1,217 @@ +"""Scan for TOC anchor pages and render their PNGs for VLM inspection.""" + +from __future__ import annotations + +import gc +import os +import time +from collections import Counter +from pathlib import Path +from typing import Any + +from app.services.document_agent.manifest import TocAnchorPage, ToolContext, ToolResult +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) +from loguru import logger + +# CJK and English TOC keywords used for anchor detection. +TOC_KEYWORDS = {"目录", "目次", "contents", "tableofcontents", "table of contents"} + +# If a TOC keyword fingerprint appears on more than this fraction of total +# pages, it is treated as a recurring navigation element (header/footer link) +# rather than real TOC content. +RECURRING_ELEMENT_THRESHOLD = 0.30 + +# Hard cap on the number of candidate anchor pages sent to VLM. A real +# document never has more than ~30 TOC start pages. +MAX_ANCHOR_CANDIDATES = 30 + + +def _normalize_for_toc(text: str) -> str: + """Collapse whitespace for keyword matching.""" + return text.replace(" ", "").replace("\u3000", "").lower() + + +@worker +def _render_pages_worker( + queue, pdf_path: str, pages: list[int], output_dir: str, dpi: int +) -> None: + import pymupdf # type: ignore[import] + + results: list[dict[str, Any]] = [] + try: + doc = pymupdf.open(pdf_path) + for page_num in pages: + idx = page_num - 1 + if 0 <= idx < doc.page_count: + page = doc[idx] + mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) + pix = page.get_pixmap(matrix=mat) + png_name = f"toc_anchor_page_{page_num}.png" + png_path = os.path.join(output_dir, png_name) + pix.save(png_path) + results.append({"page": page_num, "png_path": png_path}) + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "results": results}) + + +def _filter_recurring_elements( + matches: list[tuple[int, str, int]], + total_pages: int, +) -> set[int]: + """Remove pages whose TOC keyword pattern is a recurring navigation element. + + Each *match* is ``(page, raw_line, line_index)``. We build a **composite + fingerprint** per page by joining all its matches as + ``"raw_line@line_idx\\n..."``. If the same composite fingerprint appears + on more than ``RECURRING_ELEMENT_THRESHOLD`` of all pages, those pages are + header/footer false-positives. + + Example (SpaceX S-1): + - page 17 → ``"Table of Contents@0\\nTABLE OF CONTENTS@1"`` (unique → keeps) + - page 18 → ``"Table of Contents@1"`` (376 pages → recurring → filtered) + - page 401 → ``"Table of Contents@0"`` (~31 pages → VLM decides) + """ + # Collect all matches per page + page_matches: dict[int, list[tuple[str, int]]] = {} + for page, raw_line, line_idx in matches: + page_matches.setdefault(page, []).append((raw_line, line_idx)) + + # Build composite fingerprint per page (sorted by line_idx for stability) + page_fingerprints: dict[int, str] = {} + for page, hits in page_matches.items(): + hits_sorted = sorted(hits, key=lambda h: h[1]) + page_fingerprints[page] = "\n".join( + f"{raw}@{idx}" for raw, idx in hits_sorted + ) + + # Group pages by composite fingerprint + fp_groups: dict[str, list[int]] = {} + for page, fp in page_fingerprints.items(): + fp_groups.setdefault(fp, []).append(page) + + threshold = max(int(total_pages * RECURRING_ELEMENT_THRESHOLD), 1) + + surviving: set[int] = set() + for fp, pages in fp_groups.items(): + if len(pages) > threshold: + logger.info( + "[find.toc_anchor_pages] recurring pattern filtered: " + "{!r} appears on {}/{} pages", + fp[:60], + len(pages), + total_pages, + ) + else: + surviving.update(pages) + + return surviving + + +@register_tool( + name="find.toc_anchor_pages", + description=( + "Scan page text previews for TOC keywords, filter recurring " + "navigation elements, then render candidate PNGs for VLM confirmation." + ), + allowed_states={DocumentAgentState.CLASSIFIED}, +) +def find_toc_anchor_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + total_pages = ctx.blackboard.page_count + + # Scan text previews for TOC keywords and record per-line matches. + # Each entry: (page, raw_line_text, line_index) + # We use raw (original) text for fingerprinting so that casing + # differences (e.g. "Table of Contents" vs "TABLE OF CONTENTS") + # naturally produce distinct composite fingerprints. + keyword_matches: list[tuple[int, str, int]] = [] + raw_hit_pages: set[int] = set() + + for feature in ctx.blackboard.page_features: + for line_idx, raw_line in enumerate(feature.text_lines_preview): + norm_line = _normalize_for_toc(raw_line) + for keyword in TOC_KEYWORDS: + if keyword in norm_line: + keyword_matches.append((feature.page, raw_line.strip(), line_idx)) + raw_hit_pages.add(feature.page) + break # one match per line is enough + + # Apply recurring element fingerprint filter + if keyword_matches: + anchor_pages = _filter_recurring_elements(keyword_matches, total_pages) + else: + anchor_pages = set() + + logger.info( + "[find.toc_anchor_pages] keyword scan: {} raw hits → {} after " + "fingerprint filter", + len(raw_hit_pages), + len(anchor_pages), + ) + + # Hard cap: a document never has more than ~30 real TOC start candidates. + if len(anchor_pages) > MAX_ANCHOR_CANDIDATES: + logger.warning( + "[find.toc_anchor_pages] {} candidates exceed cap of {}, truncating", + len(anchor_pages), + MAX_ANCHOR_CANDIDATES, + ) + anchor_pages = set(sorted(anchor_pages)[:MAX_ANCHOR_CANDIDATES]) + + if not anchor_pages: + logger.info("[find.toc_anchor_pages] no TOC keyword pages found") + ctx.blackboard.toc_anchor_pages = [] + return ToolResult( + status="ok", + payload={"anchor_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={"anchor_count": 0, "pages": []}, + ) + + # Render candidate pages as PNGs for downstream VLM confirmation + sorted_pages = sorted(anchor_pages) + output_dir = str( + Path(ctx.output_dir or os.path.expanduser("~/.knowhere/_debug_profile")) + / "toc_pages" + ) + os.makedirs(output_dir, exist_ok=True) + + dpi = int(ctx.settings.get("toc_png_dpi", "144")) + result = run_in_child_process( + _render_pages_worker, ctx.pdf_path, sorted_pages, output_dir, dpi, timeout=120 + ) + + anchors: list[TocAnchorPage] = [] + for item in result.get("results") or []: + page = int(item["page"]) + anchors.append( + TocAnchorPage(page=page, png_path=item["png_path"], source="text_scan") + ) + + ctx.blackboard.toc_anchor_pages = anchors + logger.info( + "[find.toc_anchor_pages] found {} anchor pages: {}", + len(anchors), + [a.page for a in anchors], + ) + + return ToolResult( + status="ok", + payload={"anchor_count": len(anchors)}, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={ + "anchor_count": len(anchors), + "pages": [a.to_dict() for a in anchors], + }, + ) + diff --git a/apps/worker/app/services/document_agent/tools/match_h1_pages.py b/apps/worker/app/services/document_agent/tools/match_h1_pages.py new file mode 100644 index 000000000..ac97c0a00 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/match_h1_pages.py @@ -0,0 +1,196 @@ +"""Match TOC level-1 headings to body pages via PyMuPDF text search.""" + +from __future__ import annotations + +import re +import time +import unicodedata +from typing import Any + +from app.services.document_agent.manifest import ( + H1BoundaryResult, + H1Candidate, + ToolContext, + ToolResult, +) +from app.services.document_agent.pdf_text import read_page_texts +from app.services.document_agent.registry import register_tool +from app.services.document_agent.state import DocumentAgentState +from loguru import logger + + +# ── Text normalization for matching ────────────────────────────────────── + +_LEADING_NUMBER_RE = re.compile( + r"""^ + (?: + [#]+\s* + | 第\s*[零一二三四五六七八九十百千\d]+\s*[章节篇部分] + | [零一二三四五六七八九十百千]+\s*[、。，,] + | [（(]\s*[零一二三四五六七八九十百千\d]+\s*[）)] + | \d+(?:\.\d+)*\.?\s* + | [IVXLCDM]+\.?\s+ + | [A-Za-z]\.\s+ + | Chapter\s+\w+\s* + ) + """, + re.VERBOSE | re.IGNORECASE, +) + +_PAGE_SUFFIX_RE = re.compile(r"[\s\.\-·…]+\d+\s*$") + + +def _normalize(text: str) -> str: + """Normalize text for fuzzy heading matching.""" + text = unicodedata.normalize("NFKC", text or "") + text = re.sub(r"\s+", " ", text).strip() + return text + + +def _clean_toc_title(title: str) -> str: + """Remove leading numbering/hashes and trailing page numbers from a TOC title.""" + cleaned = _PAGE_SUFFIX_RE.sub("", title or "").strip() + cleaned = _LEADING_NUMBER_RE.sub("", cleaned).strip() + return cleaned + + +def _extract_level1_titles(toc_hierarchies: list[dict[str, Any]]) -> list[str]: + """Extract level-1 titles from toc_hierarchies. + + Each hierarchy dict contains ``toc_tree`` – a nested dict where top-level + keys are level-1 headings (values are sub-heading dicts). + """ + titles: list[str] = [] + for hier in toc_hierarchies: + toc_tree = hier.get("toc_tree") or {} + for raw_title in toc_tree.keys(): + cleaned = _clean_toc_title(raw_title) + if cleaned and len(cleaned) >= 2: + titles.append(cleaned) + return titles + + +@register_tool( + name="match.h1_pages", + description=( + "Match TOC level-1 headings to body pages using PyMuPDF substring search. " + "Produces H1Candidate list for downstream shard planning." + ), + allowed_states={DocumentAgentState.CLASSIFIED}, +) +def match_h1_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + + if not ctx.blackboard.toc_hierarchies: + logger.info("[match.h1_pages] no toc_hierarchies, skipping") + ctx.blackboard.h1_result = H1BoundaryResult( + method="none", + notes="No toc_hierarchies available for H1 matching", + ) + return ToolResult( + status="ok", + payload={"h1_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + ) + + level1_titles = _extract_level1_titles(ctx.blackboard.toc_hierarchies) + if not level1_titles: + logger.info("[match.h1_pages] no level-1 titles in toc_hierarchies") + ctx.blackboard.h1_result = H1BoundaryResult( + method="toc_grep", + notes="toc_hierarchies contained no level-1 entries", + ) + return ToolResult( + status="ok", + payload={"h1_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={"level1_titles": level1_titles}, + ) + + # Build exclusion set: TOC pages should not be searched + toc_page_set: set[int] = set() + if ctx.blackboard.toc_result: + toc_page_set.update(ctx.blackboard.toc_result.toc_pages) + + # Read text for all non-TOC pages + search_pages = sorted( + p + for p in range(1, ctx.blackboard.page_count + 1) + if p not in toc_page_set + ) + page_texts = read_page_texts(ctx.pdf_path, search_pages, timeout=300) + + # Strict substring matching: for each level-1 title, find the first body page + h1_candidates: list[H1Candidate] = [] + matched_titles: list[str] = [] + unmatched_titles: list[str] = [] + + for title in level1_titles: + normalized_title = _normalize(title) + found = False + for page in search_pages: + text = page_texts.get(page, "") + normalized_text = _normalize(text) + if normalized_title in normalized_text: + # Find the matched line for evidence + matched_line = "" + for line in text.splitlines(): + if normalized_title in _normalize(line): + matched_line = line.strip()[:100] + break + + h1_candidates.append( + H1Candidate( + title=title, + page=page, + confidence=0.88, + matched_line=matched_line, + source="toc_exact_top", + evidence={ + "normalized_needle": normalized_title, + "page_text_length": len(text), + }, + ) + ) + matched_titles.append(title) + found = True + break # Only first match per title + + if not found: + unmatched_titles.append(title) + + # Deduplicate: if multiple titles map to the same page, keep the first + seen_pages: set[int] = set() + deduped: list[H1Candidate] = [] + for candidate in h1_candidates: + if candidate.page not in seen_pages: + seen_pages.add(candidate.page) + deduped.append(candidate) + h1_candidates = deduped + + ctx.blackboard.h1_result = H1BoundaryResult( + h1_candidates=h1_candidates, + method="toc_grep", + notes=( + f"Matched {len(matched_titles)}/{len(level1_titles)} level-1 titles. " + f"Unmatched: {unmatched_titles[:5]}" + ), + ) + + logger.info( + "[match.h1_pages] matched {}/{} level-1 titles to body pages: {}", + len(matched_titles), + len(level1_titles), + [(c.title[:20], c.page) for c in h1_candidates], + ) + + return ToolResult( + status="ok", + payload={"h1_count": len(h1_candidates)}, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={ + "level1_titles": level1_titles, + "matched": [(c.title, c.page) for c in h1_candidates], + "unmatched": unmatched_titles, + }, + ) diff --git a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py index 48ed6ea66..b48900a02 100644 --- a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py @@ -39,6 +39,7 @@ def build_anatomy_map(ctx: ToolContext) -> PageAnatomyMap: hierarchy_assist=ctx.blackboard.hierarchy_assist, shard_plan=ctx.blackboard.shard_plan, boundary_candidates=ctx.blackboard.boundary_candidates, + toc_hierarchies=ctx.blackboard.toc_hierarchies, page_processing_plan=None, global_signals=ctx.blackboard.global_signals, trace_summary={ diff --git a/apps/worker/app/services/document_agent/tools/test_vlm_toc_extract.py b/apps/worker/app/services/document_agent/tools/test_vlm_toc_extract.py new file mode 100644 index 000000000..100d1dfd3 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/test_vlm_toc_extract.py @@ -0,0 +1,817 @@ +#!/usr/bin/env python3 +"""Standalone test: VLM-direct TOC extraction from page PNGs. + +Usage: + cd knowhereapi-main + python apps/worker/app/services/document_agent/tools/test_vlm_toc_extract.py + +Requires: + - ALI_API_KEYS env var (or in apps/worker/.env) + - pymupdf, openai installed + - Test PDFs accessible on disk +""" + +from __future__ import annotations + +import base64 +import json +import os +import sys +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +# --------------------------------------------------------------------------- +# VLM prompt — no hardcoded examples, no overfitting +# --------------------------------------------------------------------------- + +VLM_TOC_EXTRACT_PROMPT = """\ +You are analyzing a Table of Contents (TOC) page from a document. + +Your task is to extract every TOC entry visible on this page. + +Each entry consists of three parts: +1. **title** — the section or chapter name, copied verbatim from the page. + - EXCLUDE any trailing dots (…·.), dashes (—-), or leader lines that connect the title to its page number. + - If one entry's title wraps across multiple printed lines, combine them into a single string. + - Include any numbering prefix that is part of the title text (e.g. "1.", "第三章", "(二)"). +2. **page_number** — the page reference at the right side of the entry. + - Use an integer when the reference is a plain number (e.g. 26). + - Use a string when the reference is non-numeric (e.g. "iv", "F-1", "A-3"). + - Use null when no page reference is visible for that entry. +3. **level** — the hierarchy depth of the entry, determined by visual formatting cues: + - level 1: top-level entries — no indentation, or the largest / boldest text. + - level 2: sub-entries — indented under a level-1 entry, or in a noticeably smaller font. + - level 3+: deeper indentation, if present. + - Category headers or group labels that are visually distinct (centered, larger font, different style) and do NOT have a page number should be treated as level 1. + +Additional rules: +- Extract ALL entries, even if the page only shows a partial continuation of the TOC. +- Do NOT include the TOC page's own heading (e.g. a "TABLE OF CONTENTS" or "目 录" title at the top) as an entry. +- Do NOT include column headers (e.g. a standalone "Page" label) as entries. +- Preserve the original language and wording of each title exactly. + +Return strict JSON with no markdown fences: +{"entries": [{"title": "...", "page_number": ..., "level": ...}, ...]} +""" + +# Mirror the production boundary step from extract_toc_with_boundaries.py +BOUNDARY_STEP_PAGES = 5 + +VLM_TOC_CONTINUATION_CONTEXT = """\ + +--- IMPORTANT: Continuation Context --- +This is a CONTINUATION page of a multi-page Table of Contents. +The previous page(s) already extracted the following entries: + +{previous_summary} + +The LAST active category/section before this page was: + Level {last_l1_level}: "{last_l1_title}" + +Entries on THIS page that visually continue as sub-items under that +category (same indentation, same numbering sequence) must keep their +correct subordinate level — do NOT promote them to level 1 just because +the parent heading is not visible on this page. +""" + + +# --------------------------------------------------------------------------- +# Test case definitions +# --------------------------------------------------------------------------- + + +@dataclass +class TocTestCase: + """One test document for VLM TOC extraction.""" + + name: str + pdf_path: str | None # None means PNGs are pre-rendered + toc_page_nums: list[int] # 1-based page numbers to extract + output_dir: str + pre_rendered_pngs: list[str] = field(default_factory=list) + expected_entry_count_range: tuple[int, int] = (1, 999) + description: str = "" + toc_output_path: str | None = None # Override toc_hierarchies.json output path + + +def _build_test_cases() -> list[TocTestCase]: + debug_root = os.path.expanduser("~/.knowhere/_debug_profile") + cases: list[TocTestCase] = [] + + # Case 1: SpaceX S-1 — single TOC page, English + spacex_pdf = "/Users/wuchengke/Downloads/spacex-s1.pdf" + spacex_out = os.path.join(debug_root, "spacex-s1", "toc_pages") + if os.path.exists(spacex_pdf): + cases.append( + TocTestCase( + name="SpaceX S-1 (English)", + pdf_path=spacex_pdf, + toc_page_nums=[17], # VLM-confirmed TOC start, boundary says 17 only + output_dir=spacex_out, + expected_entry_count_range=(20, 30), + description="Single flat TOC page, all level-1 entries, page nums on right", + ) + ) + else: + # Fall back to pre-rendered PNG + png17 = os.path.join(spacex_out, "toc_anchor_page_17.png") + if os.path.exists(png17): + cases.append( + TocTestCase( + name="SpaceX S-1 (English, pre-rendered)", + pdf_path=None, + toc_page_nums=[17], + output_dir=spacex_out, + pre_rendered_pngs=[png17], + expected_entry_count_range=(20, 30), + description="Single flat TOC page from pre-rendered PNG", + ) + ) + + # Case 2: Chinese corporate regulations — multi-page TOC with categories + cn_pdf = "/Users/wuchengke/Desktop/temp/test_docs/SJSYJ-SC-2024 企业制度汇编（上册）.pdf" + cn_out = os.path.join(debug_root, "chinese_corp", "toc_pages") + if os.path.exists(cn_pdf): + cases.append( + TocTestCase( + name="企业制度汇编 (Chinese, multi-page TOC)", + pdf_path=cn_pdf, + toc_page_nums=[5, 6], # TOC spans pages 5-6 + output_dir=cn_out, + expected_entry_count_range=(25, 45), + description=( + "Multi-page TOC with category headers (经营类/生产类/安全类/...), " + "subcategory codes (SJSYJ-SC101-2024), numbered entries, " + "multi-line wrapping, and Chinese dot leaders" + ), + ) + ) + + return cases + + +# --------------------------------------------------------------------------- +# PNG rendering +# --------------------------------------------------------------------------- + + +def _render_page_png(pdf_path: str, page_num: int, output_dir: str, dpi: int = 144) -> str: + """Render a single page to PNG. Returns the PNG path.""" + import pymupdf # type: ignore[import] + + os.makedirs(output_dir, exist_ok=True) + png_path = os.path.join(output_dir, f"toc_page_{page_num}.png") + + doc = pymupdf.open(pdf_path) + try: + idx = page_num - 1 + if 0 <= idx < doc.page_count: + page = doc[idx] + mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) + pix = page.get_pixmap(matrix=mat) + pix.save(png_path) + else: + raise ValueError(f"Page {page_num} out of range (total: {doc.page_count})") + finally: + doc.close() + + return png_path + + +# --------------------------------------------------------------------------- +# VLM call +# --------------------------------------------------------------------------- + + +def _build_openai_client(model: str): + """Build an OpenAI SDK client for the given model, without requiring AppConfig. + + Mirrors the routing logic in shared.services.ai.openai_compatible_client_sync + but reads env vars directly so the test can run standalone. + """ + from openai import OpenAI + + model_lower = model.lower() + + if "qwen" in model_lower: + # Aliyun DashScope + api_key = os.environ.get("ALI_API_KEYS", "").strip() + # ALI_API_KEYS can be JSON array or comma-separated; grab the first one + if api_key.startswith("["): + import re + keys = re.findall(r'"([^"]+)"', api_key) + api_key = keys[0] if keys else "" + elif "," in api_key: + api_key = api_key.split(",")[0].strip() + # Handle token_id=api_key format + if "=" in api_key and not api_key.startswith("sk-"): + api_key = api_key.split("=", 1)[1] + base_url = os.environ.get( + "ALI_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1" + ) + elif "deepseek" in model_lower: + api_key = os.environ.get("DS_KEY", "") + base_url = os.environ.get("DS_URL", "https://api.deepseek.com/v1") + else: + api_key = os.environ.get("OPENAI_API_KEY", "") + base_url = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + + # Strip /chat/completions suffix if present + if base_url.rstrip("/").endswith("/chat/completions"): + base_url = base_url.rstrip("/").removesuffix("/chat/completions") + + if not api_key: + raise RuntimeError( + f"No API key found for model {model!r}. " + "Set ALI_API_KEYS, DS_KEY, or OPENAI_API_KEY env var." + ) + + return OpenAI(api_key=api_key, base_url=base_url, timeout=120, max_retries=2) + + +def _build_continuation_context(previous_entries: list[dict[str, Any]]) -> str: + """Build a concise context string from previously extracted entries. + + Returns the continuation context to append to the VLM prompt, + or empty string if no previous entries exist. + """ + if not previous_entries: + return "" + + # Build a compact summary: last N entries with their levels + # Show the last 8 entries to give enough context + tail = previous_entries[-8:] + summary_lines = [] + for e in tail: + lvl = e.get("level", "?") + title = e.get("title", "?") + pn = e.get("page_number") + pn_str = f" → p.{pn}" if pn is not None else "" + summary_lines.append(f" L{lvl}: {title}{pn_str}") + + if len(previous_entries) > 8: + summary_lines.insert(0, f" ... ({len(previous_entries) - 8} earlier entries omitted)") + + previous_summary = "\n".join(summary_lines) + + # Find the last L1 entry (the active parent category) + last_l1 = None + for e in reversed(previous_entries): + if e.get("level") == 1: + last_l1 = e + break + + if last_l1 is None: + # No L1 found — still provide the summary but skip the "last active" part + return f"\n\n--- IMPORTANT: Continuation Context ---\nThis is a CONTINUATION page. Previous entries:\n{previous_summary}\n" + + return VLM_TOC_CONTINUATION_CONTEXT.format( + previous_summary=previous_summary, + last_l1_level=last_l1.get("level", 1), + last_l1_title=last_l1.get("title", "?"), + ) + + +def _vlm_extract_toc_entries( + png_path: str, + page_num: int, + model: str, + previous_entries: list[dict[str, Any]] | None = None, +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Call VLM to extract TOC entries from a single page PNG. + + Args: + png_path: path to the page PNG image. + page_num: 1-based page number. + model: VLM model name. + previous_entries: entries already extracted from earlier TOC pages. + Used to build continuation context so the model can correctly + assign hierarchy levels on continuation pages. + + Returns: + (entries, meta) where meta has token usage and timing info. + """ + client = _build_openai_client(model) + + with open(png_path, "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + + # Build prompt: base + optional continuation context + prompt_text = VLM_TOC_EXTRACT_PROMPT + continuation_ctx = _build_continuation_context(previous_entries or []) + if continuation_ctx: + prompt_text += continuation_ctx + + content_parts: list[dict[str, Any]] = [ + {"type": "text", "text": prompt_text}, + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + }, + ] + + messages = [{"role": "user", "content": content_parts}] + + t0 = time.monotonic() + response = client.chat.completions.create( + model=model, + messages=messages, + temperature=0.1, + max_tokens=4096, + response_format={"type": "json_object"}, + extra_body={"enable_thinking": False}, + ) + elapsed_ms = int((time.monotonic() - t0) * 1000) + + raw = response.choices[0].message.content or "" + usage_obj = response.usage + usage = { + "prompt_tokens": getattr(usage_obj, "prompt_tokens", 0) or 0, + "completion_tokens": getattr(usage_obj, "completion_tokens", 0) or 0, + "total_tokens": getattr(usage_obj, "total_tokens", 0) or 0, + } + + meta = { + "page": page_num, + "model": model, + "elapsed_ms": elapsed_ms, + "usage": usage, + "raw_response_length": len(raw), + "has_continuation_context": bool(continuation_ctx), + } + + # Parse JSON + data = json.loads(raw) + if isinstance(data, dict): + entries = data.get("entries", []) + elif isinstance(data, list): + entries = data + else: + entries = [] + + return entries, meta + + +# --------------------------------------------------------------------------- +# VLM entries → toc_hierarchies.json conversion +# --------------------------------------------------------------------------- + + +def _build_toc_tree(entries: list[dict[str, Any]]) -> dict[str, Any]: + """Build a nested dict tree from VLM entries, same algorithm as + toc_hierarchy.build_tree_tocs() but self-contained for standalone use. + + Args: + entries: list of {"title": str, "level": int, "page_number": ...} + + Returns: + Nested dict: {heading: {child_heading: {...}, ...}, ...} + """ + if not entries: + return {} + + root: dict[str, Any] = {} + stack: list[tuple[dict[str, Any], int]] = [(root, 0)] + + positive_levels = [e["level"] for e in entries if isinstance(e.get("level"), int) and e["level"] > 0] + level_for_minus_one = max(positive_levels) + 1 if positive_levels else 1 + + for entry in entries: + heading = entry.get("title", "").strip() + if not heading: + continue + + original_level = entry.get("level", 1) + normalized_level = level_for_minus_one if original_level == -1 else original_level + + while len(stack) > 1 and stack[-1][1] >= normalized_level: + stack.pop() + + parent_dict = stack[-1][0] + parent_dict[heading] = {} + stack.append((parent_dict[heading], normalized_level)) + + return root + + +def _build_toc_with_level_md(entries: list[dict[str, Any]]) -> str: + """Build a markdown table string compatible with existing toc_with_level format.""" + if not entries: + return "" + + lines = ["| id | heading | level |"] + lines.append("|----|---------|-------|") + for i, e in enumerate(entries, 1): + heading = e.get("title", "").strip().replace("|", "\\|") + level = e.get("level", 1) + # Pad heading for readability + lines.append(f"| {i:<2} | {heading:<60} | {level:<5} |") + return "\n".join(lines) + + +def vlm_entries_to_toc_hierarchies( + all_entries: list[dict[str, Any]], + toc_page_nums: list[int], + scan_end_page: int | None = None, + page_count: int | None = None, +) -> list[dict[str, Any]]: + """Convert VLM extraction results to the standard toc_hierarchies.json schema. + + Args: + all_entries: merged VLM entries from all TOC pages. + Each entry: {"title": str, "page_number": int|str|None, "level": int} + toc_page_nums: 1-based page numbers of the TOC pages. + scan_end_page: 1-based page number of the lookahead boundary check. + In production, this is `anchor_page + BOUNDARY_STEP_PAGES - 1`. + If the boundary check page was NOT TOC, it still counts as part + of the scan range (it was the page we "looked at" to decide). + Defaults to `max(toc_page_nums) + BOUNDARY_STEP_PAGES - 1`. + page_count: total pages in the document (used to clamp scan_end_page). + + Returns: + List with one dict per TOC region (usually 1), matching the schema: + { + "toc_range": [start_page, end_page], + "toc_range_unit": "page", + "scan_range": [start_page, scan_end_page], + "source": "vlm", + "toc_with_level": [structured list], + "toc_with_level_md": "| id | heading | level |\n...", + "toc_tree": {nested dict} + } + """ + if not all_entries or not toc_page_nums: + return [] + + # Build structured toc_with_level list + toc_with_level: list[dict[str, Any]] = [] + for i, entry in enumerate(all_entries, 1): + item: dict[str, Any] = { + "id": i, + "heading": entry.get("title", "").strip(), + "level": entry.get("level", 1), + } + # Include page_number from VLM extraction + pn = entry.get("page_number") + item["page_number"] = pn + toc_with_level.append(item) + + # Build toc_tree from clean entries + toc_tree = _build_toc_tree(all_entries) + + # Build markdown table for backward compatibility + toc_with_level_md = _build_toc_with_level_md(all_entries) + + # Determine page ranges + start_page = min(toc_page_nums) + end_page = max(toc_page_nums) + + # scan_range: the full lookahead window used during boundary detection. + # In production: anchor + BOUNDARY_STEP_PAGES - 1, clamped to page_count. + # The anchor is start_page (not end_page), matching _detect_toc_range_for_anchor(). + if scan_end_page is None: + scan_end_page = start_page + BOUNDARY_STEP_PAGES - 1 + if page_count is not None: + scan_end_page = min(scan_end_page, page_count) + + return [ + { + "toc_range": [start_page, end_page], + "toc_range_unit": "page", + "scan_range": [start_page, scan_end_page], + "source": "vlm", + "toc_with_level": toc_with_level, + "toc_with_level_md": toc_with_level_md, + "toc_tree": toc_tree, + } + ] + + +# --------------------------------------------------------------------------- +# Comparison with existing toc_hierarchies.json +# --------------------------------------------------------------------------- + + +def _compare_with_existing( + vlm_entries: list[dict[str, Any]], + test_case: TocTestCase, +) -> dict[str, Any]: + """Compare VLM results with existing toc_hierarchies.json if available.""" + comparison: dict[str, Any] = {"available": False} + + # Try to find existing toc_hierarchies.json + parent_dir = str(Path(test_case.output_dir).parent) + existing_path = os.path.join(parent_dir, "toc_hierarchies.json") + if not os.path.exists(existing_path): + return comparison + + with open(existing_path, "r", encoding="utf-8") as f: + existing = json.load(f) + + comparison["available"] = True + comparison["existing_path"] = existing_path + + # Extract titles from existing toc_tree + existing_titles: list[str] = [] + if isinstance(existing, list) and existing: + tree = existing[0].get("toc_tree", {}) + for key, sub in tree.items(): + existing_titles.append(key) + if isinstance(sub, dict): + for subkey in sub: + existing_titles.append(subkey) + + vlm_titles = [e.get("title", "") for e in vlm_entries] + + # Quality checks + issues: list[str] = [] + + # Check 1: titles with residual page numbers (MinerU artifact) + import re + residual_pattern = re.compile(r"\s+\d+\s*$") + existing_with_residual = [t for t in existing_titles if residual_pattern.search(t)] + vlm_with_residual = [t for t in vlm_titles if residual_pattern.search(t)] + + comparison["existing_residual_page_nums"] = existing_with_residual + comparison["vlm_residual_page_nums"] = vlm_with_residual + + if existing_with_residual and not vlm_with_residual: + issues.append( + f"✅ VLM fixed {len(existing_with_residual)} residual page numbers " + f"in titles (e.g. '{existing_with_residual[0]}')" + ) + elif vlm_with_residual: + issues.append( + f"⚠️ VLM still has {len(vlm_with_residual)} titles with residual " + f"page numbers: {vlm_with_residual[:3]}" + ) + + # Check 2: broken multi-line titles + # In existing data: "MANAGEMENT'S DISCUSSION..." and "OF OPERATIONS 74" are separate + broken_line_keywords = ["OF OPERATIONS", "CLASS A COMMON STOCK"] + existing_broken = [ + t for t in existing_titles + if any(t.strip().startswith(kw) for kw in broken_line_keywords) + ] + vlm_broken = [ + t for t in vlm_titles + if any(t.strip().startswith(kw) for kw in broken_line_keywords) + ] + + if existing_broken and not vlm_broken: + issues.append( + f"✅ VLM merged {len(existing_broken)} broken multi-line titles " + f"(e.g. '{existing_broken[0]}')" + ) + elif vlm_broken: + issues.append( + f"⚠️ VLM still has {len(vlm_broken)} broken multi-line titles: " + f"{vlm_broken}" + ) + + # Check 3: false positive entries (e.g. "Page" as a title) + false_positive_patterns = {"Page", "页码"} + existing_fp = [t for t in existing_titles if t.strip() in false_positive_patterns] + vlm_fp = [ + e for e in vlm_entries + if e.get("title", "").strip() in false_positive_patterns + ] + + if existing_fp and not vlm_fp: + issues.append( + f"✅ VLM removed {len(existing_fp)} false positive entries " + f"(e.g. '{existing_fp[0]}')" + ) + elif vlm_fp: + issues.append( + f"⚠️ VLM still has false positive entries: " + f"{[e['title'] for e in vlm_fp]}" + ) + + # Check 4: entry count comparison + comparison["existing_entry_count"] = len(existing_titles) + comparison["vlm_entry_count"] = len(vlm_entries) + comparison["quality_checks"] = issues + + return comparison + + +# --------------------------------------------------------------------------- +# Main test runner +# --------------------------------------------------------------------------- + + +def run_test(test_case: TocTestCase, model: str) -> dict[str, Any]: + """Run VLM TOC extraction for one test case.""" + print(f"\n{'='*70}") + print(f"TEST: {test_case.name}") + print(f" {test_case.description}") + print(f" TOC pages: {test_case.toc_page_nums}") + print(f"{'='*70}") + + os.makedirs(test_case.output_dir, exist_ok=True) + + # Detect page count for scan_range calculation + pdf_page_count: int | None = None + if test_case.pdf_path: + try: + import pymupdf + with pymupdf.open(test_case.pdf_path) as doc: + pdf_page_count = len(doc) + print(f" PDF page count: {pdf_page_count}") + except Exception: + pass + + # Step 1: Prepare PNGs + png_paths: list[tuple[int, str]] = [] + + if test_case.pre_rendered_pngs: + for i, png in enumerate(test_case.pre_rendered_pngs): + png_paths.append((test_case.toc_page_nums[i], png)) + print(f" [png] Using pre-rendered: {png}") + elif test_case.pdf_path: + for page_num in test_case.toc_page_nums: + png = _render_page_png(test_case.pdf_path, page_num, test_case.output_dir) + png_paths.append((page_num, png)) + print(f" [png] Rendered page {page_num}: {png}") + + # Step 2: VLM extraction per page + all_entries: list[dict[str, Any]] = [] + all_meta: list[dict[str, Any]] = [] + total_elapsed_ms = 0 + + for page_num, png_path in png_paths: + is_continuation = len(all_entries) > 0 + ctx_label = " (with context)" if is_continuation else "" + print(f"\n [vlm] Extracting page {page_num}{ctx_label}...") + try: + entries, meta = _vlm_extract_toc_entries( + png_path, page_num, model, + previous_entries=all_entries if is_continuation else None, + ) + all_entries.extend(entries) + all_meta.append(meta) + total_elapsed_ms += meta["elapsed_ms"] + print(f" [vlm] Page {page_num}: {len(entries)} entries, {meta['elapsed_ms']}ms") + + # Show first few entries + for e in entries[:5]: + title = e.get("title", "?") + pn = e.get("page_number", "?") + lv = e.get("level", "?") + print(f" L{lv}: {title!r} → p.{pn}") + if len(entries) > 5: + print(f" ... ({len(entries)} total)") + + except Exception as exc: + print(f" [vlm] ❌ FAILED for page {page_num}: {exc}") + all_meta.append({"page": page_num, "error": str(exc)}) + + # Step 3: Summary + print(f"\n --- Summary ---") + print(f" Total entries: {len(all_entries)}") + print(f" Total VLM time: {total_elapsed_ms}ms") + expected_lo, expected_hi = test_case.expected_entry_count_range + count_ok = expected_lo <= len(all_entries) <= expected_hi + print( + f" Entry count check: {len(all_entries)} " + f"(expected {expected_lo}-{expected_hi}) → " + f"{'✅' if count_ok else '⚠️ OUT OF RANGE'}" + ) + + # Step 4: Quality analysis + level_dist = {} + for e in all_entries: + lv = e.get("level", "?") + level_dist[lv] = level_dist.get(lv, 0) + 1 + print(f" Level distribution: {level_dist}") + + entries_with_page = [e for e in all_entries if e.get("page_number") is not None] + entries_no_page = [e for e in all_entries if e.get("page_number") is None] + print(f" Entries with page number: {len(entries_with_page)}") + print(f" Entries without page number: {len(entries_no_page)}") + if entries_no_page: + for e in entries_no_page[:5]: + print(f" → {e.get('title', '?')!r} (level={e.get('level')})") + + # Step 5: Compare with existing + comparison = _compare_with_existing(all_entries, test_case) + if comparison.get("available"): + print(f"\n --- Comparison with existing toc_hierarchies.json ---") + print(f" Existing entries: {comparison['existing_entry_count']}") + print(f" VLM entries: {comparison['vlm_entry_count']}") + for check in comparison.get("quality_checks", []): + print(f" {check}") + + # Step 6: Save raw VLM results + result = { + "test_name": test_case.name, + "model": model, + "toc_pages": test_case.toc_page_nums, + "total_entries": len(all_entries), + "total_elapsed_ms": total_elapsed_ms, + "entries": all_entries, + "per_page_meta": all_meta, + "level_distribution": level_dist, + "comparison": comparison if comparison.get("available") else None, + } + + parent_dir = str(Path(test_case.output_dir).parent) + output_path = os.path.join(parent_dir, "vlm_toc_entries.json") + with open(output_path, "w", encoding="utf-8") as f: + json.dump(result, f, ensure_ascii=False, indent=2) + print(f"\n 📄 Raw results saved to: {output_path}") + + # Step 7: Generate toc_hierarchies.json in production schema + toc_hierarchies = vlm_entries_to_toc_hierarchies( + all_entries, test_case.toc_page_nums, + page_count=pdf_page_count, + ) + + toc_hier_path = test_case.toc_output_path or os.path.join( + parent_dir, "toc_hierarchies.json" + ) + os.makedirs(os.path.dirname(toc_hier_path), exist_ok=True) + with open(toc_hier_path, "w", encoding="utf-8") as f: + json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) + print(f" 📄 toc_hierarchies.json saved to: {toc_hier_path}") + + # Print schema summary + if toc_hierarchies: + h = toc_hierarchies[0] + print(f" --- Schema Summary ---") + print(f" toc_range: {h['toc_range']}") + print(f" toc_range_unit: {h['toc_range_unit']}") + print(f" source: {h['source']}") + print(f" toc_with_level entries: {len(h['toc_with_level'])}") + print(f" toc_tree L1 keys: {list(h['toc_tree'].keys())[:8]}") + + # Step 8: Print full entry table + print(f"\n --- Full Entry Table ---") + print(f" {'#':>3} {'Lv':>3} {'Page':>6} Title") + print(f" {'─'*3} {'─'*3} {'─'*6} {'─'*50}") + for i, e in enumerate(all_entries, 1): + title = e.get("title", "?") + pn = e.get("page_number", "—") + lv = e.get("level", "?") + # Truncate long titles for display + disp_title = title if len(title) <= 60 else title[:57] + "..." + print(f" {i:>3} {lv:>3} {str(pn):>6} {disp_title}") + + return result + + +def main() -> None: + # Load env vars from worker/.env if available + script_dir = os.path.dirname(os.path.abspath(__file__)) + # Walk up to find the repo root: test_vlm_toc_extract.py is at + # apps/worker/app/services/document_agent/tools/ + repo_root = os.path.abspath(os.path.join(script_dir, "..", "..", "..", "..", "..", "..")) + env_file = os.path.join(repo_root, "apps", "worker", ".env") + if os.path.exists(env_file): + print(f"Loading env from: {env_file}") + with open(env_file) as f: + for line in f: + line = line.strip() + if line and not line.startswith("#") and "=" in line: + key, _, value = line.partition("=") + key = key.strip() + value = value.strip().strip("'\"") + if key and key not in os.environ: # don't override existing + os.environ[key] = value + + model = os.environ.get("IMAGE_MODEL", "qwen3.5-flash") + print(f"VLM model: {model}") + + test_cases = _build_test_cases() + if not test_cases: + print("No test cases found! Check PDF paths.") + sys.exit(1) + + print(f"Found {len(test_cases)} test case(s)") + + results: list[dict[str, Any]] = [] + for tc in test_cases: + try: + result = run_test(tc, model) + results.append(result) + except Exception as exc: + print(f"\n❌ Test '{tc.name}' FAILED: {exc}") + import traceback + traceback.print_exc() + + # Final summary + print(f"\n{'='*70}") + print("FINAL SUMMARY") + print(f"{'='*70}") + for r in results: + name = r["test_name"] + count = r["total_entries"] + ms = r["total_elapsed_ms"] + print(f" {name}: {count} entries, {ms}ms") + + +if __name__ == "__main__": + main() From e5d8d7ea31960d7a2df59a46d310d614804192ac Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 26 May 2026 04:18:10 +0800 Subject: [PATCH 06/11] refactor: restructure document agent into modular bootstrap, executor, planner, and persist components --- .gitignore | 1 + Makefile | 5 +- .../f8a9b0c1d2e3_add_parse_agent_tables.py | 2 - .../app/services/document_agent/__init__.py | 2 - .../document_agent/bootstrap/__init__.py | 7 + .../bootstrap/aggregate_stats.py | 119 +++ .../document_agent/bootstrap/classify.py | 5 + .../document_agent/bootstrap/probe.py | 5 + .../app/services/document_agent/budget.py | 58 +- .../services/document_agent/coordinator.py | 190 ++-- .../document_agent/executor/__init__.py | 13 + .../document_agent/executor/prompts.py | 15 + .../document_agent/executor/react_loop.py | 296 +++++++ .../app/services/document_agent/manifest.py | 162 ++-- .../app/services/document_agent/pdf_text.py | 2 +- .../document_agent/persist/__init__.py | 5 + .../document_agent/persist/persist.py | 8 + .../document_agent/planner/__init__.py | 13 + .../document_agent/planner/planner.py | 294 +++++++ .../document_agent/planner/prompts.py | 16 + .../prompts/coordinator_system.md | 12 - .../prompts/propose_hierarchy_assist.md | 11 - .../app/services/document_agent/registry.py | 92 +- .../app/services/document_agent/state.py | 24 +- .../services/document_agent/tools/__init__.py | 8 +- .../tools/classify_page_kinds.py | 34 +- .../tools/collect_boundary_candidates.py | 129 --- .../tools/extract_toc_with_boundaries.py | 215 ++--- .../tools/find_toc_anchor_pages.py | 6 +- .../document_agent/tools/grep_text.py | 79 ++ .../document_agent/tools/inspect_pages.py | 108 +++ .../document_agent/tools/match_h1_pages.py | 5 +- .../tools/persist_anatomy_map.py | 13 +- .../tools/probe_page_features.py | 7 - .../tools/propose_hierarchy_assist.py | 97 --- .../tools/propose_shard_plan.py | 176 ++-- .../tools/test_vlm_toc_extract.py | 817 ------------------ .../tools/validate_anatomy_map.py | 16 +- .../services/document_agent/tools/verdict.py | 39 + .../document_agent/tools/vlm_toc_extractor.py | 237 +++++ .../app/services/document_agent/trace.py | 29 +- .../app/services/document_agent/validators.py | 42 +- .../app/services/document_agent/visual.py | 86 ++ apps/worker/build_manifest_sjsyj.py | 159 ---- apps/worker/run_hierarchy_sjsyj.py | 117 --- .../models/database/document_page_plan.py | 5 - 46 files changed, 1838 insertions(+), 1943 deletions(-) create mode 100644 apps/worker/app/services/document_agent/bootstrap/__init__.py create mode 100644 apps/worker/app/services/document_agent/bootstrap/aggregate_stats.py create mode 100644 apps/worker/app/services/document_agent/bootstrap/classify.py create mode 100644 apps/worker/app/services/document_agent/bootstrap/probe.py create mode 100644 apps/worker/app/services/document_agent/executor/__init__.py create mode 100644 apps/worker/app/services/document_agent/executor/prompts.py create mode 100644 apps/worker/app/services/document_agent/executor/react_loop.py create mode 100644 apps/worker/app/services/document_agent/persist/__init__.py create mode 100644 apps/worker/app/services/document_agent/persist/persist.py create mode 100644 apps/worker/app/services/document_agent/planner/__init__.py create mode 100644 apps/worker/app/services/document_agent/planner/planner.py create mode 100644 apps/worker/app/services/document_agent/planner/prompts.py delete mode 100644 apps/worker/app/services/document_agent/prompts/coordinator_system.md delete mode 100644 apps/worker/app/services/document_agent/prompts/propose_hierarchy_assist.md delete mode 100644 apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py create mode 100644 apps/worker/app/services/document_agent/tools/grep_text.py create mode 100644 apps/worker/app/services/document_agent/tools/inspect_pages.py delete mode 100644 apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py delete mode 100644 apps/worker/app/services/document_agent/tools/test_vlm_toc_extract.py create mode 100644 apps/worker/app/services/document_agent/tools/verdict.py create mode 100644 apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py create mode 100644 apps/worker/app/services/document_agent/visual.py delete mode 100644 apps/worker/build_manifest_sjsyj.py delete mode 100644 apps/worker/run_hierarchy_sjsyj.py diff --git a/.gitignore b/.gitignore index 0cc60df6b..8a56b4f57 100644 --- a/.gitignore +++ b/.gitignore @@ -69,6 +69,7 @@ test_*.csv !requirements.csv # Local debugging scripts +apps/worker/scripts/ apps/worker/start_celery_worker.py apps/worker/start_celery_debug.sh apps/worker/clear_celery_queues.sh diff --git a/Makefile b/Makefile index b6488b017..073def0e8 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: lint lint-fix typecheck check +.PHONY: lint lint-fix typecheck check test-doc-agent UV := uv REPO_UV_CACHE_DIR := $(CURDIR)/.uv-cache @@ -29,3 +29,6 @@ typecheck: $(PYRIGHT) --project pyproject.toml $(PYRIGHT_PATHS) check: lint typecheck + +test-doc-agent: + cd apps/worker && $(UV_RUN_ENV) $(UV) run pytest tests/document_agent diff --git a/apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py b/apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py index 8ba86f35e..5ae58232b 100644 --- a/apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py +++ b/apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py @@ -61,9 +61,7 @@ def upgrade() -> None: sa.Column("page_plan_id", sa.String(length=36), nullable=False), sa.Column("job_id", sa.String(length=36), nullable=False), sa.Column("page_count", sa.Integer(), nullable=False, server_default="0"), - sa.Column("hierarchy_assist", sa.JSON(), nullable=True), sa.Column("shard_plan", sa.JSON(), nullable=True), - sa.Column("page_processing_plan", sa.JSON(), nullable=True), sa.Column("global_signals", sa.JSON(), nullable=True), sa.Column("created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), sa.ForeignKeyConstraint(["job_id"], ["jobs.job_id"], ondelete="CASCADE"), diff --git a/apps/worker/app/services/document_agent/__init__.py b/apps/worker/app/services/document_agent/__init__.py index 02d729084..ee28ee9f6 100644 --- a/apps/worker/app/services/document_agent/__init__.py +++ b/apps/worker/app/services/document_agent/__init__.py @@ -1,7 +1,6 @@ """Page anatomy agent for hierarchy-first PDF profiling.""" from app.services.document_agent.manifest import ( - HierarchyAssistPlan, PageAnatomyMap, PageFeature, PageLabel, @@ -10,7 +9,6 @@ from app.services.document_agent.profile_agent import ProfileAgent __all__ = [ - "HierarchyAssistPlan", "PageAnatomyMap", "PageFeature", "PageLabel", diff --git a/apps/worker/app/services/document_agent/bootstrap/__init__.py b/apps/worker/app/services/document_agent/bootstrap/__init__.py new file mode 100644 index 000000000..b8494b99f --- /dev/null +++ b/apps/worker/app/services/document_agent/bootstrap/__init__.py @@ -0,0 +1,7 @@ +"""Deterministic bootstrap steps for the document profile agent.""" + +from app.services.document_agent.bootstrap.aggregate_stats import aggregate_doc_stats +from app.services.document_agent.bootstrap.classify import classify_page_kinds +from app.services.document_agent.bootstrap.probe import probe_page_features + +__all__ = ["aggregate_doc_stats", "classify_page_kinds", "probe_page_features"] diff --git a/apps/worker/app/services/document_agent/bootstrap/aggregate_stats.py b/apps/worker/app/services/document_agent/bootstrap/aggregate_stats.py new file mode 100644 index 000000000..2177cbf6c --- /dev/null +++ b/apps/worker/app/services/document_agent/bootstrap/aggregate_stats.py @@ -0,0 +1,119 @@ +"""Aggregate page-feature statistics for VLM profile planning.""" + +from __future__ import annotations + +import statistics +import time +from typing import Any + +from app.services.document_agent.manifest import PageFeature, ToolContext, ToolResult + +PROFILE_METRICS = ( + "raw_text_length", + "text_density", + "image_coverage", + "table_count", + "drawings_count", +) + +EXTREMA_ROLES = { + "raw_text_length": ("min", "max"), + "text_density": ("min", "max"), + "image_coverage": ("max",), + "table_count": ("max",), + "drawings_count": ("max",), +} + +EXTREMA_LABELS = { + "raw_text_length": "text_length", + "text_density": "text_density", + "image_coverage": "image_heavy", + "table_count": "table_heavy", + "drawings_count": "drawing_heavy", +} + + +def _percentile(values: list[float], percentile: float) -> float: + if not values: + return 0.0 + if len(values) == 1: + return values[0] + ordered = sorted(values) + index = (len(ordered) - 1) * percentile + lower = int(index) + upper = min(lower + 1, len(ordered) - 1) + weight = index - lower + return ordered[lower] * (1 - weight) + ordered[upper] * weight + + +def _metric_value(feature: PageFeature, metric: str) -> float: + return float(getattr(feature, metric)) + + +def aggregate_doc_stats(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + features = list(ctx.blackboard.page_features) + stats: dict[str, Any] = {} + extrema_pages: list[int] = [] + extrema_samples: list[dict[str, Any]] = [] + for metric in PROFILE_METRICS: + pairs = [(feature.page, _metric_value(feature, metric)) for feature in features] + values = [value for _, value in pairs] + if not pairs: + stats[metric] = { + "mean": 0.0, + "p50": 0.0, + "p90": 0.0, + "min": {"page": None, "value": 0.0}, + "max": {"page": None, "value": 0.0}, + } + continue + min_page, min_value = min(pairs, key=lambda item: (item[1], item[0])) + max_page, max_value = max(pairs, key=lambda item: (item[1], -item[0])) + stats[metric] = { + "mean": round(statistics.fmean(values), 4), + "p50": round(_percentile(values, 0.5), 4), + "p90": round(_percentile(values, 0.9), 4), + "min": {"page": min_page, "value": round(min_value, 4)}, + "max": {"page": max_page, "value": round(max_value, 4)}, + } + extrema_by_role = { + "min": (min_page, min_value), + "max": (max_page, max_value), + } + for role in EXTREMA_ROLES[metric]: + page, value = extrema_by_role[role] + extrema_pages.append(page) + extrema_samples.append( + { + "page": page, + "metric": metric, + "label": EXTREMA_LABELS[metric], + "role": role, + "value": round(value, 4), + } + ) + + deduped_extrema = sorted(set(extrema_pages)) + ctx.blackboard.doc_stats = stats + ctx.blackboard.extrema_pages = deduped_extrema + ctx.blackboard.global_signals["doc_stats"] = stats + ctx.blackboard.global_signals["extrema_pages"] = deduped_extrema + ctx.blackboard.global_signals["extrema_samples"] = extrema_samples + return ToolResult( + status="ok", + payload={ + "metric_count": len(PROFILE_METRICS), + "extrema_pages": deduped_extrema, + "extrema_samples": extrema_samples, + }, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={ + "doc_stats": stats, + "extrema_pages": deduped_extrema, + "extrema_samples": extrema_samples, + }, + ) + + +__all__ = ["PROFILE_METRICS", "aggregate_doc_stats"] diff --git a/apps/worker/app/services/document_agent/bootstrap/classify.py b/apps/worker/app/services/document_agent/bootstrap/classify.py new file mode 100644 index 000000000..c8dc808c7 --- /dev/null +++ b/apps/worker/app/services/document_agent/bootstrap/classify.py @@ -0,0 +1,5 @@ +"""Bootstrap wrapper for deterministic page classification.""" + +from app.services.document_agent.tools.classify_page_kinds import classify_page_kinds + +__all__ = ["classify_page_kinds"] diff --git a/apps/worker/app/services/document_agent/bootstrap/probe.py b/apps/worker/app/services/document_agent/bootstrap/probe.py new file mode 100644 index 000000000..14877b1ba --- /dev/null +++ b/apps/worker/app/services/document_agent/bootstrap/probe.py @@ -0,0 +1,5 @@ +"""Bootstrap wrapper for deterministic page probing.""" + +from app.services.document_agent.tools.probe_page_features import probe_page_features + +__all__ = ["probe_page_features"] diff --git a/apps/worker/app/services/document_agent/budget.py b/apps/worker/app/services/document_agent/budget.py index d5ee9e80d..6edbc1309 100644 --- a/apps/worker/app/services/document_agent/budget.py +++ b/apps/worker/app/services/document_agent/budget.py @@ -17,49 +17,55 @@ def remaining(self) -> int: class BudgetTracker: - """A minimal counter with the same public shape as retrieval's ledger.""" + """A minimal synchronous ledger with plan and visual pools.""" - def __init__(self, *, plan_budget: int = 5000, max_tool_calls: int = 12) -> None: + def __init__( + self, + *, + plan_budget: int = 5000, + visual_budget: int = 8000, + ) -> None: self._plan = BudgetPool(capacity=max(int(plan_budget), 0)) - self._max_tool_calls = max(int(max_tool_calls), 1) - self._tool_calls = 0 - - def increment_tool_call(self) -> bool: - if self._tool_calls >= self._max_tool_calls: - return False - self._tool_calls += 1 - return True + self._visual = BudgetPool(capacity=max(int(visual_budget), 0)) def try_reserve(self, pool: str, est: int) -> bool: - if pool != "plan": + if pool not in {"plan", "visual"}: return True est = max(int(est), 0) - if self._plan.remaining < est: + budget_pool = self._pool(pool) + if budget_pool.remaining < est: return False - self._plan.reserved += est + budget_pool.reserved += est return True def commit(self, pool: str, *, actual: int, est: int) -> None: - if pool != "plan": + if pool not in {"plan", "visual"}: return est = max(int(est), 0) actual = max(int(actual), 0) - self._plan.reserved = max(self._plan.reserved - est, 0) - self._plan.used = min(self._plan.capacity, self._plan.used + actual) + budget_pool = self._pool(pool) + budget_pool.reserved = max(budget_pool.reserved - est, 0) + budget_pool.used = min(budget_pool.capacity, budget_pool.used + actual) def refund(self, pool: str, *, est: int) -> None: - if pool != "plan": + if pool not in {"plan", "visual"}: return - self._plan.reserved = max(self._plan.reserved - max(int(est), 0), 0) + budget_pool = self._pool(pool) + budget_pool.reserved = max(budget_pool.reserved - max(int(est), 0), 0) + + def _pool(self, pool: str) -> BudgetPool: + return self._visual if pool == "visual" else self._plan + + def _pool_snapshot(self, pool: BudgetPool) -> dict[str, int]: + return { + "capacity": pool.capacity, + "used": pool.used, + "reserved": pool.reserved, + "remaining": pool.remaining, + } def snapshot(self) -> dict[str, object]: return { - "plan": { - "capacity": self._plan.capacity, - "used": self._plan.used, - "reserved": self._plan.reserved, - "remaining": self._plan.remaining, - }, - "tool_calls": self._tool_calls, - "max_tool_calls": self._max_tool_calls, + "plan": self._pool_snapshot(self._plan), + "visual": self._pool_snapshot(self._visual), } diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 5330b0fa6..eedd989b6 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -1,73 +1,28 @@ -"""Synchronous ReAct-style coordinator for the document profile agent.""" +"""Plan-then-act coordinator for the document profile agent.""" from __future__ import annotations -import json import os from typing import Any from loguru import logger +from app.services.document_agent.bootstrap import ( + aggregate_doc_stats, + classify_page_kinds, + probe_page_features, +) from app.services.document_agent.budget import BudgetTracker +from app.services.document_agent.executor import ReActExecutor from app.services.document_agent.manifest import PageAnatomyMap, ToolContext +from app.services.document_agent.persist import build_anatomy_map, persist_anatomy_map +from app.services.document_agent.planner import ProfilePlanner from app.services.document_agent.registry import REGISTRY from app.services.document_agent.state import AgentBlackboard, DocumentAgentState -from app.services.document_agent.tools.persist_anatomy_map import build_anatomy_map +from app.services.document_agent import tools as _registered_tools # noqa: F401 from app.services.document_agent.trace import ParseRunRecorder -TRANSITIONS: dict[str, DocumentAgentState] = { - "probe.page_features": DocumentAgentState.PROBED, - "classify.page_kinds": DocumentAgentState.PROBED, - "find.toc_anchor_pages": DocumentAgentState.CLASSIFIED, - "extract.toc_with_boundaries": DocumentAgentState.CLASSIFIED, - "match.h1_pages": DocumentAgentState.H1_FOUND, - "collect.boundary_candidates": DocumentAgentState.H1_FOUND, - "propose.hierarchy_assist": DocumentAgentState.H1_FOUND, - "propose.shard_plan": DocumentAgentState.H1_FOUND, - "validate.anatomy_map": DocumentAgentState.VALIDATED, - "persist.anatomy_map": DocumentAgentState.PERSISTED, -} - -REQUIRED_TOOLS = [ - "probe.page_features", - "classify.page_kinds", - "find.toc_anchor_pages", - "extract.toc_with_boundaries", - "match.h1_pages", - "collect.boundary_candidates", - "propose.hierarchy_assist", - "propose.shard_plan", - "validate.anatomy_map", - "persist.anatomy_map", -] - - -def _parse_tool_calls(response: Any) -> list[dict[str, Any]]: - choices = getattr(response, "choices", None) - if not choices: - return [] - message = choices[0].message - calls = getattr(message, "tool_calls", None) or [] - parsed: list[dict[str, Any]] = [] - for call in calls: - function = getattr(call, "function", None) - if function is None: - continue - try: - args = json.loads(getattr(function, "arguments", "{}") or "{}") - except json.JSONDecodeError: - args = {} - parsed.append( - { - "id": getattr(call, "id", ""), - "name": getattr(function, "name", ""), - "args": args, - } - ) - return parsed - - class ProfileCoordinator: def __init__( self, @@ -82,8 +37,8 @@ def __init__( self.state = DocumentAgentState.INIT self.blackboard = AgentBlackboard() self.budget = BudgetTracker( - plan_budget=int(os.environ.get("PARSE_AGENT_PLAN_BUDGET", "5000")), - max_tool_calls=int(os.environ.get("PARSE_AGENT_MAX_TOOL_CALLS", "15")), + plan_budget=int(os.environ.get("PARSE_AGENT_PLAN_BUDGET", "50000")), + visual_budget=int(os.environ.get("PARSE_AGENT_VISUAL_BUDGET", "80000")), ) effective_settings = settings or {} if model: @@ -102,58 +57,99 @@ def __init__( self.round_index = 0 def run(self) -> PageAnatomyMap: - # The deterministic tool chain is the primary execution contract. LLM is - # used inside proposal tools where it adds judgement, not to own ordering. - final_status = "ready" try: - for tool_name in REQUIRED_TOOLS: - if not self.budget.increment_tool_call(): - final_status = "fallback" - break - result = REGISTRY.dispatch(tool_name, self.ctx, {}, self.state) - self.trace.record_step( - round_index=self.round_index, - actor=f"tool:{tool_name}", - action_type="tool_call", - result=result, - tool_name=tool_name, - tool_args={}, + self.state = DocumentAgentState.RUNNING + self._run_bootstrap() + self._run_toc_pipeline() + profile, initial_decision, planner_result = ProfilePlanner(self.ctx).propose() + self.blackboard.document_profile = profile + self.blackboard.global_signals["document_profile"] = profile.to_dict() + self.trace.record_step( + round_index=self.round_index, + actor="planner", + action_type="plan", + result=planner_result, + tool_name=None, + tool_args={}, + ) + self.round_index += 1 + + executor_result = ReActExecutor( + self.ctx, + registry=REGISTRY, + max_rounds=int(self.ctx.settings.get("max_rounds", 30)), + initial_decision=initial_decision, + ).run() + if executor_result.verdict.status != "success": + raise RuntimeError( + f"profile aborted: {executor_result.verdict.rationale}" ) - if result.status not in {"ok", "invalid"}: - raise RuntimeError(result.error or f"{tool_name} failed") - self._advance(tool_name) - self._maybe_advance_composite_state() - self.round_index += 1 - if self.state == DocumentAgentState.PERSISTED: - self.state = DocumentAgentState.READY anatomy = build_anatomy_map(self.ctx) + persist_result = persist_anatomy_map(self.ctx, {}) + self.trace.record_step( + round_index=self.round_index, + actor="persist", + action_type="persist", + result=persist_result, + tool_name="persist.anatomy_map", + tool_args={}, + ) + self.state = DocumentAgentState.READY + self.trace.write_trace_artifact( + self.ctx.output_dir, + final_status="ready", + summary=anatomy.trace_summary | self.trace.summary(), + ) self.trace.flush( - final_status=final_status, + final_status="ready", summary=anatomy.trace_summary | self.trace.summary(), ) return anatomy except Exception as exc: logger.error(f"[document_agent] profile failed: {exc}") self.state = DocumentAgentState.FAILED + self.trace.write_trace_artifact( + self.ctx.output_dir, + final_status="failed", + summary={"error": str(exc), "budget": self.ctx.budget.snapshot()}, + ) self.trace.flush(final_status="failed", summary={"error": str(exc)}) raise - def _advance(self, tool_name: str) -> None: - next_state = TRANSITIONS.get(tool_name, self.state) - self.state = next_state - self.blackboard.mark(self.state) - - def _maybe_advance_composite_state(self) -> None: - if ( - self.state == DocumentAgentState.PROBED - and self.blackboard.page_labels + def _run_bootstrap(self) -> None: + for tool_name, handler in ( + ("probe.page_features", probe_page_features), + ("classify.page_kinds", classify_page_kinds), + ("aggregate.doc_stats", aggregate_doc_stats), ): - self.state = DocumentAgentState.CLASSIFIED - self.blackboard.mark(self.state) - if ( - self.state == DocumentAgentState.H1_FOUND - and self.blackboard.hierarchy_assist is not None - and self.blackboard.shard_plan is not None + result = handler(self.ctx, {}) + self.trace.record_step( + round_index=self.round_index, + actor=f"bootstrap:{tool_name}", + action_type="bootstrap", + result=result, + tool_name=tool_name, + tool_args={}, + ) + if result.status != "ok": + raise RuntimeError(result.error or f"{tool_name} failed") + self.round_index += 1 + + def _run_toc_pipeline(self) -> None: + for tool_name in ( + "find.toc_anchor_pages", + "extract.toc_with_boundaries", + "match.h1_pages", ): - self.state = DocumentAgentState.PLANNED - self.blackboard.mark(self.state) + result = REGISTRY.dispatch(tool_name, self.ctx, {}) + self.trace.record_step( + round_index=self.round_index, + actor=f"toc:{tool_name}", + action_type="toc", + result=result, + tool_name=tool_name, + tool_args={}, + ) + if result.status not in {"ok", "invalid"}: + raise RuntimeError(result.error or f"{tool_name} failed") + self.round_index += 1 diff --git a/apps/worker/app/services/document_agent/executor/__init__.py b/apps/worker/app/services/document_agent/executor/__init__.py new file mode 100644 index 000000000..e5656b258 --- /dev/null +++ b/apps/worker/app/services/document_agent/executor/__init__.py @@ -0,0 +1,13 @@ +"""ReAct executor for the document profile agent.""" + +from app.services.document_agent.executor.react_loop import ( + ExecutorResult, + ReActExecutor, + _parse_decision, +) + +__all__ = [ + "ExecutorResult", + "ReActExecutor", + "_parse_decision", +] diff --git a/apps/worker/app/services/document_agent/executor/prompts.py b/apps/worker/app/services/document_agent/executor/prompts.py new file mode 100644 index 000000000..8c225d2cc --- /dev/null +++ b/apps/worker/app/services/document_agent/executor/prompts.py @@ -0,0 +1,15 @@ +"""Prompts for executor reflexion.""" + +REFLEXION_INSTRUCTIONS = ( + "You are the executor of a document profiling agent. Decide the next action " + "from the blackboard facts and available tools. Return strict JSON with keys: " + "action (tool_call or verdict_now), rationale, optional tool_name/tool_args, " + "optional verdict {status, rationale}. 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 verdict only after validation succeeds. If a " + "tool failed or validation is invalid, either gather targeted evidence and " + "retry the relevant tool or abort with a clear rationale." +) + +__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 new file mode 100644 index 000000000..2d95a62ae --- /dev/null +++ b/apps/worker/app/services/document_agent/executor/react_loop.py @@ -0,0 +1,296 @@ +"""ReAct-style executor for the document profile agent.""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass +from typing import Any + +from app.services.document_agent.manifest import ( + AgentVerdict, + ReflexionDecision, + ToolContext, + ToolResult, +) +from app.services.document_agent.executor.prompts import REFLEXION_INSTRUCTIONS +from app.services.document_agent.registry import ToolRegistry +from shared.utils.token_estimate import estimate_tokens + + +@dataclass +class ExecutorResult: + verdict: AgentVerdict + rounds: int + + +def _compact_blackboard(ctx: ToolContext) -> dict[str, Any]: + return { + "page_count": ctx.blackboard.page_count, + "page_kind_counts": ctx.blackboard.global_signals.get("page_kind_counts", {}), + "doc_stats": ctx.blackboard.doc_stats, + "extrema_pages": ctx.blackboard.extrema_pages, + "document_profile": ctx.blackboard.document_profile.to_dict() + if ctx.blackboard.document_profile + else None, + "toc_anchor_pages": [anchor.page for anchor in ctx.blackboard.toc_anchor_pages], + "toc_pages": ( + ctx.blackboard.toc_result.toc_pages if ctx.blackboard.toc_result else [] + ), + "toc_hierarchies_count": len(ctx.blackboard.toc_hierarchies or []), + "h1_count": ( + len(ctx.blackboard.h1_result.h1_candidates) + if ctx.blackboard.h1_result + else 0 + ), + "shard_plan": ctx.blackboard.shard_plan.to_dict() + if ctx.blackboard.shard_plan + else None, + "validation_report": ctx.blackboard.validation_report, + "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(), + } + + +def _parse_decision(raw: str) -> ReflexionDecision: + data = json.loads(raw) + action = str(data.get("action") or "tool_call") + if action not in {"tool_call", "verdict_now"}: + action = "tool_call" + verdict = None + if isinstance(data.get("verdict"), dict): + verdict_data = data["verdict"] + status = str(verdict_data.get("status") or "abort") + if status not in {"success", "abort"}: + status = "abort" + verdict = AgentVerdict( + status=status, # type: ignore[arg-type] + rationale=str(verdict_data.get("rationale") or data.get("rationale") or ""), + ) + return ReflexionDecision( + action=action, # type: ignore[arg-type] + rationale=str(data.get("rationale") or ""), + tool_name=data.get("tool_name"), + tool_args=dict(data.get("tool_args") or {}), + verdict=verdict, + ) + + +class ReActExecutor: + def __init__( + self, + ctx: ToolContext, + *, + registry: ToolRegistry, + max_rounds: int = 30, + initial_decision: ReflexionDecision | None = None, + ) -> None: + self.ctx = ctx + self.registry = registry + self.max_rounds = max_rounds + self._initial_decision = initial_decision + + def run(self) -> ExecutorResult: + for round_index in range(self.max_rounds): + pending_recovery_verdict: AgentVerdict | None = None + decision, result = self._next_decision(round_index) + self.ctx.blackboard.global_signals.setdefault("reflexion_decisions", []).append( + decision.to_dict() + ) + if self.ctx.trace: + self.ctx.trace.record_step( + round_index=round_index, + actor=f"executor:r{round_index}", + action_type="reflexion", + result=result, + tool_name=decision.tool_name, + tool_args=decision.tool_args, + ) + + tool_name: str | None = None + tool_args: dict[str, Any] = {} + if decision.action == "verdict_now": + verdict = decision.verdict or AgentVerdict( + status="abort", + rationale=decision.rationale or "Executor stopped without verdict.", + ) + if verdict.status == "success" and not ( + self.ctx.blackboard.validation_report + and self.ctx.blackboard.validation_report.get("valid") is True + ): + decision = ReflexionDecision( + action="tool_call", + rationale=( + "Validate the anatomy map before accepting a success verdict." + ), + tool_name="validate.anatomy_map", + tool_args={}, + ) + tool_name, tool_args = self._resolve_tool_call(decision) + else: + self.ctx.blackboard.verdict = verdict + return ExecutorResult(verdict=verdict, rounds=round_index + 1) + else: + tool_name, tool_args = self._resolve_tool_call(decision) + + if not tool_name: + verdict = AgentVerdict( + status="abort", + rationale="Executor did not choose a tool.", + ) + self.ctx.blackboard.verdict = verdict + return ExecutorResult(verdict=verdict, rounds=round_index + 1) + + tool_result = self.registry.dispatch(tool_name, self.ctx, tool_args) + if self.ctx.trace: + self.ctx.trace.record_step( + round_index=round_index, + actor=f"tool:{tool_name}", + action_type="tool_call", + result=tool_result, + tool_name=tool_name, + tool_args=tool_args, + ) + self.ctx.blackboard.step_history.append( + { + "round": round_index, + "tool_name": tool_name, + "tool_args": tool_args, + "status": tool_result.status, + "error": tool_result.error, + } + ) + if tool_result.status == "error": + pending_recovery_verdict = AgentVerdict( + status="abort", + rationale=tool_result.error or f"{tool_name} failed", + ) + elif tool_result.status == "precondition_unmet": + pending_recovery_verdict = AgentVerdict( + status="abort", + rationale=tool_result.error or f"{tool_name} precondition unmet", + ) + + if self.ctx.blackboard.verdict is not None: + return ExecutorResult( + verdict=self.ctx.blackboard.verdict, + rounds=round_index + 1, + ) + if pending_recovery_verdict is not None and self._is_deterministic_mode(): + self.ctx.blackboard.verdict = pending_recovery_verdict + return ExecutorResult( + verdict=pending_recovery_verdict, + rounds=round_index + 1, + ) + + verdict = AgentVerdict(status="abort", rationale="Maximum executor rounds reached.") + self.ctx.blackboard.verdict = verdict + return ExecutorResult(verdict=verdict, rounds=self.max_rounds) + + def _resolve_tool_call( + self, + decision: ReflexionDecision, + ) -> tuple[str | None, dict[str, Any]]: + if decision.action == "tool_call" and decision.tool_name: + return decision.tool_name, decision.tool_args + return None, {} + + def _is_deterministic_mode(self) -> bool: + return not (self.ctx.settings.get("executor_model") or self.ctx.settings.get("model")) + + def _next_decision(self, round_index: int) -> tuple[ReflexionDecision, ToolResult]: + if round_index == 0 and self._initial_decision is not None: + decision = self._initial_decision + return decision, ToolResult(status="ok", payload=decision.to_dict()) + model = self.ctx.settings.get("executor_model") or self.ctx.settings.get("model") + if not model: + decision = self._deterministic_decision() + return decision, ToolResult(status="ok", payload=decision.to_dict()) + + payload = { + "blackboard": _compact_blackboard(self.ctx), + "history_tail": self.ctx.blackboard.step_history[-6:], + "available_tools": self.registry.openai_specs(self.ctx.blackboard), + "round_index": round_index, + } + prompt = REFLEXION_INSTRUCTIONS + "\nPayload:\n" + json.dumps( + payload, + ensure_ascii=False, + ) + est = estimate_tokens(prompt) + if not self.ctx.budget.try_reserve("plan", est): + decision = ReflexionDecision( + action="verdict_now", + rationale="Planner budget exhausted.", + verdict=AgentVerdict(status="abort", rationale="Planner budget exhausted."), + ) + return decision, ToolResult( + status="ok", + payload=decision.to_dict(), + input_summary=payload, + ) + start = time.monotonic() + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=[{"role": "user", "content": prompt}], + model=model, + temperature=0.0, + max_tokens=1200, + response_format={"type": "json_object"}, + ) + self.ctx.budget.commit( + "plan", + actual=usage.get("total_tokens", est), + est=est, + ) + decision = _parse_decision(raw) + return decision, ToolResult( + status="ok", + payload=decision.to_dict(), + latency_ms=int((time.monotonic() - start) * 1000), + tokens_used=usage.get("total_tokens", 0), + input_summary=payload, + debug={"prompt_text": prompt, "raw_response": raw}, + ) + except Exception: + self.ctx.budget.refund("plan", est=est) + raise + + def _deterministic_decision(self) -> ReflexionDecision: + if self.ctx.blackboard.shard_plan is None: + return ReflexionDecision( + action="tool_call", + rationale="Create a shard plan.", + tool_name="propose.shard_plan", + tool_args={}, + ) + if not self.ctx.blackboard.validation_report: + return ReflexionDecision( + action="tool_call", + rationale="Validate the current shard plan.", + tool_name="validate.anatomy_map", + tool_args={}, + ) + if self.ctx.blackboard.validation_report.get("valid") is True: + return ReflexionDecision( + action="tool_call", + rationale="Validation succeeded; finish profile run.", + tool_name="verdict", + tool_args={ + "status": "success", + "rationale": "Validation succeeded; finishing profile run.", + }, + ) + return ReflexionDecision( + action="verdict_now", + rationale="Validation failed in deterministic mode.", + verdict=AgentVerdict(status="abort", rationale="Validation failed."), + ) + diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index 9bb5ea12a..ebc3c4e3e 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -7,29 +7,10 @@ from typing import Any, Literal -PageKind = Literal[ - "cover", - "toc", - "preface", - "normal", - "chapter_start", - "section_start", - "table_heavy", - "single_image", - "blank", - "separator", - "appendix", - "landscape", - "sparse", -] - -BoundaryCandidateKind = Literal[ - "h1", - "toc", - "blank", - "sparse", - "separator", -] +PageKind = Literal["normal", "table_heavy", "image_heavy", "low_content", "landscape"] + +ReflexionAction = Literal["tool_call", "verdict_now"] +VerdictStatus = Literal["success", "abort"] @dataclass @@ -62,6 +43,45 @@ def to_dict(self) -> dict[str, Any]: return asdict(self) +@dataclass +class DocumentProfile: + is_scanned: bool + category: str + category_rationale: str = "" + language: str = "unknown" + rationale: str = "" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class AgentVerdict: + status: VerdictStatus + rationale: str + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class ReflexionDecision: + action: ReflexionAction + rationale: str + tool_name: str | None = None + tool_args: dict[str, Any] = field(default_factory=dict) + verdict: AgentVerdict | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "action": self.action, + "rationale": self.rationale, + "tool_name": self.tool_name, + "tool_args": dict(self.tool_args), + "verdict": self.verdict.to_dict() if self.verdict else None, + } + + @dataclass class TocCandidate: title: str @@ -80,7 +100,7 @@ class TocAnchorPage: page: int # 1-based page number png_path: str # local PNG path for VLM inspection - source: Literal["page_label", "text_scan"] # how this anchor was discovered + source: Literal["page_label", "text_scan", "visual_scan"] # how this anchor was discovered def to_dict(self) -> dict[str, Any]: return asdict(self) @@ -90,7 +110,7 @@ def to_dict(self) -> dict[str, Any]: class TocResult: toc_pages: list[int] = field(default_factory=list) candidates: list[TocCandidate] = field(default_factory=list) - method: Literal["toc_marker", "vlm_progressive", "none"] = "none" + method: Literal["toc_marker", "vlm_progressive", "visual_scan", "none"] = "none" notes: str = "" def to_dict(self) -> dict[str, Any]: @@ -124,55 +144,6 @@ def to_dict(self) -> dict[str, Any]: return data -@dataclass -class BoundaryHint: - page: int - anchor_type: Literal["h1_boundary", "blank_separator", "separator", "forced_max_size"] - confidence: float - evidence: dict[str, Any] = field(default_factory=dict) - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -@dataclass -class BoundaryCandidate: - page: int - kind: BoundaryCandidateKind - priority: int - confidence: float - evidence: dict[str, Any] = field(default_factory=dict) - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -@dataclass -class HierarchyAssistPlan: - exclude_pages_from_title_candidates: list[int] = field(default_factory=list) - prefer_h1_start_pages: list[H1Candidate] = field(default_factory=list) - suppress_title_pages: list[int] = field(default_factory=list) - section_boundary_hints: list[BoundaryHint] = field(default_factory=list) - smart_parse_recommendation: Literal["off", "normal", "aggressive"] = "normal" - rationale: str = "" - - def to_dict(self) -> dict[str, Any]: - return { - "exclude_pages_from_title_candidates": list( - self.exclude_pages_from_title_candidates - ), - "prefer_h1_start_pages": [ - candidate.to_dict() for candidate in self.prefer_h1_start_pages - ], - "suppress_title_pages": list(self.suppress_title_pages), - "section_boundary_hints": [ - hint.to_dict() for hint in self.section_boundary_hints - ], - "smart_parse_recommendation": self.smart_parse_recommendation, - "rationale": self.rationale, - } - - @dataclass class ValidationReport: valid: bool @@ -221,32 +192,6 @@ def to_dict(self) -> dict[str, Any]: } -@dataclass -class PagePlanEntry: - page_index: int - strategy: Literal["vlm_detail", "vlm_lite", "text_only", "skip_tagging"] - expected_kind: str - rationale: str - estimated_cost_tokens: int = 0 - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -@dataclass -class PageProcessingPlan: - entries: list[PagePlanEntry] = field(default_factory=list) - global_strategy_summary: dict[str, int] = field(default_factory=dict) - rationale: str = "" - - def to_dict(self) -> dict[str, Any]: - return { - "entries": [entry.to_dict() for entry in self.entries], - "global_strategy_summary": dict(self.global_strategy_summary), - "rationale": self.rationale, - } - - @dataclass class PageAnatomyMap: job_id: str @@ -256,11 +201,9 @@ class PageAnatomyMap: page_labels: list[PageLabel] toc_result: TocResult h1_result: H1BoundaryResult - hierarchy_assist: HierarchyAssistPlan shard_plan: ShardPlan - boundary_candidates: list[BoundaryCandidate] = field(default_factory=list) + document_profile: DocumentProfile | None = None toc_hierarchies: list[dict[str, Any]] | None = None - page_processing_plan: PageProcessingPlan | None = None global_signals: dict[str, Any] = field(default_factory=dict) trace_summary: dict[str, Any] = field(default_factory=dict) created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) @@ -276,17 +219,11 @@ def to_dict(self) -> dict[str, Any]: "page_labels": [label.to_dict() for label in self.page_labels], "toc_result": self.toc_result.to_dict(), "h1_result": self.h1_result.to_dict(), - "hierarchy_assist": self.hierarchy_assist.to_dict(), "shard_plan": self.shard_plan.to_dict(), - "boundary_candidates": [ - candidate.to_dict() for candidate in self.boundary_candidates - ], + "document_profile": self.document_profile.to_dict() + if self.document_profile + else None, "toc_hierarchies": self.toc_hierarchies, - "page_processing_plan": ( - self.page_processing_plan.to_dict() - if self.page_processing_plan is not None - else None - ), "global_signals": dict(self.global_signals), "trace_summary": dict(self.trace_summary), "created_at": self.created_at.isoformat(), @@ -305,7 +242,6 @@ class ToolResult: warnings: list[str] = field(default_factory=list) debug: dict[str, Any] | None = None - @dataclass class ToolContext: pdf_path: str diff --git a/apps/worker/app/services/document_agent/pdf_text.py b/apps/worker/app/services/document_agent/pdf_text.py index 3fc2dd544..ca9684370 100644 --- a/apps/worker/app/services/document_agent/pdf_text.py +++ b/apps/worker/app/services/document_agent/pdf_text.py @@ -25,7 +25,7 @@ def _read_page_texts_worker(queue, pdf_path: str, pages: list[int]) -> None: for page in pages: idx = page - 1 if 0 <= idx < doc.page_count: - texts[page] = doc[idx].get_text() or "" + texts[page] = str(doc[idx].get_text() or "") finally: try: doc.close() diff --git a/apps/worker/app/services/document_agent/persist/__init__.py b/apps/worker/app/services/document_agent/persist/__init__.py new file mode 100644 index 000000000..56ef9f3e1 --- /dev/null +++ b/apps/worker/app/services/document_agent/persist/__init__.py @@ -0,0 +1,5 @@ +"""Deterministic persistence for document-agent outputs.""" + +from app.services.document_agent.persist.persist import build_anatomy_map, persist_anatomy_map + +__all__ = ["build_anatomy_map", "persist_anatomy_map"] diff --git a/apps/worker/app/services/document_agent/persist/persist.py b/apps/worker/app/services/document_agent/persist/persist.py new file mode 100644 index 000000000..b83ef8244 --- /dev/null +++ b/apps/worker/app/services/document_agent/persist/persist.py @@ -0,0 +1,8 @@ +"""Compatibility wrapper for deterministic anatomy-map persistence.""" + +from app.services.document_agent.tools.persist_anatomy_map import ( + build_anatomy_map, + persist_anatomy_map, +) + +__all__ = ["build_anatomy_map", "persist_anatomy_map"] diff --git a/apps/worker/app/services/document_agent/planner/__init__.py b/apps/worker/app/services/document_agent/planner/__init__.py new file mode 100644 index 000000000..09de37f57 --- /dev/null +++ b/apps/worker/app/services/document_agent/planner/__init__.py @@ -0,0 +1,13 @@ +"""One-shot VLM profile planner.""" + +from app.services.document_agent.planner.planner import ( + PAGE_KIND_DEFINITIONS, + ProfilePlanner, + _sample_pages, +) + +__all__ = [ + "PAGE_KIND_DEFINITIONS", + "ProfilePlanner", + "_sample_pages", +] diff --git a/apps/worker/app/services/document_agent/planner/planner.py b/apps/worker/app/services/document_agent/planner/planner.py new file mode 100644 index 000000000..5b361673b --- /dev/null +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -0,0 +1,294 @@ +"""Initial VLM profile planner for the document profile agent.""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from typing import Any, cast + +from loguru import logger + +from app.services.document_agent.manifest import ( + DocumentProfile, + ReflexionDecision, + ToolContext, + ToolResult, +) +from app.services.document_agent.planner.prompts import PLANNER_INSTRUCTIONS +from app.services.document_agent.visual import render_pages +from shared.utils.token_estimate import estimate_tokens + +PAGE_KIND_DEFINITIONS = { + "normal": ( + "Page with enough extractable native text and no dominant table/image " + "structure." + ), + "table_heavy": ( + "Page with detected tables or many vector drawings, often financial " + "tables or dense tabular layout." + ), + "image_heavy": ( + "Page dominated by image coverage with little extractable native text; " + "may be scanned, infographic, photo, or rendered page." + ), + "low_content": ( + "Page with very little extractable text and little visual/table content; " + "may be blank, separator, short heading page, or sparse transition page." + ), + "landscape": "Landscape-oriented page, often wide tables, drawings, slides, or diagrams.", +} + + +def _feature_rows(ctx: ToolContext, pages: list[int]) -> list[dict[str, Any]]: + labels_by_page = {label.page: label for label in ctx.blackboard.page_labels} + selected = [] + for feature in ctx.blackboard.page_features: + if feature.page not in pages: + continue + label = labels_by_page.get(feature.page) + selected.append( + { + "page": feature.page, + "kind": label.kind if label else None, + "confidence": label.confidence if label else None, + "raw_text_length": feature.raw_text_length, + "text_density": feature.text_density, + "image_coverage": feature.image_coverage, + "image_count": feature.image_count, + "table_count": feature.table_count, + "drawings_count": feature.drawings_count, + "orientation": feature.orientation, + "is_blank_like": feature.is_blank_like, + } + ) + return selected + + +def _segment_sample(candidates: list[int], count: int) -> list[int]: + if count <= 0 or not candidates: + return [] + if len(candidates) <= count: + return candidates + if count == 1: + return [candidates[len(candidates) // 2]] + step = (len(candidates) - 1) / (count - 1) + return [candidates[round(index * step)] for index in range(count)] + + +def _sample_pages(page_count: int, extrema_pages: list[int]) -> list[int]: + if page_count <= 0: + return [] + extrema = [page for page in extrema_pages if 1 <= page <= page_count] + remaining = [page for page in range(1, page_count + 1) if page not in set(extrema)] + if not remaining: + return sorted(set(extrema)) + third = max(len(remaining) // 3, 1) + front = remaining[:third] + middle = remaining[third : third * 2] + back = remaining[third * 2 :] + sampled = ( + _segment_sample(front, 4) + + _segment_sample(middle or remaining, 3) + + _segment_sample(back or remaining, 3) + ) + ordered = [] + for page in extrema + sampled: + if page not in ordered: + ordered.append(page) + return ordered[:20] + + +def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDecision]: + data = json.loads(raw) + if not isinstance(data, dict): + raise ValueError("planner output must be a JSON object") + category = " ".join(str(data.get("category") or "unknown document").split()[:5]) + raw_is_scanned = data.get("is_scanned") + if isinstance(raw_is_scanned, bool): + is_scanned = raw_is_scanned + elif isinstance(raw_is_scanned, str): + is_scanned = raw_is_scanned.strip().lower() in {"true", "yes", "1", "scanned"} + else: + is_scanned = bool(raw_is_scanned) + profile = DocumentProfile( + is_scanned=is_scanned, + category=category or "unknown document", + category_rationale=str(data.get("category_rationale") or ""), + language=str(data.get("language") or "unknown"), + rationale=str(data.get("rationale") or ""), + ) + next_action = str(data.get("next_action") or "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: + query = str(data.get("grep_query") or "").strip() + if query: + tool_name = "grep.text" + tool_args = {"query": query, "max_results": 20} + elif next_action == "verdict_now": + return profile, ReflexionDecision( + action="verdict_now", + rationale=profile.rationale, + verdict=None, + ) + if tool_name: + return profile, ReflexionDecision( + action="tool_call", + rationale=profile.rationale, + tool_name=tool_name, + tool_args=tool_args, + ) + return profile, ReflexionDecision( + action="tool_call", + rationale=profile.rationale, + tool_name="propose.shard_plan", + tool_args={}, + ) + + +class ProfilePlanner: + """One-shot VLM planner that profiles the document and proposes the first action.""" + + def __init__(self, ctx: ToolContext) -> None: + self.ctx = ctx + + def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: + start = time.monotonic() + model = ( + self.ctx.settings.get("planner_model") + or self.ctx.settings.get("vlm_model") + or os.environ.get("IMAGE_MODEL") + ) + pages = _sample_pages(self.ctx.blackboard.page_count, self.ctx.blackboard.extrema_pages) + pngs = render_pages( + self.ctx, + pages, + folder_name="planner_pages", + prefix="planner", + timeout=180, + ) + feature_summary = _feature_rows(self.ctx, pages) + payload = { + "page_count": self.ctx.blackboard.page_count, + "page_kind_counts": self.ctx.blackboard.global_signals.get( + "page_kind_counts", + {}, + ), + "page_kind_definitions": PAGE_KIND_DEFINITIONS, + "doc_stats": self.ctx.blackboard.doc_stats, + "extrema_samples": self.ctx.blackboard.global_signals.get( + "extrema_samples", + [], + ), + "sampled_page_features": feature_summary, + "toc_pages": self.ctx.blackboard.toc_result.toc_pages + if self.ctx.blackboard.toc_result + else [], + "h1_pages": [ + {"title": item.title, "page": item.page} + for item in ( + self.ctx.blackboard.h1_result.h1_candidates + if self.ctx.blackboard.h1_result + else [] + ) + ], + "available_actions": [ + "inspect.pages", + "grep.text", + "propose.shard_plan", + "validate.anatomy_map", + "verdict", + ], + } + prompt_text = PLANNER_INSTRUCTIONS + "\nPayload:\n" + json.dumps( + payload, + ensure_ascii=False, + ) + prompt_tokens_est = estimate_tokens(prompt_text) + len(pngs) * 800 + if not model: + profile = DocumentProfile( + is_scanned=False, + category="unknown document", + rationale="No planner model configured.", + ) + decision = ReflexionDecision( + action="tool_call", + rationale=profile.rationale, + tool_name="propose.shard_plan", + tool_args={}, + ) + return profile, decision, ToolResult( + status="ok", + payload={"source": "deterministic", "sampled_pages": pages}, + latency_ms=int((time.monotonic() - start) * 1000), + warnings=["No planner model configured; using conservative profile."], + input_summary=payload, + output_summary={"profile": profile.to_dict(), "decision": decision.to_dict()}, + ) + if not self.ctx.budget.try_reserve("plan", prompt_tokens_est): + raise RuntimeError("Insufficient planner budget for profile planning.") + + content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt_text}] + for item in pngs: + try: + with open(str(item["png_path"]), "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + content_parts.append( + {"type": "text", "text": f"\n--- Page {item['page']} ---"} + ) + content_parts.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + } + ) + except Exception as exc: + logger.warning("[document_agent] planner png attach failed: {}", exc) + + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=model, + temperature=0.0, + max_tokens=1800, + response_format={"type": "json_object"}, + ) + self.ctx.budget.commit( + "plan", + actual=usage.get("total_tokens", prompt_tokens_est), + est=prompt_tokens_est, + ) + profile, decision = _parse_profile_and_decision(raw) + return profile, decision, ToolResult( + status="ok", + payload={ + "source": "llm", + "sampled_pages": pages, + "first_action": decision.tool_name, + }, + latency_ms=int((time.monotonic() - start) * 1000), + tokens_used=usage.get("total_tokens", 0), + input_summary=payload, + output_summary={"profile": profile.to_dict(), "decision": decision.to_dict()}, + debug={ + "prompt_text": prompt_text, + "sampled_pngs": pngs, + "raw_response": raw, + }, + ) + except Exception: + self.ctx.budget.refund("plan", est=prompt_tokens_est) + raise + diff --git a/apps/worker/app/services/document_agent/planner/prompts.py b/apps/worker/app/services/document_agent/planner/prompts.py new file mode 100644 index 000000000..49feab396 --- /dev/null +++ b/apps/worker/app/services/document_agent/planner/prompts.py @@ -0,0 +1,16 @@ +"""Prompts for the document profile planner.""" + +PLANNER_INSTRUCTIONS = ( + "You are a document profile agent. Use global page-feature statistics, " + "TOC/H1 evidence, and page screenshots to classify the document and decide " + "whether enough evidence exists to continue toward sharding. Return strict " + "JSON only with keys: is_scanned, category, category_rationale, language, " + "rationale, next_action, inspect_pages, grep_query. category must be at " + "most 5 English words. next_action must be one of inspect_more, grep_text, " + "ready_to_shard, verdict_now. Use inspect_more only when specific extra " + "page screenshots are needed. Use grep_text only for native PDFs when a " + "global text search would clarify structure. Do not output a fixed step " + "plan." +) + +__all__ = ["PLANNER_INSTRUCTIONS"] diff --git a/apps/worker/app/services/document_agent/prompts/coordinator_system.md b/apps/worker/app/services/document_agent/prompts/coordinator_system.md deleted file mode 100644 index dcfb075d0..000000000 --- a/apps/worker/app/services/document_agent/prompts/coordinator_system.md +++ /dev/null @@ -1,12 +0,0 @@ -You are coordinating page anatomy profiling for a PDF. - -Goal: -- identify special pages that affect hierarchy extraction -- identify reliable H1 starts and pages that should not become title candidates -- produce a shard plan for long PDF-to-Markdown execution - -Rules: -- use only the available tools for the current state -- do not invent page numbers or section titles -- prefer evidence from TOC pages and page-start heading matches -- treat tool outputs as structured evidence, not prose diff --git a/apps/worker/app/services/document_agent/prompts/propose_hierarchy_assist.md b/apps/worker/app/services/document_agent/prompts/propose_hierarchy_assist.md deleted file mode 100644 index be6803aef..000000000 --- a/apps/worker/app/services/document_agent/prompts/propose_hierarchy_assist.md +++ /dev/null @@ -1,11 +0,0 @@ -Return strict JSON for hierarchy assistance. - -Use only page numbers present in the input payload. -Do not invent headings. -Keep reasons brief and evidence-based. - -Required fields: -- exclude_pages_from_title_candidates: integer array -- suppress_title_pages: integer array -- smart_parse_recommendation: one of off, normal, aggressive -- rationale: short string diff --git a/apps/worker/app/services/document_agent/registry.py b/apps/worker/app/services/document_agent/registry.py index 5bf9d4cd2..80a17531f 100644 --- a/apps/worker/app/services/document_agent/registry.py +++ b/apps/worker/app/services/document_agent/registry.py @@ -1,4 +1,4 @@ -"""Tool registry with state-based exposure.""" +"""Agent tool registry with blackboard-based preconditions.""" from __future__ import annotations @@ -6,9 +6,14 @@ from typing import Any, Callable from app.services.document_agent.manifest import ToolContext, ToolResult -from app.services.document_agent.state import DocumentAgentState +from app.services.document_agent.state import AgentBlackboard ToolHandler = Callable[[ToolContext, dict[str, Any]], ToolResult] +Precondition = Callable[[AgentBlackboard], tuple[bool, str]] + + +def _always(_blackboard: AgentBlackboard) -> tuple[bool, str]: + return True, "" @dataclass(frozen=True) @@ -16,7 +21,7 @@ class ToolSpec: name: str description: str parameters: dict[str, Any] - allowed_states: frozenset[DocumentAgentState] + preconditions: tuple[Precondition, ...] handler: ToolHandler def to_openai_schema(self) -> dict[str, Any]: @@ -40,35 +45,50 @@ def register(self, spec: ToolSpec) -> None: def get(self, name: str) -> ToolSpec | None: return self._tools.get(name) - def catalogue_for_state(self, state: DocumentAgentState) -> list[dict[str, Any]]: + def openai_specs(self, blackboard: AgentBlackboard) -> list[dict[str, Any]]: return [ tool.to_openai_schema() for tool in self._tools.values() - if state in tool.allowed_states + if self._preconditions_met(tool, blackboard)[0] ] - def allowed_names(self, state: DocumentAgentState) -> list[str]: + def allowed_names(self, blackboard: AgentBlackboard) -> list[str]: return [ name for name, tool in self._tools.items() - if state in tool.allowed_states + if self._preconditions_met(tool, blackboard)[0] ] + def _preconditions_met( + self, + tool: ToolSpec, + blackboard: AgentBlackboard, + ) -> tuple[bool, str]: + for check in tool.preconditions: + ok, reason = check(blackboard) + if not ok: + return False, reason + return True, "" + def dispatch( self, name: str, ctx: ToolContext, args: dict[str, Any], - state: DocumentAgentState, ) -> ToolResult: tool = self.get(name) if tool is None: return ToolResult(status="error", error=f"unknown tool: {name}") - if state not in tool.allowed_states: + ok, reason = self._preconditions_met(tool, ctx.blackboard) + if not ok: return ToolResult( - status="state_error", - payload={"allowed_tools": self.allowed_names(state), "state": state.value}, - error=f"tool {name} is not allowed in state {state.value}", + status="precondition_unmet", + payload={ + "allowed_tools": self.allowed_names(ctx.blackboard), + "tool": name, + "reason": reason, + }, + error=reason, ) return tool.handler(ctx, args) @@ -81,7 +101,7 @@ def register_tool( name: str, description: str, parameters: dict[str, Any] | None = None, - allowed_states: set[DocumentAgentState], + preconditions: tuple[Precondition, ...] | None = None, ) -> Callable[[ToolHandler], ToolHandler]: def _decorator(handler: ToolHandler) -> ToolHandler: REGISTRY.register( @@ -90,10 +110,54 @@ def _decorator(handler: ToolHandler) -> ToolHandler: description=description, parameters=parameters or {"type": "object", "properties": {}, "required": []}, - allowed_states=frozenset(allowed_states), + preconditions=preconditions or (_always,), handler=handler, ) ) return handler return _decorator + + +def has_page_features(blackboard: AgentBlackboard) -> tuple[bool, str]: + return bool(blackboard.page_features), "page_features missing; run bootstrap probe first" + + +def has_page_labels(blackboard: AgentBlackboard) -> tuple[bool, str]: + return bool(blackboard.page_labels), "page_labels missing; run bootstrap classify first" + + +def has_doc_stats(blackboard: AgentBlackboard) -> tuple[bool, str]: + return bool(blackboard.doc_stats), "doc_stats missing; run bootstrap aggregate first" + + +def has_document_profile(blackboard: AgentBlackboard) -> tuple[bool, str]: + return blackboard.document_profile is not None, "document_profile missing; run planner first" + + +def not_is_scanned(blackboard: AgentBlackboard) -> tuple[bool, str]: + profile = blackboard.document_profile + return ( + profile is not None and not profile.is_scanned, + "document is scanned or profile is missing; text grep is unavailable", + ) + + +def has_toc_anchors(blackboard: AgentBlackboard) -> tuple[bool, str]: + return bool(blackboard.toc_anchor_pages), "toc anchors missing; call find_toc_anchors first" + + +def has_toc_result(blackboard: AgentBlackboard) -> tuple[bool, str]: + return blackboard.toc_result is not None, "toc_result missing; call extract_toc first" + + +def has_toc_hierarchies(blackboard: AgentBlackboard) -> tuple[bool, str]: + return bool(blackboard.toc_hierarchies), "toc_hierarchies missing; call extract_toc first" + + +def has_h1_result(blackboard: AgentBlackboard) -> tuple[bool, str]: + return blackboard.h1_result is not None, "h1_result missing; call match_h1 first" + + +def has_shard_plan(blackboard: AgentBlackboard) -> tuple[bool, str]: + return blackboard.shard_plan is not None, "shard_plan missing; call propose_shard first" diff --git a/apps/worker/app/services/document_agent/state.py b/apps/worker/app/services/document_agent/state.py index f54278fa6..bcdb1a1c5 100644 --- a/apps/worker/app/services/document_agent/state.py +++ b/apps/worker/app/services/document_agent/state.py @@ -7,9 +7,9 @@ from typing import Any from app.services.document_agent.manifest import ( - BoundaryCandidate, + AgentVerdict, + DocumentProfile, H1BoundaryResult, - HierarchyAssistPlan, PageFeature, PageLabel, ShardPlan, @@ -20,34 +20,28 @@ class DocumentAgentState(str, Enum): INIT = "init" - PROBED = "probed" - CLASSIFIED = "classified" - H1_FOUND = "h1_found" - PLANNED = "planned" - VALIDATED = "validated" - PERSISTED = "persisted" + RUNNING = "running" READY = "ready" FAILED = "failed" - PROCESSING_PLAN_PROPOSED = "processing_plan_proposed" @dataclass class AgentBlackboard: page_count: int = 0 + document_profile: DocumentProfile | None = None page_features: list[PageFeature] = field(default_factory=list) page_labels: list[PageLabel] = field(default_factory=list) + doc_stats: dict[str, Any] = field(default_factory=dict) + extrema_pages: list[int] = field(default_factory=list) toc_anchor_pages: list[TocAnchorPage] = field(default_factory=list) toc_result: TocResult | None = None toc_hierarchies: list[dict[str, Any]] | None = None h1_result: H1BoundaryResult | None = None - boundary_candidates: list[BoundaryCandidate] = field(default_factory=list) - hierarchy_assist: HierarchyAssistPlan | None = None shard_plan: ShardPlan | None = None validation_report: dict[str, Any] | None = None + verdict: AgentVerdict | None = None + step_history: list[dict[str, Any]] = field(default_factory=list) + page_full_text_cache: dict[int, str] = field(default_factory=dict) global_signals: dict[str, Any] = field(default_factory=dict) errors: list[str] = field(default_factory=list) - state_trace: list[str] = field(default_factory=list) - - def mark(self, state: DocumentAgentState) -> None: - self.state_trace.append(state.value) diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index d212eaa14..00260327d 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -2,15 +2,13 @@ from app.services.document_agent.registry import REGISTRY -from . import classify_page_kinds as classify_page_kinds # noqa: F401 -from . import collect_boundary_candidates as collect_boundary_candidates # noqa: F401 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 match_h1_pages as match_h1_pages # noqa: F401 -from . import persist_anatomy_map as persist_anatomy_map # noqa: F401 -from . import probe_page_features as probe_page_features # noqa: F401 -from . import propose_hierarchy_assist as propose_hierarchy_assist # 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 __all__ = ["REGISTRY"] diff --git a/apps/worker/app/services/document_agent/tools/classify_page_kinds.py b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py index 44fb98f41..bd7daf67e 100644 --- a/apps/worker/app/services/document_agent/tools/classify_page_kinds.py +++ b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py @@ -7,8 +7,6 @@ from typing import Any from app.services.document_agent.manifest import PageFeature, PageLabel, ToolContext, ToolResult -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState def _joined_preview(feature: PageFeature) -> str: @@ -16,20 +14,16 @@ def _joined_preview(feature: PageFeature) -> str: def _label_feature(feature: PageFeature) -> PageLabel: - preview = _joined_preview(feature) page = feature.page - if any(marker in preview.replace(" ", "") for marker in ("目录", "目次", "contents")): + if ( + feature.raw_text_length < 80 + and feature.image_coverage < 0.02 + and feature.drawings_count < 5 + ): return PageLabel( page=page, - kind="toc", - confidence=0.86, - evidence={"signal": "toc_marker"}, - ) - if feature.is_blank_like: - return PageLabel( - page=page, - kind="blank", - confidence=0.92, + kind="low_content", + confidence=0.78, evidence={"signal": "low_text_image_drawings"}, ) if feature.orientation == "landscape": @@ -42,7 +36,7 @@ def _label_feature(feature: PageFeature) -> PageLabel: if feature.image_coverage >= 0.35 and feature.raw_text_length < 250: return PageLabel( page=page, - kind="single_image", + kind="image_heavy", confidence=0.84, evidence={"image_coverage": feature.image_coverage}, ) @@ -56,21 +50,9 @@ def _label_feature(feature: PageFeature) -> PageLabel: "drawings_count": feature.drawings_count, }, ) - if feature.raw_text_length < 80: - return PageLabel( - page=page, - kind="sparse", - confidence=0.67, - evidence={"raw_text_length": feature.raw_text_length}, - ) return PageLabel(page=page, kind="normal", confidence=0.65, evidence={}) -@register_tool( - name="classify.page_kinds", - description="Classify every probed page into structural page kinds using deterministic signals.", - allowed_states={DocumentAgentState.PROBED}, -) def classify_page_kinds(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() labels = [_label_feature(feature) for feature in ctx.blackboard.page_features] diff --git a/apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py b/apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py deleted file mode 100644 index 89e72610f..000000000 --- a/apps/worker/app/services/document_agent/tools/collect_boundary_candidates.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Collect candidate split pages without deciding where to split.""" - -from __future__ import annotations - -import time -from collections import Counter -from typing import Any - -from app.services.document_agent.manifest import ( - BoundaryCandidate, - H1BoundaryResult, - TocResult, - ToolContext, - ToolResult, -) -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState - - -BOUNDARY_PAGE_KINDS = {"blank", "sparse", "toc"} -PRIORITY_BY_KIND = { - "h1": 100, - "toc": 80, - "blank": 45, - "sparse": 40, -} - - -def _feature_by_page(ctx: ToolContext) -> dict[int, Any]: - return {feature.page: feature for feature in ctx.blackboard.page_features} - - -def _candidate_evidence(ctx: ToolContext, page: int, kind: str) -> dict[str, Any]: - feature = _feature_by_page(ctx).get(page) - if feature is None: - return {} - return { - "source": "page_label", - "label_kind": kind, - "position_ratio": round(page / max(ctx.blackboard.page_count, 1), 4), - "raw_text_length": feature.raw_text_length, - "image_coverage": feature.image_coverage, - "table_count": feature.table_count, - "drawings_count": feature.drawings_count, - "text_preview": feature.text_lines_preview[:5], - } - - -@register_tool( - name="collect.boundary_candidates", - description="Collect sparse, blank, TOC, and H1 pages as split candidates without making split decisions.", - allowed_states={DocumentAgentState.H1_FOUND}, -) -def collect_boundary_candidates(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: - start = time.monotonic() - candidates: list[BoundaryCandidate] = [] - seen: set[tuple[int, str]] = set() - - # Upstream tools (extract.toc_with_boundaries, match.h1_pages) populate - # these before collect runs. Provide safe defaults if they were skipped. - if ctx.blackboard.toc_result is None: - ctx.blackboard.toc_result = TocResult(method="none") - if ctx.blackboard.h1_result is None: - ctx.blackboard.h1_result = H1BoundaryResult(method="none") - - for label in ctx.blackboard.page_labels: - if label.kind not in BOUNDARY_PAGE_KINDS: - continue - key = (label.page, label.kind) - if key in seen: - continue - seen.add(key) - candidates.append( - BoundaryCandidate( - page=label.page, - kind=label.kind, # type: ignore[arg-type] - priority=PRIORITY_BY_KIND.get(label.kind, 0), - confidence=label.confidence, - evidence=_candidate_evidence(ctx, label.page, label.kind), - ) - ) - - # Add H1 candidates from upstream match.h1_pages as highest-priority - # boundary hints for shard planning. - if ctx.blackboard.h1_result: - for h1 in ctx.blackboard.h1_result.h1_candidates: - key = (h1.page, "h1") - if key in seen: - continue - seen.add(key) - candidates.append( - BoundaryCandidate( - page=h1.page, - kind="h1", - priority=PRIORITY_BY_KIND["h1"], - confidence=h1.confidence, - evidence={ - "source": h1.source, - "title": h1.title, - "matched_line": h1.matched_line, - "position_ratio": round( - h1.page / max(ctx.blackboard.page_count, 1), 4 - ), - **h1.evidence, - }, - ) - ) - - candidates.sort(key=lambda item: (item.page, -item.priority)) - ctx.blackboard.boundary_candidates = candidates - counts = Counter(candidate.kind for candidate in candidates) - ctx.blackboard.global_signals["boundary_candidate_counts"] = dict(counts) - - return ToolResult( - status="ok", - payload={ - "candidate_count": len(candidates), - "candidate_counts": dict(counts), - }, - latency_ms=int((time.monotonic() - start) * 1000), - input_summary={ - "page_count": ctx.blackboard.page_count, - "page_kind_counts": ctx.blackboard.global_signals.get("page_kind_counts", {}), - }, - output_summary={ - "candidate_counts": dict(counts), - "sample_candidates": [candidate.to_dict() for candidate in candidates[:20]], - }, - ) diff --git a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py index fe0dc617c..ac457ab41 100644 --- a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py +++ b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py @@ -1,14 +1,13 @@ -"""VLM-driven progressive TOC boundary detection + mineru local MD + toc_parser reuse.""" +"""VLM-driven TOC anchor, boundary, and entry extraction.""" from __future__ import annotations import gc import json import os -import subprocess import time from pathlib import Path -from typing import Any +from typing import Any, cast from app.services.document_agent.manifest import ( TocAnchorPage, @@ -17,7 +16,10 @@ ToolResult, ) from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState +from app.services.document_agent.tools.vlm_toc_extractor import ( + vlm_entries_to_toc_hierarchies, + vlm_extract_toc_entries, +) from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( run_in_child_process, worker, @@ -29,7 +31,6 @@ BOUNDARY_STEP_PAGES = 5 MAX_BOUNDARY_ROUNDS = 6 MAX_TOC_PAGES = BOUNDARY_STEP_PAGES * MAX_BOUNDARY_ROUNDS # 30 -MINERU_TIMEOUT_SECONDS = 180 # -- PyMuPDF workers (must be top-level for multiprocessing pickle) ------------ @@ -64,12 +65,12 @@ def _render_single_page_worker( def _vlm_confirm_anchors( anchor_pages: list[TocAnchorPage], model: str, -) -> list[TocAnchorPage]: +) -> tuple[list[TocAnchorPage], bool]: """Phase 1: send all anchor PNGs to VLM, ask which are real TOC starts.""" from shared.services.ai.openai_compatible_client_sync import get_openai_client if not anchor_pages: - return [] + return [], False import base64 @@ -113,7 +114,7 @@ def _vlm_confirm_anchors( } ) - messages = [{"role": "user", "content": content_parts}] + messages = cast(Any, [{"role": "user", "content": content_parts}]) try: client = get_openai_client(model=model) @@ -146,21 +147,21 @@ def _vlm_confirm_anchors( len(confirmed), rejected, ) - return confirmed + return confirmed, False except Exception as exc: logger.warning( "[extract.toc] VLM anchor confirmation failed: {}, " "falling back to no confirmed anchors (safe degradation)", exc, ) - return [] + return [], True def _vlm_check_boundary_page( png_path: str, page_num: int, model: str, -) -> bool: +) -> tuple[bool, bool]: """Phase 2: check if a single page still contains TOC content.""" from shared.services.ai.openai_compatible_client_sync import get_openai_client @@ -200,20 +201,20 @@ def _vlm_check_boundary_page( try: client = get_openai_client(model=model) raw, _ = client.chat_completion_with_usage( - messages=[{"role": "user", "content": content_parts}], + messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, temperature=0.1, max_tokens=200, response_format={"type": "json_object"}, ) data = json.loads(raw) - return bool(data.get("still_toc")) + return bool(data.get("still_toc")), False except Exception as exc: logger.warning( "[extract.toc] boundary VLM check failed for page {}: {}", page_num, exc ) # Conservative: stop expansion on failure - return False + return False, True # -- Progressive boundary detection -------------------------------------------- @@ -227,7 +228,7 @@ def _detect_toc_range_for_anchor( output_dir: str, dpi: int, model: str, -) -> tuple[int, int, list[dict[str, Any]]]: +) -> tuple[int, int, list[dict[str, Any]], list[str]]: """Progressively expand from anchor_page to find the TOC end boundary. Returns: @@ -236,6 +237,7 @@ def _detect_toc_range_for_anchor( start_page = anchor_page current_end = min(anchor_page + BOUNDARY_STEP_PAGES - 1, page_count) trace_rounds: list[dict[str, Any]] = [] + warnings: list[str] = [] for round_idx in range(MAX_BOUNDARY_ROUNDS): check_page = current_end @@ -249,7 +251,9 @@ def _detect_toc_range_for_anchor( timeout=60, ) - still_toc = _vlm_check_boundary_page(png_path, check_page, model) + still_toc, failed = _vlm_check_boundary_page(png_path, check_page, model) + if failed: + warnings.append(f"vlm_boundary_check_failed:p{check_page}") trace_rounds.append( { "round": round_idx, @@ -266,7 +270,7 @@ def _detect_toc_range_for_anchor( ) if not still_toc: - # Boundary page is NOT TOC; TOC ends at previous page + # Boundary page is NOT TOC; TOC ends at previous page. current_end = max(check_page - 1, start_page) break @@ -275,70 +279,7 @@ def _detect_toc_range_for_anchor( break current_end = next_end - return start_page, current_end, trace_rounds - - -# -- MinerU local extraction --------------------------------------------------- - - -def _run_mineru_local( - pdf_path: str, - start_page_0based: int, - end_page_0based: int, - output_dir: str, -) -> list[str]: - """Run mineru CLI on a page range and return the resulting markdown lines.""" - os.makedirs(output_dir, exist_ok=True) - cmd = [ - "mineru", - "-p", - pdf_path, - "-o", - output_dir, - "-s", - str(start_page_0based), - "-e", - str(end_page_0based), - "-t", - "false", # skip table parsing for speed - "-b", - "pipeline", - "-m", - "txt", - ] - logger.info( - "[extract.toc] mineru local: pages {}-{}, cmd: {}", - start_page_0based, - end_page_0based, - " ".join(cmd), - ) - - try: - proc = subprocess.run( - cmd, - capture_output=True, - text=True, - timeout=MINERU_TIMEOUT_SECONDS, - ) - if proc.returncode != 0: - logger.error( - "[extract.toc] mineru failed (code={}): {}", - proc.returncode, - proc.stderr[-500:] if proc.stderr else "", - ) - return [] - except subprocess.TimeoutExpired: - logger.error("[extract.toc] mineru timed out after {}s", MINERU_TIMEOUT_SECONDS) - return [] - - md_lines: list[str] = [] - output_path = Path(output_dir) - for md_file in sorted(output_path.rglob("*.md")): - content = md_file.read_text(encoding="utf-8", errors="replace") - md_lines.extend(content.splitlines()) - - logger.info("[extract.toc] mineru produced {} markdown lines", len(md_lines)) - return md_lines + return start_page, current_end, trace_rounds, warnings # -- Main tool ----------------------------------------------------------------- @@ -348,9 +289,8 @@ def _run_mineru_local( name="extract.toc_with_boundaries", description=( "VLM-confirms TOC anchor pages, progressively detects TOC boundaries, " - "runs mineru local extraction, then reuses toc_parser for hierarchy." + "then extracts TOC entries directly from rendered pages with VLM." ), - allowed_states={DocumentAgentState.CLASSIFIED}, ) def extract_toc_with_boundaries( ctx: ToolContext, _args: dict[str, Any] @@ -384,7 +324,9 @@ def extract_toc_with_boundaries( os.makedirs(output_dir, exist_ok=True) # -- Phase 1: VLM confirm anchors ----------------------------------------- - confirmed = _vlm_confirm_anchors(anchors, model) + confirmed, confirm_failed = _vlm_confirm_anchors(anchors, model) + if confirm_failed: + warnings.append("vlm_anchor_confirmation_failed") debug_info["phase1_confirmed"] = [a.page for a in confirmed] debug_info["phase1_rejected"] = [ a.page for a in anchors if a not in confirmed @@ -408,7 +350,7 @@ def extract_toc_with_boundaries( all_trace_rounds: list[dict[str, Any]] = [] for anchor in confirmed: - toc_start, toc_end, trace_rounds = _detect_toc_range_for_anchor( + toc_start, toc_end, trace_rounds, boundary_warnings = _detect_toc_range_for_anchor( anchor_page=anchor.page, pdf_path=ctx.pdf_path, page_count=page_count, @@ -416,6 +358,7 @@ def extract_toc_with_boundaries( dpi=dpi, model=model, ) + warnings.extend(boundary_warnings) toc_ranges.append((toc_start, toc_end)) all_trace_rounds.extend(trace_rounds) logger.info( @@ -425,66 +368,49 @@ def extract_toc_with_boundaries( debug_info["phase2_ranges"] = toc_ranges debug_info["phase2_trace_rounds"] = all_trace_rounds - # -- Phase 3: mineru local extraction -------------------------------------- - all_md_lines: list[str] = [] - for i, (toc_start, toc_end) in enumerate(toc_ranges): - region_dir = os.path.join(output_dir, f"mineru_region_{i}") - md_lines = _run_mineru_local( - pdf_path=ctx.pdf_path, - start_page_0based=toc_start - 1, # mineru uses 0-based - end_page_0based=toc_end - 1, - output_dir=region_dir, - ) - if md_lines: - all_md_lines.extend(md_lines) - else: - warnings.append( - f"mineru produced no output for region {i} (pages {toc_start}-{toc_end})" - ) - - debug_info["phase3_md_line_count"] = len(all_md_lines) - - if not all_md_lines: - ctx.blackboard.toc_result = TocResult( - toc_pages=[p for s, e in toc_ranges for p in range(s, e + 1)], - method="vlm_progressive", - notes="VLM detected TOC ranges but mineru produced no markdown", - ) - warnings.append("mineru produced no markdown for any TOC region") - return ToolResult( - status="ok", - payload={"toc_count": 0}, - latency_ms=int((time.monotonic() - start) * 1000), - warnings=warnings, - debug=debug_info, + # -- Phase 3: VLM entry extraction ----------------------------------------- + all_toc_pages = sorted( + {p for s, e in toc_ranges for p in range(s, e + 1)} + ) + all_entries: list[dict[str, Any]] = [] + per_page_meta: list[dict[str, Any]] = [] + rendered_pages: list[dict[str, Any]] = [] + for page_num in all_toc_pages: + png_path = os.path.join(output_dir, f"toc_page_{page_num}.png") + render_result = run_in_child_process( + _render_single_page_worker, + ctx.pdf_path, + page_num, + png_path, + dpi, + timeout=60, ) - - # -- Phase 4: toc_parser reuse --------------------------------------------- - try: - from app.services.document_parser.structure.toc_parser import ( - detect_tocs_in_texts, + rendered_pages.append(render_result) + entries, meta = vlm_extract_toc_entries( + png_path=str(render_result.get("png_path") or png_path), + page_num=page_num, + model=model, + previous_entries=all_entries, ) + all_entries.extend(entries) + per_page_meta.append(meta) - hierarchy_model = ctx.settings.get("model") or os.environ.get( - "HIERARCHY_LLM_MODEL" - ) or os.environ.get("NORMOL_MODEL") - - toc_hierarchies, _filtered = detect_tocs_in_texts( - all_md_lines, - model_name=hierarchy_model, - hierarchy_model_name=hierarchy_model, - branch="normal", - limit_=150, - ) - except Exception as exc: - logger.error("[extract.toc] toc_parser failed: {}", exc) - toc_hierarchies = None - warnings.append(f"toc_parser failed: {exc}") + if not all_entries: + raise RuntimeError("VLM TOC extractor returned no entries for confirmed TOC pages") - # -- Write results to blackboard ------------------------------------------- - all_toc_pages = sorted( - {p for s, e in toc_ranges for p in range(s, e + 1)} + scan_end_page = max( + (round_info.get("check_page", 0) for round_info in all_trace_rounds), + default=max(all_toc_pages), ) + toc_hierarchies = vlm_entries_to_toc_hierarchies( + all_entries, + toc_page_nums=all_toc_pages, + scan_end_page=int(scan_end_page), + page_count=page_count, + ) + debug_info["phase3_vlm_entry_count"] = len(all_entries) + debug_info["phase3_vlm_per_page_meta"] = per_page_meta + debug_info["phase3_rendered_pages"] = rendered_pages ctx.blackboard.toc_result = TocResult( toc_pages=all_toc_pages, @@ -495,6 +421,13 @@ def extract_toc_with_boundaries( ), ) ctx.blackboard.toc_hierarchies = toc_hierarchies if toc_hierarchies else None + ctx.blackboard.global_signals["vlm_toc_entries"] = { + "model": model, + "toc_pages": all_toc_pages, + "total_entries": len(all_entries), + "entries": all_entries, + "per_page_meta": per_page_meta, + } # Persist toc_hierarchies to disk for inspection / downstream reuse if toc_hierarchies and ctx.output_dir: @@ -509,6 +442,8 @@ def extract_toc_with_boundaries( toc_summary: dict[str, Any] = { "toc_ranges": toc_ranges, "toc_page_count": len(all_toc_pages), + "toc_entry_count": len(all_entries), + "toc_source": "vlm", } if toc_hierarchies: for i, hier in enumerate(toc_hierarchies): diff --git a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py index 6d9bc4437..f6cfdae17 100644 --- a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py +++ b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py @@ -5,13 +5,11 @@ import gc import os import time -from collections import Counter from pathlib import Path from typing import Any from app.services.document_agent.manifest import TocAnchorPage, ToolContext, ToolResult -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState +from app.services.document_agent.registry import has_page_labels, register_tool from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( run_in_child_process, worker, @@ -123,7 +121,7 @@ def _filter_recurring_elements( "Scan page text previews for TOC keywords, filter recurring " "navigation elements, then render candidate PNGs for VLM confirmation." ), - allowed_states={DocumentAgentState.CLASSIFIED}, + preconditions=(has_page_labels,), ) def find_toc_anchor_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() diff --git a/apps/worker/app/services/document_agent/tools/grep_text.py b/apps/worker/app/services/document_agent/tools/grep_text.py new file mode 100644 index 000000000..6c59f5268 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/grep_text.py @@ -0,0 +1,79 @@ +"""Generic full-document text grep for native PDFs.""" + +from __future__ import annotations + +import re +import time +from typing import Any + +from app.services.document_agent.manifest import ToolContext, ToolResult +from app.services.document_agent.pdf_text import read_page_texts +from app.services.document_agent.registry import has_page_features, not_is_scanned, register_tool + + +def _load_page_texts(ctx: ToolContext) -> dict[int, str]: + if ctx.blackboard.page_full_text_cache: + return dict(ctx.blackboard.page_full_text_cache) + pages = list(range(1, ctx.blackboard.page_count + 1)) + texts = read_page_texts(ctx.pdf_path, pages, timeout=300) + ctx.blackboard.page_full_text_cache = texts + return texts + + +@register_tool( + name="grep.text", + description="Search full PDF text for a substring or regex. Available only for native PDFs.", + parameters={ + "type": "object", + "properties": { + "query": {"type": "string"}, + "regex": {"type": "boolean", "default": False}, + "case_sensitive": {"type": "boolean", "default": False}, + "max_results": {"type": "integer", "default": 30}, + "context_chars": {"type": "integer", "default": 80}, + }, + "required": ["query"], + }, + preconditions=(has_page_features, not_is_scanned), +) +def grep_text(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + query = str(args.get("query") or "").strip() + if not query: + return ToolResult( + status="error", + error="grep.text requires query", + latency_ms=int((time.monotonic() - start) * 1000), + ) + use_regex = bool(args.get("regex", False)) + case_sensitive = bool(args.get("case_sensitive", False)) + max_results = max(1, min(int(args.get("max_results") or 30), 100)) + context_chars = max(20, min(int(args.get("context_chars") or 80), 300)) + flags = 0 if case_sensitive else re.IGNORECASE + pattern = re.compile(query if use_regex else re.escape(query), flags) + results: list[dict[str, Any]] = [] + for page, text in sorted(_load_page_texts(ctx).items()): + for match in pattern.finditer(text): + start_idx = max(match.start() - context_chars, 0) + end_idx = min(match.end() + context_chars, len(text)) + results.append( + { + "page": page, + "char_offset": match.start(), + "snippet": text[start_idx:end_idx].replace("\n", " "), + } + ) + if len(results) >= max_results: + break + if len(results) >= max_results: + break + summary = {"query": query, "hit_count": len(results), "results": results} + ctx.blackboard.global_signals.setdefault("grep_history", []).append( + {"query": query, "hit_count": len(results), "sample_pages": [item["page"] for item in results[:10]]} + ) + return ToolResult( + status="ok", + payload=summary, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={"query": query, "hit_count": len(results)}, + ) diff --git a/apps/worker/app/services/document_agent/tools/inspect_pages.py b/apps/worker/app/services/document_agent/tools/inspect_pages.py new file mode 100644 index 000000000..102f4ae1c --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/inspect_pages.py @@ -0,0 +1,108 @@ +"""Generic VLM inspection tool for selected PDF pages.""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from typing import Any, cast + +from app.services.document_agent.manifest import ToolContext, ToolResult +from app.services.document_agent.registry import has_page_features, register_tool +from app.services.document_agent.visual import render_pages +from shared.utils.token_estimate import estimate_tokens + + +@register_tool( + name="inspect.pages", + description="Render arbitrary PDF pages and ask the VLM a custom profiling question.", + parameters={ + "type": "object", + "properties": { + "pages": {"type": "array", "items": {"type": "integer"}}, + "question": {"type": "string"}, + }, + "required": ["pages", "question"], + }, + preconditions=(has_page_features,), +) +def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + pages = sorted( + { + int(page) + for page in (args.get("pages") or []) + if 1 <= int(page) <= ctx.blackboard.page_count + } + )[:10] + if not pages: + return ToolResult( + status="error", + error="inspect.pages requires at least one valid page", + latency_ms=int((time.monotonic() - start) * 1000), + ) + question = str(args.get("question") or "Describe the document structure visible on these pages.") + pngs = render_pages(ctx, pages, folder_name="inspect_pages", prefix="inspect") + model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") + prompt = ( + "You are inspecting PDF page screenshots for a document profiling agent. " + "Answer strict JSON with keys: observations, implications, recommended_next_action. " + "observations must be an array of {page, summary, visual_kind}. " + f"Question: {question}" + ) + est = estimate_tokens(prompt) + len(pngs) * 800 + if not model: + payload = {"pages": pages, "pngs": pngs, "note": "No VLM model configured."} + ctx.blackboard.global_signals.setdefault("visual_inspections", []).append(payload) + return ToolResult( + status="ok", + payload=payload, + latency_ms=int((time.monotonic() - start) * 1000), + warnings=["No VLM model configured; returned rendered page paths only."], + ) + if not ctx.budget.try_reserve("visual", est): + return ToolResult( + status="error", + error="insufficient visual budget", + latency_ms=int((time.monotonic() - start) * 1000), + ) + + content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + for item in pngs: + with open(str(item["png_path"]), "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + content_parts.append({"type": "text", "text": f"\n--- Page {item['page']} ---"}) + content_parts.append( + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_b64}"}} + ) + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=model, + temperature=0.0, + max_tokens=1200, + response_format={"type": "json_object"}, + ) + ctx.budget.commit("visual", actual=usage.get("total_tokens", est), est=est) + try: + payload = json.loads(raw) + except json.JSONDecodeError: + payload = {"raw": raw} + if isinstance(payload, dict): + payload.setdefault("pages", pages) + else: + payload = {"result": payload, "pages": pages} + ctx.blackboard.global_signals.setdefault("visual_inspections", []).append(payload) + return ToolResult( + status="ok", + payload=payload, + latency_ms=int((time.monotonic() - start) * 1000), + tokens_used=usage.get("total_tokens", 0), + ) + except Exception: + ctx.budget.refund("visual", est=est) + raise diff --git a/apps/worker/app/services/document_agent/tools/match_h1_pages.py b/apps/worker/app/services/document_agent/tools/match_h1_pages.py index ac97c0a00..a60ad2695 100644 --- a/apps/worker/app/services/document_agent/tools/match_h1_pages.py +++ b/apps/worker/app/services/document_agent/tools/match_h1_pages.py @@ -14,8 +14,7 @@ ToolResult, ) from app.services.document_agent.pdf_text import read_page_texts -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState +from app.services.document_agent.registry import has_toc_result, register_tool from loguru import logger @@ -76,7 +75,7 @@ def _extract_level1_titles(toc_hierarchies: list[dict[str, Any]]) -> list[str]: "Match TOC level-1 headings to body pages using PyMuPDF substring search. " "Produces H1Candidate list for downstream shard planning." ), - allowed_states={DocumentAgentState.CLASSIFIED}, + preconditions=(has_toc_result,), ) def match_h1_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() diff --git a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py index b48900a02..7ae7f0195 100644 --- a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py @@ -9,8 +9,6 @@ from typing import Any from app.services.document_agent.manifest import PageAnatomyMap, ToolContext, ToolResult -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState def _artifact_dir(ctx: ToolContext) -> Path: @@ -24,7 +22,6 @@ def build_anatomy_map(ctx: ToolContext) -> PageAnatomyMap: if not ( ctx.blackboard.toc_result and ctx.blackboard.h1_result - and ctx.blackboard.hierarchy_assist and ctx.blackboard.shard_plan ): raise ValueError("cannot build anatomy map from incomplete blackboard") @@ -36,25 +33,17 @@ def build_anatomy_map(ctx: ToolContext) -> PageAnatomyMap: page_labels=ctx.blackboard.page_labels, toc_result=ctx.blackboard.toc_result, h1_result=ctx.blackboard.h1_result, - hierarchy_assist=ctx.blackboard.hierarchy_assist, shard_plan=ctx.blackboard.shard_plan, - boundary_candidates=ctx.blackboard.boundary_candidates, + document_profile=ctx.blackboard.document_profile, toc_hierarchies=ctx.blackboard.toc_hierarchies, - page_processing_plan=None, global_signals=ctx.blackboard.global_signals, trace_summary={ "budget": ctx.budget.snapshot(), - "state_trace": list(ctx.blackboard.state_trace), "validation": ctx.blackboard.validation_report, }, ) -@register_tool( - name="persist.anatomy_map", - description="Persist anatomy_map.json and buffer SQL trace payloads.", - allowed_states={DocumentAgentState.VALIDATED}, -) def persist_anatomy_map(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() anatomy = build_anatomy_map(ctx) 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 6eac30a0d..678557fad 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 @@ -8,8 +8,6 @@ from app.services.document_agent.manifest import PageFeature, ToolContext, ToolResult from app.services.document_agent.pdf_text import top_lines -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( run_in_child_process, worker, @@ -104,11 +102,6 @@ def _probe_worker(queue, pdf_path: str) -> None: queue.put({"ok": True, "page_count": page_count, "features": features}) -@register_tool( - name="probe.page_features", - description="Probe every PDF page for structural signals without parsing content semantically.", - allowed_states={DocumentAgentState.INIT}, -) def probe_page_features(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() try: diff --git a/apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py b/apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py deleted file mode 100644 index 50cd350bf..000000000 --- a/apps/worker/app/services/document_agent/tools/propose_hierarchy_assist.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Build hierarchy hints from page anatomy signals.""" - -from __future__ import annotations - -import time -from typing import Any - -from app.services.document_agent.manifest import ( - BoundaryHint, - H1Candidate, - HierarchyAssistPlan, - PageLabel, - ToolContext, - ToolResult, -) -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState -from app.services.document_agent.validators import repair_hierarchy_assist - - -def _rule_based_plan(labels: list[PageLabel], h1_candidates: list[H1Candidate]) -> HierarchyAssistPlan: - exclude = [ - label.page - for label in labels - if label.kind in {"toc", "blank", "separator", "landscape"} - ] - suppress = [ - label.page - for label in labels - if label.kind in {"table_heavy", "single_image", "scan_like", "image_heavy"} - ] - boundary_hints = [ - BoundaryHint( - page=candidate.page, - anchor_type="h1_boundary", - confidence=candidate.confidence, - evidence=candidate.evidence | {"title": candidate.title}, - ) - for candidate in h1_candidates - ] - scan_like_count = sum(1 for label in labels if label.kind in {"scan_like", "single_image"}) - recommendation = "aggressive" if scan_like_count > max(len(labels) // 3, 0) else "normal" - if h1_candidates and scan_like_count == 0: - recommendation = "normal" - if not h1_candidates and scan_like_count == 0: - recommendation = "aggressive" - return HierarchyAssistPlan( - exclude_pages_from_title_candidates=sorted(set(exclude)), - prefer_h1_start_pages=sorted(h1_candidates, key=lambda item: item.page), - suppress_title_pages=sorted(set(suppress)), - section_boundary_hints=boundary_hints, - smart_parse_recommendation=recommendation, # type: ignore[arg-type] - rationale="Derived from page labels and H1 boundary evidence.", - ) - - -@register_tool( - name="propose.hierarchy_assist", - description="Produce hierarchy hints used by section skeleton extraction and title parsing.", - allowed_states={DocumentAgentState.H1_FOUND}, -) -def propose_hierarchy_assist(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: - start = time.monotonic() - h1_result = ctx.blackboard.h1_result - labels = ctx.blackboard.page_labels - fallback = _rule_based_plan(labels, h1_result.h1_candidates if h1_result else []) - toc_pages = set(ctx.blackboard.toc_result.toc_pages if ctx.blackboard.toc_result else []) - plan = fallback - - repaired = repair_hierarchy_assist( - plan, - page_count=ctx.blackboard.page_count, - toc_pages=toc_pages, - ) - ctx.blackboard.hierarchy_assist = repaired - return ToolResult( - status="ok", - payload={ - "exclude_pages": repaired.exclude_pages_from_title_candidates, - "suppress_pages": repaired.suppress_title_pages, - "h1_hint_count": len(repaired.prefer_h1_start_pages), - }, - latency_ms=int((time.monotonic() - start) * 1000), - input_summary={ - "page_kind_counts": ctx.blackboard.global_signals.get("page_kind_counts", {}), - "boundary_candidate_counts": ctx.blackboard.global_signals.get( - "boundary_candidate_counts", {} - ), - }, - output_summary={ - "exclude_pages": repaired.exclude_pages_from_title_candidates[:50], - "suppress_pages_count": len(repaired.suppress_title_pages), - "h1_hint_count": len(repaired.prefer_h1_start_pages), - "smart_parse_recommendation": repaired.smart_parse_recommendation, - "rationale": repaired.rationale, - }, - ) 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 91836936c..dc5bfda52 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -1,4 +1,4 @@ -"""LLM-guided long-PDF shard planning from candidate split pages.""" +"""LLM-guided long-PDF shard planning from document profile evidence.""" from __future__ import annotations @@ -9,14 +9,12 @@ from typing import Any from app.services.document_agent.manifest import ( - BoundaryCandidate, Shard, ShardPlan, ToolContext, ToolResult, ) -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState +from app.services.document_agent.registry import has_doc_stats, has_h1_result, has_toc_result, register_tool from app.services.document_agent.validators import single_shard_plan, validate_shard_plan from shared.utils.token_estimate import estimate_tokens @@ -70,73 +68,52 @@ def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> return shards -def _compact_candidate(candidate: BoundaryCandidate, page_count: int) -> dict[str, Any]: - evidence = candidate.evidence or {} - return { - "page": candidate.page, - "kind": candidate.kind, - "priority": candidate.priority, - "confidence": candidate.confidence, - "position_ratio": round(candidate.page / max(page_count, 1), 4), - "raw_text_length": evidence.get("raw_text_length"), - "image_coverage": evidence.get("image_coverage"), - "table_count": evidence.get("table_count"), - "drawings_count": evidence.get("drawings_count"), - "text_preview": evidence.get("text_preview", [])[:3], - "title": evidence.get("title"), - } - - def _build_prompt( *, page_count: int, min_pages: int, max_pages: int, - candidates: list[BoundaryCandidate], + doc_stats: dict[str, Any], page_kind_counts: dict[str, int], + toc_pages: list[int], + h1_pages: list[dict[str, Any]], + 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, - "candidate_priority": { - "h1": "highest semantic priority, but still decide using document size and spacing", - "toc": "high priority marker, usually not a cut by itself unless it indicates nearby structure", - "separator": "explicit sparse separator page", - "blank": "sparse/blank structural gap candidate", - "sparse": "low-density candidate, useful when no stronger signal exists", - }, - "candidates": [ - _compact_candidate(candidate, page_count) for candidate in candidates - ], + "doc_stats": doc_stats, + "toc_pages": toc_pages, + "h1_pages": h1_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 only the provided candidate pages and document-scale " - "features.\n" + "and where to split it using document-scale features, TOC/H1 evidence, and " + "recent agent observations.\n" "Rules:\n" "- Return strict JSON only.\n" - "- Do not invent pages. Every cut_after_page must be one of the candidate pages, " - "or candidate_page - 1 when the candidate is a semantic start page such as h1.\n" - "- H1 candidates have the highest semantic priority, but do not blindly split on " - "every H1. Consider total page_count, candidate spacing, min/max shard sizes, and " - "whether splitting would over-fragment the document.\n" - "- Blank and sparse pages are valid split candidates because they often mark section " - "gaps, especially when TOC/H1 evidence is weak or absent.\n" + "- Prefer H1 start pages as semantic boundaries, cutting at page-1 when possible.\n" + "- Do not blindly split on every H1. Consider total page_count, spacing, min/max " + "shard sizes, and over-fragmentation.\n" "- Prefer fewer, semantically coherent shards over many tiny shards.\n" - "- Do not use domain-specific hardcoded labels or examples; decide from the supplied " - "features and positions only. Do not quote business/category words from text_preview " - "in rationale; refer to them generically as sparse separator text.\n" - "- Every resulting shard length must be <= max_pages_per_shard unless enabled=false. " - "Check each segment length exactly before returning.\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\" | \"separator\" | \"forced_max_size\", " + "\"blank_separator\" | \"forced_max_size\", " "\"confidence\": number, \"rationale\": string}\n" " ],\n" ' "reason": "llm_boundary_decision" | "not_needed" | "too_large",\n' @@ -155,21 +132,39 @@ def _sanitize_rationale(text: str) -> str: return sanitized -def _validate_cut_lengths(cuts: list[tuple[int, str, str, float]], page_count: int, max_pages: int) -> None: +def _validate_cut_lengths( + cuts: list[tuple[int, str, str, float]], + page_count: int, + min_pages: int, + max_pages: int, +) -> None: previous = 0 for cut_page, *_ in cuts: + if cut_page - previous < min_pages: + raise ValueError( + f"LLM cut plan creates shard length {cut_page - previous} < min_pages={min_pages}" + ) if cut_page - previous > max_pages: raise ValueError( f"LLM cut plan creates shard length {cut_page - previous} > max_pages={max_pages}" ) previous = cut_page + if page_count - previous < min_pages and cuts: + raise ValueError( + f"LLM cut plan creates final shard length {page_count - previous} < min_pages={min_pages}" + ) if page_count - previous > max_pages: raise ValueError( f"LLM cut plan creates final shard length {page_count - previous} > max_pages={max_pages}" ) -def _parse_llm_plan(raw: str, page_count: int, max_pages: int) -> tuple[bool, list[tuple[int, str, str, float]], str, str]: +def _parse_llm_plan( + raw: str, + page_count: int, + min_pages: int, + max_pages: int, +) -> tuple[bool, list[tuple[int, str, str, float]], str, str]: data = json.loads(raw) if not isinstance(data, dict): raise ValueError("LLM shard plan is not an object") @@ -180,17 +175,20 @@ def _parse_llm_plan(raw: str, page_count: int, max_pages: int) -> tuple[bool, li for item in data.get("cuts") or []: if not isinstance(item, dict): continue - cut_page = int(item.get("cut_after_page")) + raw_cut_page = item.get("cut_after_page") + if raw_cut_page is None: + continue + cut_page = int(raw_cut_page) if not 1 <= cut_page < page_count: continue - anchor_type = str(item.get("anchor_type") or "separator") - if anchor_type not in {"h1_boundary", "blank_separator", "separator", "forced_max_size"}: - anchor_type = "separator" + anchor_type = str(item.get("anchor_type") or "forced_max_size") + if anchor_type not in {"h1_boundary", "blank_separator", "forced_max_size"}: + anchor_type = "forced_max_size" confidence = float(item.get("confidence") or 0.5) cuts.append((cut_page, anchor_type, _sanitize_rationale(str(item.get("rationale") or rationale)), confidence)) cuts = sorted({cut[0]: cut for cut in cuts}.values(), key=lambda cut: cut[0]) if enabled: - _validate_cut_lengths(cuts, page_count, max_pages) + _validate_cut_lengths(cuts, page_count, min_pages, max_pages) return enabled, cuts, reason, rationale @@ -198,20 +196,19 @@ def _deterministic_guardrail_plan( *, page_count: int, max_pages: int, - candidates: list[BoundaryCandidate], + h1_pages: list[int], ) -> tuple[list[tuple[int, str, str, float]], str]: cuts: list[tuple[int, str, str, float]] = [] previous = 0 while page_count - previous > max_pages: target = previous + max_pages eligible = [ - candidate for candidate in candidates if previous < candidate.page <= target + page for page in h1_pages if previous + 1 < page <= target ] if eligible: - chosen = max(eligible, key=lambda item: (item.priority, item.page)) - cut_page = chosen.page - 1 if chosen.kind == "h1" and chosen.page > previous + 1 else chosen.page - anchor_type = "h1_boundary" if chosen.kind == "h1" else "blank_separator" - cuts.append((cut_page, anchor_type, f"guardrail candidate {chosen.kind} at page {chosen.page}", 0.35)) + chosen = max(eligible) + cut_page = chosen - 1 + cuts.append((cut_page, "h1_boundary", f"guardrail H1 start page {chosen}", 0.35)) previous = cut_page else: cuts.append((target, "forced_max_size", "guardrail max shard size", 0.25)) @@ -221,8 +218,8 @@ def _deterministic_guardrail_plan( @register_tool( name="propose.shard_plan", - description="Ask the LLM to decide whether and where to split using candidate boundary pages.", - allowed_states={DocumentAgentState.H1_FOUND}, + description="Ask the LLM to decide whether and where to split using profile, TOC, and H1 evidence.", + preconditions=(has_doc_stats, has_toc_result, has_h1_result), ) def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() @@ -237,14 +234,24 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) - candidates = list(ctx.blackboard.boundary_candidates) + h1_candidates = ( + ctx.blackboard.h1_result.h1_candidates if ctx.blackboard.h1_result else [] + ) + h1_pages = [{"title": item.title, "page": item.page} for item in h1_candidates] model = ctx.settings.get("model") prompt = _build_prompt( page_count=page_count, min_pages=min_pages, max_pages=max_pages, - candidates=candidates, + doc_stats=ctx.blackboard.doc_stats, page_kind_counts=ctx.blackboard.global_signals.get("page_kind_counts", {}), + toc_pages=ctx.blackboard.toc_result.toc_pages if ctx.blackboard.toc_result else [], + h1_pages=h1_pages, + profile=ctx.blackboard.document_profile.to_dict() + if ctx.blackboard.document_profile + else None, + visual_evidence=ctx.blackboard.global_signals.get("visual_inspections", []), + grep_history=ctx.blackboard.global_signals.get("grep_history", []), ) prompt_tokens_est = estimate_tokens(prompt) warnings: list[str] = [] @@ -265,30 +272,46 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: response_format={"type": "json_object"}, ) ctx.budget.commit("plan", actual=usage.get("total_tokens", prompt_tokens_est), est=prompt_tokens_est) - enabled, cuts, reason, rationale = _parse_llm_plan(raw_response, page_count, max_pages) + enabled, cuts, reason, rationale = _parse_llm_plan(raw_response, page_count, min_pages, max_pages) if not enabled: cuts = [] reason = "not_needed" except Exception as exc: ctx.budget.refund("plan", est=prompt_tokens_est) - warnings.append(f"LLM shard decision rejected, using guardrail plan: {exc}") + warnings.append(f"LLM shard decision failed; using guardrail plan: {exc}") + ctx.blackboard.global_signals.setdefault("degraded_reasons", []).append( + "shard_plan: llm_parse_failed" + ) cuts, reason = _deterministic_guardrail_plan( page_count=page_count, max_pages=max_pages, - candidates=candidates, + h1_pages=[item["page"] for item in h1_pages], ) - rationale = "Guardrail plan after LLM shard decision failure." + rationale = "Guardrail plan after malformed LLM shard decision." else: if not model: warnings.append("No model configured for shard decision; using guardrail plan.") + ctx.blackboard.global_signals.setdefault("degraded_reasons", []).append( + "shard_plan: no model" + ) + cuts, reason = _deterministic_guardrail_plan( + page_count=page_count, + max_pages=max_pages, + h1_pages=[item["page"] for item in h1_pages], + ) + rationale = "Guardrail plan without configured shard model." else: - warnings.append("Insufficient plan budget for shard decision; using guardrail plan.") - cuts, reason = _deterministic_guardrail_plan( - page_count=page_count, - max_pages=max_pages, - candidates=candidates, - ) - rationale = "Guardrail plan without LLM decision." + return ToolResult( + status="error", + error="Insufficient plan budget for shard decision.", + latency_ms=int((time.monotonic() - start) * 1000), + warnings=warnings, + debug={ + "prompt_excerpt": prompt[:4000], + "raw_response_excerpt": raw_response[:4000], + "llm_attempted": llm_attempted, + }, + ) shards = _cuts_to_shards(cuts, page_count) enabled = len(shards) > 1 @@ -318,8 +341,7 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: tokens_used=ctx.budget.snapshot()["plan"]["used"], input_summary={ "page_count": page_count, - "candidate_count": len(candidates), - "candidate_counts": ctx.blackboard.global_signals.get("boundary_candidate_counts", {}), + "h1_count": len(h1_pages), "model": model, }, output_summary={ diff --git a/apps/worker/app/services/document_agent/tools/test_vlm_toc_extract.py b/apps/worker/app/services/document_agent/tools/test_vlm_toc_extract.py deleted file mode 100644 index 100d1dfd3..000000000 --- a/apps/worker/app/services/document_agent/tools/test_vlm_toc_extract.py +++ /dev/null @@ -1,817 +0,0 @@ -#!/usr/bin/env python3 -"""Standalone test: VLM-direct TOC extraction from page PNGs. - -Usage: - cd knowhereapi-main - python apps/worker/app/services/document_agent/tools/test_vlm_toc_extract.py - -Requires: - - ALI_API_KEYS env var (or in apps/worker/.env) - - pymupdf, openai installed - - Test PDFs accessible on disk -""" - -from __future__ import annotations - -import base64 -import json -import os -import sys -import time -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any - -# --------------------------------------------------------------------------- -# VLM prompt — no hardcoded examples, no overfitting -# --------------------------------------------------------------------------- - -VLM_TOC_EXTRACT_PROMPT = """\ -You are analyzing a Table of Contents (TOC) page from a document. - -Your task is to extract every TOC entry visible on this page. - -Each entry consists of three parts: -1. **title** — the section or chapter name, copied verbatim from the page. - - EXCLUDE any trailing dots (…·.), dashes (—-), or leader lines that connect the title to its page number. - - If one entry's title wraps across multiple printed lines, combine them into a single string. - - Include any numbering prefix that is part of the title text (e.g. "1.", "第三章", "(二)"). -2. **page_number** — the page reference at the right side of the entry. - - Use an integer when the reference is a plain number (e.g. 26). - - Use a string when the reference is non-numeric (e.g. "iv", "F-1", "A-3"). - - Use null when no page reference is visible for that entry. -3. **level** — the hierarchy depth of the entry, determined by visual formatting cues: - - level 1: top-level entries — no indentation, or the largest / boldest text. - - level 2: sub-entries — indented under a level-1 entry, or in a noticeably smaller font. - - level 3+: deeper indentation, if present. - - Category headers or group labels that are visually distinct (centered, larger font, different style) and do NOT have a page number should be treated as level 1. - -Additional rules: -- Extract ALL entries, even if the page only shows a partial continuation of the TOC. -- Do NOT include the TOC page's own heading (e.g. a "TABLE OF CONTENTS" or "目 录" title at the top) as an entry. -- Do NOT include column headers (e.g. a standalone "Page" label) as entries. -- Preserve the original language and wording of each title exactly. - -Return strict JSON with no markdown fences: -{"entries": [{"title": "...", "page_number": ..., "level": ...}, ...]} -""" - -# Mirror the production boundary step from extract_toc_with_boundaries.py -BOUNDARY_STEP_PAGES = 5 - -VLM_TOC_CONTINUATION_CONTEXT = """\ - ---- IMPORTANT: Continuation Context --- -This is a CONTINUATION page of a multi-page Table of Contents. -The previous page(s) already extracted the following entries: - -{previous_summary} - -The LAST active category/section before this page was: - Level {last_l1_level}: "{last_l1_title}" - -Entries on THIS page that visually continue as sub-items under that -category (same indentation, same numbering sequence) must keep their -correct subordinate level — do NOT promote them to level 1 just because -the parent heading is not visible on this page. -""" - - -# --------------------------------------------------------------------------- -# Test case definitions -# --------------------------------------------------------------------------- - - -@dataclass -class TocTestCase: - """One test document for VLM TOC extraction.""" - - name: str - pdf_path: str | None # None means PNGs are pre-rendered - toc_page_nums: list[int] # 1-based page numbers to extract - output_dir: str - pre_rendered_pngs: list[str] = field(default_factory=list) - expected_entry_count_range: tuple[int, int] = (1, 999) - description: str = "" - toc_output_path: str | None = None # Override toc_hierarchies.json output path - - -def _build_test_cases() -> list[TocTestCase]: - debug_root = os.path.expanduser("~/.knowhere/_debug_profile") - cases: list[TocTestCase] = [] - - # Case 1: SpaceX S-1 — single TOC page, English - spacex_pdf = "/Users/wuchengke/Downloads/spacex-s1.pdf" - spacex_out = os.path.join(debug_root, "spacex-s1", "toc_pages") - if os.path.exists(spacex_pdf): - cases.append( - TocTestCase( - name="SpaceX S-1 (English)", - pdf_path=spacex_pdf, - toc_page_nums=[17], # VLM-confirmed TOC start, boundary says 17 only - output_dir=spacex_out, - expected_entry_count_range=(20, 30), - description="Single flat TOC page, all level-1 entries, page nums on right", - ) - ) - else: - # Fall back to pre-rendered PNG - png17 = os.path.join(spacex_out, "toc_anchor_page_17.png") - if os.path.exists(png17): - cases.append( - TocTestCase( - name="SpaceX S-1 (English, pre-rendered)", - pdf_path=None, - toc_page_nums=[17], - output_dir=spacex_out, - pre_rendered_pngs=[png17], - expected_entry_count_range=(20, 30), - description="Single flat TOC page from pre-rendered PNG", - ) - ) - - # Case 2: Chinese corporate regulations — multi-page TOC with categories - cn_pdf = "/Users/wuchengke/Desktop/temp/test_docs/SJSYJ-SC-2024 企业制度汇编（上册）.pdf" - cn_out = os.path.join(debug_root, "chinese_corp", "toc_pages") - if os.path.exists(cn_pdf): - cases.append( - TocTestCase( - name="企业制度汇编 (Chinese, multi-page TOC)", - pdf_path=cn_pdf, - toc_page_nums=[5, 6], # TOC spans pages 5-6 - output_dir=cn_out, - expected_entry_count_range=(25, 45), - description=( - "Multi-page TOC with category headers (经营类/生产类/安全类/...), " - "subcategory codes (SJSYJ-SC101-2024), numbered entries, " - "multi-line wrapping, and Chinese dot leaders" - ), - ) - ) - - return cases - - -# --------------------------------------------------------------------------- -# PNG rendering -# --------------------------------------------------------------------------- - - -def _render_page_png(pdf_path: str, page_num: int, output_dir: str, dpi: int = 144) -> str: - """Render a single page to PNG. Returns the PNG path.""" - import pymupdf # type: ignore[import] - - os.makedirs(output_dir, exist_ok=True) - png_path = os.path.join(output_dir, f"toc_page_{page_num}.png") - - doc = pymupdf.open(pdf_path) - try: - idx = page_num - 1 - if 0 <= idx < doc.page_count: - page = doc[idx] - mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) - pix = page.get_pixmap(matrix=mat) - pix.save(png_path) - else: - raise ValueError(f"Page {page_num} out of range (total: {doc.page_count})") - finally: - doc.close() - - return png_path - - -# --------------------------------------------------------------------------- -# VLM call -# --------------------------------------------------------------------------- - - -def _build_openai_client(model: str): - """Build an OpenAI SDK client for the given model, without requiring AppConfig. - - Mirrors the routing logic in shared.services.ai.openai_compatible_client_sync - but reads env vars directly so the test can run standalone. - """ - from openai import OpenAI - - model_lower = model.lower() - - if "qwen" in model_lower: - # Aliyun DashScope - api_key = os.environ.get("ALI_API_KEYS", "").strip() - # ALI_API_KEYS can be JSON array or comma-separated; grab the first one - if api_key.startswith("["): - import re - keys = re.findall(r'"([^"]+)"', api_key) - api_key = keys[0] if keys else "" - elif "," in api_key: - api_key = api_key.split(",")[0].strip() - # Handle token_id=api_key format - if "=" in api_key and not api_key.startswith("sk-"): - api_key = api_key.split("=", 1)[1] - base_url = os.environ.get( - "ALI_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1" - ) - elif "deepseek" in model_lower: - api_key = os.environ.get("DS_KEY", "") - base_url = os.environ.get("DS_URL", "https://api.deepseek.com/v1") - else: - api_key = os.environ.get("OPENAI_API_KEY", "") - base_url = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") - - # Strip /chat/completions suffix if present - if base_url.rstrip("/").endswith("/chat/completions"): - base_url = base_url.rstrip("/").removesuffix("/chat/completions") - - if not api_key: - raise RuntimeError( - f"No API key found for model {model!r}. " - "Set ALI_API_KEYS, DS_KEY, or OPENAI_API_KEY env var." - ) - - return OpenAI(api_key=api_key, base_url=base_url, timeout=120, max_retries=2) - - -def _build_continuation_context(previous_entries: list[dict[str, Any]]) -> str: - """Build a concise context string from previously extracted entries. - - Returns the continuation context to append to the VLM prompt, - or empty string if no previous entries exist. - """ - if not previous_entries: - return "" - - # Build a compact summary: last N entries with their levels - # Show the last 8 entries to give enough context - tail = previous_entries[-8:] - summary_lines = [] - for e in tail: - lvl = e.get("level", "?") - title = e.get("title", "?") - pn = e.get("page_number") - pn_str = f" → p.{pn}" if pn is not None else "" - summary_lines.append(f" L{lvl}: {title}{pn_str}") - - if len(previous_entries) > 8: - summary_lines.insert(0, f" ... ({len(previous_entries) - 8} earlier entries omitted)") - - previous_summary = "\n".join(summary_lines) - - # Find the last L1 entry (the active parent category) - last_l1 = None - for e in reversed(previous_entries): - if e.get("level") == 1: - last_l1 = e - break - - if last_l1 is None: - # No L1 found — still provide the summary but skip the "last active" part - return f"\n\n--- IMPORTANT: Continuation Context ---\nThis is a CONTINUATION page. Previous entries:\n{previous_summary}\n" - - return VLM_TOC_CONTINUATION_CONTEXT.format( - previous_summary=previous_summary, - last_l1_level=last_l1.get("level", 1), - last_l1_title=last_l1.get("title", "?"), - ) - - -def _vlm_extract_toc_entries( - png_path: str, - page_num: int, - model: str, - previous_entries: list[dict[str, Any]] | None = None, -) -> tuple[list[dict[str, Any]], dict[str, Any]]: - """Call VLM to extract TOC entries from a single page PNG. - - Args: - png_path: path to the page PNG image. - page_num: 1-based page number. - model: VLM model name. - previous_entries: entries already extracted from earlier TOC pages. - Used to build continuation context so the model can correctly - assign hierarchy levels on continuation pages. - - Returns: - (entries, meta) where meta has token usage and timing info. - """ - client = _build_openai_client(model) - - with open(png_path, "rb") as f: - img_b64 = base64.b64encode(f.read()).decode() - - # Build prompt: base + optional continuation context - prompt_text = VLM_TOC_EXTRACT_PROMPT - continuation_ctx = _build_continuation_context(previous_entries or []) - if continuation_ctx: - prompt_text += continuation_ctx - - content_parts: list[dict[str, Any]] = [ - {"type": "text", "text": prompt_text}, - { - "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{img_b64}"}, - }, - ] - - messages = [{"role": "user", "content": content_parts}] - - t0 = time.monotonic() - response = client.chat.completions.create( - model=model, - messages=messages, - temperature=0.1, - max_tokens=4096, - response_format={"type": "json_object"}, - extra_body={"enable_thinking": False}, - ) - elapsed_ms = int((time.monotonic() - t0) * 1000) - - raw = response.choices[0].message.content or "" - usage_obj = response.usage - usage = { - "prompt_tokens": getattr(usage_obj, "prompt_tokens", 0) or 0, - "completion_tokens": getattr(usage_obj, "completion_tokens", 0) or 0, - "total_tokens": getattr(usage_obj, "total_tokens", 0) or 0, - } - - meta = { - "page": page_num, - "model": model, - "elapsed_ms": elapsed_ms, - "usage": usage, - "raw_response_length": len(raw), - "has_continuation_context": bool(continuation_ctx), - } - - # Parse JSON - data = json.loads(raw) - if isinstance(data, dict): - entries = data.get("entries", []) - elif isinstance(data, list): - entries = data - else: - entries = [] - - return entries, meta - - -# --------------------------------------------------------------------------- -# VLM entries → toc_hierarchies.json conversion -# --------------------------------------------------------------------------- - - -def _build_toc_tree(entries: list[dict[str, Any]]) -> dict[str, Any]: - """Build a nested dict tree from VLM entries, same algorithm as - toc_hierarchy.build_tree_tocs() but self-contained for standalone use. - - Args: - entries: list of {"title": str, "level": int, "page_number": ...} - - Returns: - Nested dict: {heading: {child_heading: {...}, ...}, ...} - """ - if not entries: - return {} - - root: dict[str, Any] = {} - stack: list[tuple[dict[str, Any], int]] = [(root, 0)] - - positive_levels = [e["level"] for e in entries if isinstance(e.get("level"), int) and e["level"] > 0] - level_for_minus_one = max(positive_levels) + 1 if positive_levels else 1 - - for entry in entries: - heading = entry.get("title", "").strip() - if not heading: - continue - - original_level = entry.get("level", 1) - normalized_level = level_for_minus_one if original_level == -1 else original_level - - while len(stack) > 1 and stack[-1][1] >= normalized_level: - stack.pop() - - parent_dict = stack[-1][0] - parent_dict[heading] = {} - stack.append((parent_dict[heading], normalized_level)) - - return root - - -def _build_toc_with_level_md(entries: list[dict[str, Any]]) -> str: - """Build a markdown table string compatible with existing toc_with_level format.""" - if not entries: - return "" - - lines = ["| id | heading | level |"] - lines.append("|----|---------|-------|") - for i, e in enumerate(entries, 1): - heading = e.get("title", "").strip().replace("|", "\\|") - level = e.get("level", 1) - # Pad heading for readability - lines.append(f"| {i:<2} | {heading:<60} | {level:<5} |") - return "\n".join(lines) - - -def vlm_entries_to_toc_hierarchies( - all_entries: list[dict[str, Any]], - toc_page_nums: list[int], - scan_end_page: int | None = None, - page_count: int | None = None, -) -> list[dict[str, Any]]: - """Convert VLM extraction results to the standard toc_hierarchies.json schema. - - Args: - all_entries: merged VLM entries from all TOC pages. - Each entry: {"title": str, "page_number": int|str|None, "level": int} - toc_page_nums: 1-based page numbers of the TOC pages. - scan_end_page: 1-based page number of the lookahead boundary check. - In production, this is `anchor_page + BOUNDARY_STEP_PAGES - 1`. - If the boundary check page was NOT TOC, it still counts as part - of the scan range (it was the page we "looked at" to decide). - Defaults to `max(toc_page_nums) + BOUNDARY_STEP_PAGES - 1`. - page_count: total pages in the document (used to clamp scan_end_page). - - Returns: - List with one dict per TOC region (usually 1), matching the schema: - { - "toc_range": [start_page, end_page], - "toc_range_unit": "page", - "scan_range": [start_page, scan_end_page], - "source": "vlm", - "toc_with_level": [structured list], - "toc_with_level_md": "| id | heading | level |\n...", - "toc_tree": {nested dict} - } - """ - if not all_entries or not toc_page_nums: - return [] - - # Build structured toc_with_level list - toc_with_level: list[dict[str, Any]] = [] - for i, entry in enumerate(all_entries, 1): - item: dict[str, Any] = { - "id": i, - "heading": entry.get("title", "").strip(), - "level": entry.get("level", 1), - } - # Include page_number from VLM extraction - pn = entry.get("page_number") - item["page_number"] = pn - toc_with_level.append(item) - - # Build toc_tree from clean entries - toc_tree = _build_toc_tree(all_entries) - - # Build markdown table for backward compatibility - toc_with_level_md = _build_toc_with_level_md(all_entries) - - # Determine page ranges - start_page = min(toc_page_nums) - end_page = max(toc_page_nums) - - # scan_range: the full lookahead window used during boundary detection. - # In production: anchor + BOUNDARY_STEP_PAGES - 1, clamped to page_count. - # The anchor is start_page (not end_page), matching _detect_toc_range_for_anchor(). - if scan_end_page is None: - scan_end_page = start_page + BOUNDARY_STEP_PAGES - 1 - if page_count is not None: - scan_end_page = min(scan_end_page, page_count) - - return [ - { - "toc_range": [start_page, end_page], - "toc_range_unit": "page", - "scan_range": [start_page, scan_end_page], - "source": "vlm", - "toc_with_level": toc_with_level, - "toc_with_level_md": toc_with_level_md, - "toc_tree": toc_tree, - } - ] - - -# --------------------------------------------------------------------------- -# Comparison with existing toc_hierarchies.json -# --------------------------------------------------------------------------- - - -def _compare_with_existing( - vlm_entries: list[dict[str, Any]], - test_case: TocTestCase, -) -> dict[str, Any]: - """Compare VLM results with existing toc_hierarchies.json if available.""" - comparison: dict[str, Any] = {"available": False} - - # Try to find existing toc_hierarchies.json - parent_dir = str(Path(test_case.output_dir).parent) - existing_path = os.path.join(parent_dir, "toc_hierarchies.json") - if not os.path.exists(existing_path): - return comparison - - with open(existing_path, "r", encoding="utf-8") as f: - existing = json.load(f) - - comparison["available"] = True - comparison["existing_path"] = existing_path - - # Extract titles from existing toc_tree - existing_titles: list[str] = [] - if isinstance(existing, list) and existing: - tree = existing[0].get("toc_tree", {}) - for key, sub in tree.items(): - existing_titles.append(key) - if isinstance(sub, dict): - for subkey in sub: - existing_titles.append(subkey) - - vlm_titles = [e.get("title", "") for e in vlm_entries] - - # Quality checks - issues: list[str] = [] - - # Check 1: titles with residual page numbers (MinerU artifact) - import re - residual_pattern = re.compile(r"\s+\d+\s*$") - existing_with_residual = [t for t in existing_titles if residual_pattern.search(t)] - vlm_with_residual = [t for t in vlm_titles if residual_pattern.search(t)] - - comparison["existing_residual_page_nums"] = existing_with_residual - comparison["vlm_residual_page_nums"] = vlm_with_residual - - if existing_with_residual and not vlm_with_residual: - issues.append( - f"✅ VLM fixed {len(existing_with_residual)} residual page numbers " - f"in titles (e.g. '{existing_with_residual[0]}')" - ) - elif vlm_with_residual: - issues.append( - f"⚠️ VLM still has {len(vlm_with_residual)} titles with residual " - f"page numbers: {vlm_with_residual[:3]}" - ) - - # Check 2: broken multi-line titles - # In existing data: "MANAGEMENT'S DISCUSSION..." and "OF OPERATIONS 74" are separate - broken_line_keywords = ["OF OPERATIONS", "CLASS A COMMON STOCK"] - existing_broken = [ - t for t in existing_titles - if any(t.strip().startswith(kw) for kw in broken_line_keywords) - ] - vlm_broken = [ - t for t in vlm_titles - if any(t.strip().startswith(kw) for kw in broken_line_keywords) - ] - - if existing_broken and not vlm_broken: - issues.append( - f"✅ VLM merged {len(existing_broken)} broken multi-line titles " - f"(e.g. '{existing_broken[0]}')" - ) - elif vlm_broken: - issues.append( - f"⚠️ VLM still has {len(vlm_broken)} broken multi-line titles: " - f"{vlm_broken}" - ) - - # Check 3: false positive entries (e.g. "Page" as a title) - false_positive_patterns = {"Page", "页码"} - existing_fp = [t for t in existing_titles if t.strip() in false_positive_patterns] - vlm_fp = [ - e for e in vlm_entries - if e.get("title", "").strip() in false_positive_patterns - ] - - if existing_fp and not vlm_fp: - issues.append( - f"✅ VLM removed {len(existing_fp)} false positive entries " - f"(e.g. '{existing_fp[0]}')" - ) - elif vlm_fp: - issues.append( - f"⚠️ VLM still has false positive entries: " - f"{[e['title'] for e in vlm_fp]}" - ) - - # Check 4: entry count comparison - comparison["existing_entry_count"] = len(existing_titles) - comparison["vlm_entry_count"] = len(vlm_entries) - comparison["quality_checks"] = issues - - return comparison - - -# --------------------------------------------------------------------------- -# Main test runner -# --------------------------------------------------------------------------- - - -def run_test(test_case: TocTestCase, model: str) -> dict[str, Any]: - """Run VLM TOC extraction for one test case.""" - print(f"\n{'='*70}") - print(f"TEST: {test_case.name}") - print(f" {test_case.description}") - print(f" TOC pages: {test_case.toc_page_nums}") - print(f"{'='*70}") - - os.makedirs(test_case.output_dir, exist_ok=True) - - # Detect page count for scan_range calculation - pdf_page_count: int | None = None - if test_case.pdf_path: - try: - import pymupdf - with pymupdf.open(test_case.pdf_path) as doc: - pdf_page_count = len(doc) - print(f" PDF page count: {pdf_page_count}") - except Exception: - pass - - # Step 1: Prepare PNGs - png_paths: list[tuple[int, str]] = [] - - if test_case.pre_rendered_pngs: - for i, png in enumerate(test_case.pre_rendered_pngs): - png_paths.append((test_case.toc_page_nums[i], png)) - print(f" [png] Using pre-rendered: {png}") - elif test_case.pdf_path: - for page_num in test_case.toc_page_nums: - png = _render_page_png(test_case.pdf_path, page_num, test_case.output_dir) - png_paths.append((page_num, png)) - print(f" [png] Rendered page {page_num}: {png}") - - # Step 2: VLM extraction per page - all_entries: list[dict[str, Any]] = [] - all_meta: list[dict[str, Any]] = [] - total_elapsed_ms = 0 - - for page_num, png_path in png_paths: - is_continuation = len(all_entries) > 0 - ctx_label = " (with context)" if is_continuation else "" - print(f"\n [vlm] Extracting page {page_num}{ctx_label}...") - try: - entries, meta = _vlm_extract_toc_entries( - png_path, page_num, model, - previous_entries=all_entries if is_continuation else None, - ) - all_entries.extend(entries) - all_meta.append(meta) - total_elapsed_ms += meta["elapsed_ms"] - print(f" [vlm] Page {page_num}: {len(entries)} entries, {meta['elapsed_ms']}ms") - - # Show first few entries - for e in entries[:5]: - title = e.get("title", "?") - pn = e.get("page_number", "?") - lv = e.get("level", "?") - print(f" L{lv}: {title!r} → p.{pn}") - if len(entries) > 5: - print(f" ... ({len(entries)} total)") - - except Exception as exc: - print(f" [vlm] ❌ FAILED for page {page_num}: {exc}") - all_meta.append({"page": page_num, "error": str(exc)}) - - # Step 3: Summary - print(f"\n --- Summary ---") - print(f" Total entries: {len(all_entries)}") - print(f" Total VLM time: {total_elapsed_ms}ms") - expected_lo, expected_hi = test_case.expected_entry_count_range - count_ok = expected_lo <= len(all_entries) <= expected_hi - print( - f" Entry count check: {len(all_entries)} " - f"(expected {expected_lo}-{expected_hi}) → " - f"{'✅' if count_ok else '⚠️ OUT OF RANGE'}" - ) - - # Step 4: Quality analysis - level_dist = {} - for e in all_entries: - lv = e.get("level", "?") - level_dist[lv] = level_dist.get(lv, 0) + 1 - print(f" Level distribution: {level_dist}") - - entries_with_page = [e for e in all_entries if e.get("page_number") is not None] - entries_no_page = [e for e in all_entries if e.get("page_number") is None] - print(f" Entries with page number: {len(entries_with_page)}") - print(f" Entries without page number: {len(entries_no_page)}") - if entries_no_page: - for e in entries_no_page[:5]: - print(f" → {e.get('title', '?')!r} (level={e.get('level')})") - - # Step 5: Compare with existing - comparison = _compare_with_existing(all_entries, test_case) - if comparison.get("available"): - print(f"\n --- Comparison with existing toc_hierarchies.json ---") - print(f" Existing entries: {comparison['existing_entry_count']}") - print(f" VLM entries: {comparison['vlm_entry_count']}") - for check in comparison.get("quality_checks", []): - print(f" {check}") - - # Step 6: Save raw VLM results - result = { - "test_name": test_case.name, - "model": model, - "toc_pages": test_case.toc_page_nums, - "total_entries": len(all_entries), - "total_elapsed_ms": total_elapsed_ms, - "entries": all_entries, - "per_page_meta": all_meta, - "level_distribution": level_dist, - "comparison": comparison if comparison.get("available") else None, - } - - parent_dir = str(Path(test_case.output_dir).parent) - output_path = os.path.join(parent_dir, "vlm_toc_entries.json") - with open(output_path, "w", encoding="utf-8") as f: - json.dump(result, f, ensure_ascii=False, indent=2) - print(f"\n 📄 Raw results saved to: {output_path}") - - # Step 7: Generate toc_hierarchies.json in production schema - toc_hierarchies = vlm_entries_to_toc_hierarchies( - all_entries, test_case.toc_page_nums, - page_count=pdf_page_count, - ) - - toc_hier_path = test_case.toc_output_path or os.path.join( - parent_dir, "toc_hierarchies.json" - ) - os.makedirs(os.path.dirname(toc_hier_path), exist_ok=True) - with open(toc_hier_path, "w", encoding="utf-8") as f: - json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) - print(f" 📄 toc_hierarchies.json saved to: {toc_hier_path}") - - # Print schema summary - if toc_hierarchies: - h = toc_hierarchies[0] - print(f" --- Schema Summary ---") - print(f" toc_range: {h['toc_range']}") - print(f" toc_range_unit: {h['toc_range_unit']}") - print(f" source: {h['source']}") - print(f" toc_with_level entries: {len(h['toc_with_level'])}") - print(f" toc_tree L1 keys: {list(h['toc_tree'].keys())[:8]}") - - # Step 8: Print full entry table - print(f"\n --- Full Entry Table ---") - print(f" {'#':>3} {'Lv':>3} {'Page':>6} Title") - print(f" {'─'*3} {'─'*3} {'─'*6} {'─'*50}") - for i, e in enumerate(all_entries, 1): - title = e.get("title", "?") - pn = e.get("page_number", "—") - lv = e.get("level", "?") - # Truncate long titles for display - disp_title = title if len(title) <= 60 else title[:57] + "..." - print(f" {i:>3} {lv:>3} {str(pn):>6} {disp_title}") - - return result - - -def main() -> None: - # Load env vars from worker/.env if available - script_dir = os.path.dirname(os.path.abspath(__file__)) - # Walk up to find the repo root: test_vlm_toc_extract.py is at - # apps/worker/app/services/document_agent/tools/ - repo_root = os.path.abspath(os.path.join(script_dir, "..", "..", "..", "..", "..", "..")) - env_file = os.path.join(repo_root, "apps", "worker", ".env") - if os.path.exists(env_file): - print(f"Loading env from: {env_file}") - with open(env_file) as f: - for line in f: - line = line.strip() - if line and not line.startswith("#") and "=" in line: - key, _, value = line.partition("=") - key = key.strip() - value = value.strip().strip("'\"") - if key and key not in os.environ: # don't override existing - os.environ[key] = value - - model = os.environ.get("IMAGE_MODEL", "qwen3.5-flash") - print(f"VLM model: {model}") - - test_cases = _build_test_cases() - if not test_cases: - print("No test cases found! Check PDF paths.") - sys.exit(1) - - print(f"Found {len(test_cases)} test case(s)") - - results: list[dict[str, Any]] = [] - for tc in test_cases: - try: - result = run_test(tc, model) - results.append(result) - except Exception as exc: - print(f"\n❌ Test '{tc.name}' FAILED: {exc}") - import traceback - traceback.print_exc() - - # Final summary - print(f"\n{'='*70}") - print("FINAL SUMMARY") - print(f"{'='*70}") - for r in results: - name = r["test_name"] - count = r["total_entries"] - ms = r["total_elapsed_ms"] - print(f" {name}: {count} entries, {ms}ms") - - -if __name__ == "__main__": - main() diff --git a/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py index d37811853..c9e8273c4 100644 --- a/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py @@ -7,8 +7,12 @@ from typing import Any from app.services.document_agent.manifest import PageAnatomyMap, ToolContext, ToolResult -from app.services.document_agent.registry import register_tool -from app.services.document_agent.state import DocumentAgentState +from app.services.document_agent.registry import ( + has_h1_result, + has_shard_plan, + has_toc_result, + register_tool, +) from app.services.document_agent.validators import validate_anatomy_map @@ -27,14 +31,13 @@ def _thresholds(ctx: ToolContext) -> tuple[int, int]: @register_tool( name="validate.anatomy_map", description="Validate page anatomy, hierarchy hints, and shard coverage.", - allowed_states={DocumentAgentState.PLANNED}, + preconditions=(has_toc_result, has_h1_result, has_shard_plan), ) def validate_current_anatomy(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() if not ( ctx.blackboard.toc_result and ctx.blackboard.h1_result - and ctx.blackboard.hierarchy_assist and ctx.blackboard.shard_plan ): return ToolResult( @@ -50,9 +53,8 @@ def validate_current_anatomy(ctx: ToolContext, _args: dict[str, Any]) -> ToolRes page_labels=ctx.blackboard.page_labels, toc_result=ctx.blackboard.toc_result, h1_result=ctx.blackboard.h1_result, - hierarchy_assist=ctx.blackboard.hierarchy_assist, shard_plan=ctx.blackboard.shard_plan, - boundary_candidates=ctx.blackboard.boundary_candidates, + document_profile=ctx.blackboard.document_profile, global_signals=ctx.blackboard.global_signals, trace_summary={}, ) @@ -60,7 +62,7 @@ def validate_current_anatomy(ctx: ToolContext, _args: dict[str, Any]) -> ToolRes report = validate_anatomy_map(anatomy, min_pages=min_pages, max_pages=max_pages) ctx.blackboard.validation_report = report.to_dict() if ctx.blackboard.shard_plan: - ctx.blackboard.shard_plan.validation = report if not report.valid else ctx.blackboard.shard_plan.validation + ctx.blackboard.shard_plan.validation = report return ToolResult( status="ok" if report.valid else "invalid", payload=report.to_dict(), diff --git a/apps/worker/app/services/document_agent/tools/verdict.py b/apps/worker/app/services/document_agent/tools/verdict.py new file mode 100644 index 000000000..bada1ba22 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/verdict.py @@ -0,0 +1,39 @@ +"""Agent verdict tool.""" + +from __future__ import annotations + +import time +from typing import Any + +from app.services.document_agent.manifest import AgentVerdict, ToolContext, ToolResult +from app.services.document_agent.registry import has_shard_plan, register_tool + + +@register_tool( + name="verdict", + description="Finish the document profile run with success or abort.", + parameters={ + "type": "object", + "properties": { + "status": {"type": "string", "enum": ["success", "abort"]}, + "rationale": {"type": "string"}, + }, + "required": ["status", "rationale"], + }, + preconditions=(has_shard_plan,), +) +def verdict(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + status = str(args.get("status") or "abort") + if status not in {"success", "abort"}: + status = "abort" + ctx.blackboard.verdict = AgentVerdict( + status=status, # type: ignore[arg-type] + rationale=str(args.get("rationale") or ""), + ) + return ToolResult( + status="ok", + payload=ctx.blackboard.verdict.to_dict(), + latency_ms=int((time.monotonic() - start) * 1000), + ) + diff --git a/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py new file mode 100644 index 000000000..17821930e --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py @@ -0,0 +1,237 @@ +"""VLM-native TOC entry extraction and hierarchy conversion.""" + +from __future__ import annotations + +import base64 +import json +import time +from typing import Any, cast + + +VLM_TOC_EXTRACT_PROMPT = """\ +You are analyzing a Table of Contents (TOC) page from a document. + +Your task is to extract every TOC entry visible on this page. + +Each entry consists of three parts: +1. title — the section or chapter name, copied verbatim from the page. + - EXCLUDE any trailing dots, dashes, or leader lines that connect the title to its page number. + - If one entry's title wraps across multiple printed lines, combine them into a single string. + - Include any numbering prefix that is part of the title text. +2. page_number — the page reference at the right side of the entry. + - Use an integer when the reference is a plain number. + - Use a string when the reference is non-numeric, such as iv, F-1, or A-3. + - Use null when no page reference is visible for that entry. +3. level — the hierarchy depth of the entry, determined by visual formatting cues: + - level 1: top-level entries with no indentation, or the largest / boldest text. + - level 2: sub-entries indented under a level-1 entry, or in a noticeably smaller font. + - level 3+: deeper indentation, if present. + - Category headers or group labels that are visually distinct and do NOT have a page number should be treated as level 1. + +Additional rules: +- Extract ALL entries, even if the page only shows a partial continuation of the TOC. +- Do NOT include the TOC page's own heading, such as TABLE OF CONTENTS, 目录, or 目 录. +- Do NOT include column headers, such as a standalone Page or 页码 label. +- Preserve the original language and wording of each title exactly. +- If this screenshot is not actually a TOC page, return {"entries": []}. + +Return strict JSON with no markdown fences: +{"entries": [{"title": "...", "page_number": ..., "level": ...}, ...]} +""" + +VLM_TOC_CONTINUATION_CONTEXT = """\ + +--- IMPORTANT: Continuation Context --- +This is a CONTINUATION page of a multi-page Table of Contents. +The previous page(s) already extracted the following entries: + +{previous_summary} + +The LAST active category/section before this page was: + Level {last_l1_level}: "{last_l1_title}" + +Entries on THIS page that visually continue as sub-items under that +category must keep their correct subordinate level. +""" + + +def _build_continuation_context(previous_entries: list[dict[str, Any]]) -> str: + if not previous_entries: + return "" + + tail = previous_entries[-8:] + summary_lines = [] + for entry in tail: + level = entry.get("level", "?") + title = entry.get("title", "?") + page_number = entry.get("page_number") + suffix = f" -> p.{page_number}" if page_number is not None else "" + summary_lines.append(f" L{level}: {title}{suffix}") + + if len(previous_entries) > 8: + summary_lines.insert(0, f" ... ({len(previous_entries) - 8} earlier entries omitted)") + + previous_summary = "\n".join(summary_lines) + last_l1 = None + for entry in reversed(previous_entries): + if entry.get("level") == 1: + last_l1 = entry + break + + if last_l1 is None: + return ( + "\n\n--- IMPORTANT: Continuation Context ---\n" + f"This is a CONTINUATION page. Previous entries:\n{previous_summary}\n" + ) + + return VLM_TOC_CONTINUATION_CONTEXT.format( + previous_summary=previous_summary, + last_l1_level=last_l1.get("level", 1), + last_l1_title=last_l1.get("title", "?"), + ) + + +def vlm_extract_toc_entries( + *, + png_path: str, + page_num: int, + model: str, + previous_entries: list[dict[str, Any]] | None = None, +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Extract visible TOC entries from one rendered page screenshot.""" + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + with open(png_path, "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + + prompt_text = VLM_TOC_EXTRACT_PROMPT + _build_continuation_context( + previous_entries or [] + ) + content_parts: list[dict[str, Any]] = [ + {"type": "text", "text": prompt_text}, + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + }, + ] + start = time.monotonic() + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=model, + temperature=0.1, + max_tokens=4096, + response_format={"type": "json_object"}, + ) + elapsed_ms = int((time.monotonic() - start) * 1000) + data = json.loads(raw) + if isinstance(data, dict): + raw_entries = data.get("entries", []) + elif isinstance(data, list): + raw_entries = data + else: + raw_entries = [] + + entries: list[dict[str, Any]] = [] + for item in raw_entries: + if not isinstance(item, dict): + continue + title = str(item.get("title") or "").strip() + if not title: + continue + try: + level = int(item.get("level") or 1) + except (TypeError, ValueError): + level = 1 + entries.append( + { + "title": title, + "page_number": item.get("page_number"), + "level": level, + } + ) + + return entries, { + "page": page_num, + "model": model, + "elapsed_ms": elapsed_ms, + "usage": dict(usage), + "raw_response_length": len(raw), + "has_continuation_context": bool(previous_entries), + } + + +def build_toc_tree(entries: list[dict[str, Any]]) -> dict[str, Any]: + root: dict[str, Any] = {} + stack: list[tuple[dict[str, Any], int]] = [(root, 0)] + positive_levels = [ + int(entry["level"]) + for entry in entries + if isinstance(entry.get("level"), int) and entry["level"] > 0 + ] + level_for_minus_one = max(positive_levels) + 1 if positive_levels else 1 + + for entry in entries: + heading = str(entry.get("title") or "").strip() + if not heading: + continue + original_level = entry.get("level", 1) + level = level_for_minus_one if original_level == -1 else int(original_level or 1) + while len(stack) > 1 and stack[-1][1] >= level: + stack.pop() + parent = stack[-1][0] + parent[heading] = {} + stack.append((parent[heading], level)) + return root + + +def build_toc_with_level_md(entries: list[dict[str, Any]]) -> str: + if not entries: + return "" + lines = ["| id | heading | level |", "|----|---------|-------|"] + for index, entry in enumerate(entries, 1): + heading = str(entry.get("title") or "").strip().replace("|", "\\|") + level = entry.get("level", 1) + lines.append(f"| {index:<2} | {heading:<60} | {level:<5} |") + return "\n".join(lines) + + +def vlm_entries_to_toc_hierarchies( + entries: list[dict[str, Any]], + *, + toc_page_nums: list[int], + scan_end_page: int | None = None, + page_count: int | None = None, +) -> list[dict[str, Any]]: + if not entries or not toc_page_nums: + return [] + + toc_with_level = [] + for index, entry in enumerate(entries, 1): + toc_with_level.append( + { + "id": index, + "heading": str(entry.get("title") or "").strip(), + "level": entry.get("level", 1), + "page_number": entry.get("page_number"), + } + ) + + start_page = min(toc_page_nums) + end_page = max(toc_page_nums) + if scan_end_page is None: + scan_end_page = start_page + if page_count is not None: + scan_end_page = min(scan_end_page, page_count) + + return [ + { + "toc_range": [start_page, end_page], + "toc_range_unit": "page", + "scan_range": [start_page, scan_end_page], + "source": "vlm", + "toc_with_level": toc_with_level, + "toc_with_level_md": build_toc_with_level_md(entries), + "toc_tree": build_toc_tree(entries), + } + ] diff --git a/apps/worker/app/services/document_agent/trace.py b/apps/worker/app/services/document_agent/trace.py index 30f9b46b6..b455b63e8 100644 --- a/apps/worker/app/services/document_agent/trace.py +++ b/apps/worker/app/services/document_agent/trace.py @@ -61,7 +61,28 @@ def set_anatomy_map(self, anatomy: PageAnatomyMap, artifact_path: str) -> None: self._artifact_path = artifact_path self.write_trace_json(str(Path(artifact_path).with_name("trace.json"))) - def write_trace_json(self, trace_path: str) -> None: + def write_trace_artifact( + self, + output_dir: str | None, + *, + final_status: str, + summary: dict[str, Any] | None = None, + ) -> None: + if output_dir is None: + return + self.write_trace_json( + str(Path(output_dir) / "trace.json"), + final_status=final_status, + summary=summary, + ) + + def write_trace_json( + self, + trace_path: str, + *, + final_status: str | None = None, + summary: dict[str, Any] | None = None, + ) -> None: try: import json @@ -69,7 +90,7 @@ def write_trace_json(self, trace_path: str) -> None: for step in self._steps: item = dict(step) created_at = item.get("created_at") - if hasattr(created_at, "isoformat"): + if created_at is not None and hasattr(created_at, "isoformat"): item["created_at"] = created_at.isoformat() serializable_steps.append(item) Path(trace_path).write_text( @@ -77,6 +98,8 @@ def write_trace_json(self, trace_path: str) -> None: { "run_id": self.run_id, "job_id": self.job_id, + "final_status": final_status, + "summary": summary, "artifact_path": self._artifact_path, "steps": serializable_steps, }, @@ -137,9 +160,7 @@ def flush(self, *, final_status: str, summary: dict[str, Any] | None = None) -> page_plan_id=f"dpp_{uuid4().hex[:12]}", job_id=self.job_id, page_count=self._anatomy.page_count, - hierarchy_assist=self._anatomy.hierarchy_assist.to_dict(), shard_plan=self._anatomy.shard_plan.to_dict(), - page_processing_plan=None, global_signals=self._anatomy.global_signals, ) ) diff --git a/apps/worker/app/services/document_agent/validators.py b/apps/worker/app/services/document_agent/validators.py index 1a5da7777..2f06371ba 100644 --- a/apps/worker/app/services/document_agent/validators.py +++ b/apps/worker/app/services/document_agent/validators.py @@ -3,7 +3,6 @@ from __future__ import annotations from app.services.document_agent.manifest import ( - HierarchyAssistPlan, PageAnatomyMap, Shard, ShardPlan, @@ -11,32 +10,6 @@ ) -def valid_pages(pages: list[int], page_count: int) -> list[int]: - return sorted({page for page in pages if 1 <= int(page) <= page_count}) - - -def repair_hierarchy_assist( - plan: HierarchyAssistPlan, - *, - page_count: int, - toc_pages: set[int], -) -> HierarchyAssistPlan: - plan.exclude_pages_from_title_candidates = valid_pages( - plan.exclude_pages_from_title_candidates, - page_count, - ) - plan.suppress_title_pages = valid_pages(plan.suppress_title_pages, page_count) - plan.prefer_h1_start_pages = [ - candidate - for candidate in plan.prefer_h1_start_pages - if 1 <= candidate.page <= page_count and candidate.page not in toc_pages - ] - plan.section_boundary_hints = [ - hint for hint in plan.section_boundary_hints if 1 <= hint.page <= page_count - ] - return plan - - def validate_shard_plan( plan: ShardPlan, *, @@ -62,8 +35,8 @@ def validate_shard_plan( length = shard.page_end - shard.page_start + 1 if plan.enabled and length > max_pages: errors.append(f"shard {shard.shard_index} exceeds max_pages={max_pages}") - if plan.enabled and length < min_pages and shard.page_end != page_count: - warnings.append(f"shard {shard.shard_index} shorter than min_pages={min_pages}") + if plan.enabled and length < min_pages: + errors.append(f"shard {shard.shard_index} shorter than min_pages={min_pages}") expected_start = shard.page_end + 1 if expected_start != page_count + 1: errors.append("shard_plan does not cover full document") @@ -110,17 +83,6 @@ def validate_anatomy_map( errors.append(f"h1 candidate points to toc page {candidate.page}") if candidate.page < 1 or candidate.page > page_count: errors.append(f"h1 candidate page {candidate.page} out of range") - candidate_count = len(anatomy.boundary_candidates) - if page_count > max_pages and candidate_count == 0: - warnings.append("long document has no boundary candidates") - h1_candidate_count = sum(1 for candidate in anatomy.boundary_candidates if candidate.kind == "h1") - sparse_candidate_count = sum( - 1 - for candidate in anatomy.boundary_candidates - if candidate.kind in {"blank", "sparse", "separator"} - ) - if page_count > max_pages and h1_candidate_count == 0 and sparse_candidate_count == 0: - warnings.append("long document has neither H1 nor sparse/blank boundary candidates") shard_report = validate_shard_plan( anatomy.shard_plan, page_count=page_count, diff --git a/apps/worker/app/services/document_agent/visual.py b/apps/worker/app/services/document_agent/visual.py new file mode 100644 index 000000000..5a6e6da3c --- /dev/null +++ b/apps/worker/app/services/document_agent/visual.py @@ -0,0 +1,86 @@ +"""Shared page rendering helpers for document-agent visual reasoning.""" + +from __future__ import annotations + +import gc +import os +from pathlib import Path +from typing import Any + +from app.services.document_agent.manifest import ToolContext +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) + + +@worker +def _render_pages_worker( + queue, + pdf_path: str, + pages: list[int], + output_dir: str, + dpi: int, + prefix: str, +) -> None: + import pymupdf # type: ignore[import] + + results: list[dict[str, Any]] = [] + try: + doc = pymupdf.open(pdf_path) + for page_num in pages: + idx = page_num - 1 + if 0 <= idx < doc.page_count: + page = doc[idx] + mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) + pix = page.get_pixmap(matrix=mat) + png_name = f"{prefix}_page_{page_num}.png" + png_path = os.path.join(output_dir, png_name) + pix.save(png_path) + results.append({"page": page_num, "png_path": png_path}) + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "results": results}) + + +def visual_output_dir(ctx: ToolContext, folder_name: str = "agent_visuals") -> str: + output_dir = str( + Path(ctx.output_dir or os.path.expanduser("~/.knowhere/_debug_profile")) + / folder_name + ) + os.makedirs(output_dir, exist_ok=True) + return output_dir + + +def render_pages( + ctx: ToolContext, + pages: list[int], + *, + folder_name: str = "agent_visuals", + prefix: str = "visual", + dpi: int | None = None, + timeout: int = 120, +) -> list[dict[str, Any]]: + if not pages: + return [] + page_count = max(int(ctx.blackboard.page_count or 0), 0) + bounded_pages = sorted({page for page in pages if 1 <= page <= page_count}) + if not bounded_pages: + return [] + output_dir = visual_output_dir(ctx, folder_name=folder_name) + effective_dpi = dpi or int(ctx.settings.get("agent_png_dpi", "144")) + result = run_in_child_process( + _render_pages_worker, + ctx.pdf_path, + bounded_pages, + output_dir, + effective_dpi, + prefix, + timeout=timeout, + ) + return list(result.get("results") or []) + diff --git a/apps/worker/build_manifest_sjsyj.py b/apps/worker/build_manifest_sjsyj.py deleted file mode 100644 index 87a0319fb..000000000 --- a/apps/worker/build_manifest_sjsyj.py +++ /dev/null @@ -1,159 +0,0 @@ -""" -从 preds_5_final_output.csv 构建 HIERARCHY 树， -输出 manifest.json 并与原始 MD 的 # 标题对比 -""" - -import json -import os -import re -import pandas as pd -from collections import OrderedDict - -CSV_PATH = ( - "/Users/wuchengke/Desktop/temp/ontos/parse_comparison/" - "sjsyj_mineru_notable/SJSYJ-SC-2024 企业制度汇编（上册）/txt/" - "preds_5_final_output.csv" -) -MD_PATH = ( - "/Users/wuchengke/Desktop/temp/ontos/parse_comparison/" - "sjsyj_mineru_notable/SJSYJ-SC-2024 企业制度汇编（上册）/txt/" - "SJSYJ-SC-2024 企业制度汇编（上册）.md" -) -OUT_DIR = os.path.dirname(CSV_PATH) -SOURCE_FILE = "SJSYJ-SC-2024 企业制度汇编（上册）.pdf" - - -def clean_heading(h: str) -> str: - """去掉 Markdown # 前缀和前后空格""" - return re.sub(r"^#+\s*", "", str(h)).strip() - - -def build_hierarchy(df: pd.DataFrame) -> dict: - """从有层级的行构建嵌套字典树""" - root = OrderedDict() - root["Root"] = {} - stack = [] # [(level, dict)] - - valid = df[df["level"].astype(str).str.match(r"^[1-9]")].copy() - valid["level"] = valid["level"].astype(int) - - for _, row in valid.iterrows(): - title = clean_heading(str(row["heading"])) - level = int(row["level"]) - - # 弹栈到当前父节点 - while stack and stack[-1][0] >= level: - stack.pop() - - parent = stack[-1][1] if stack else root - # 处理重名 - key, suffix = title, 2 - while key in parent: - key = f"{title} ({suffix})" - suffix += 1 - parent[key] = OrderedDict() - stack.append((level, parent[key])) - - return root - - -def extract_md_headings(md_path: str) -> list[dict]: - """从 MD 文件提取所有 # 标题""" - headings = [] - with open(md_path, encoding="utf-8") as f: - for lineno, line in enumerate(f, 1): - m = re.match(r"^(#{1,6})\s+(.+)", line.rstrip()) - if m: - level = len(m.group(1)) - text = m.group(2).strip() - headings.append({"lineno": lineno, "level": level, "text": text}) - return headings - - -def compare_with_md(hierarchy: dict, md_headings: list[dict]): - """对比 hierarchy 树中的标题 vs MD 中的 # 标题""" - def flatten(d, prefix="", result=None): - if result is None: - result = [] - for k, v in d.items(): - result.append(k) - if isinstance(v, dict) and v: - flatten(v, prefix + k + "/", result) - return result - - tree_titles = set(flatten(hierarchy)) - tree_titles.discard("Root") - - md_texts = set(h["text"] for h in md_headings) - - in_tree_not_md = tree_titles - md_texts - in_md_not_tree = md_texts - tree_titles - both = tree_titles & md_texts - - print(f"\n{'='*60}") - print(f"📊 标题对比") - print(f"{'='*60}") - print(f" MD 中 # 标题数: {len(md_texts)}") - print(f" Hierarchy 树标题数: {len(tree_titles)}") - print(f" 完全匹配: {len(both)}") - print(f" 仅在 Hierarchy 中: {len(in_tree_not_md)}") - print(f" 仅在 MD 中: {len(in_md_not_tree)}") - - if in_tree_not_md: - print(f"\n⚠️ Hierarchy 有但 MD # 里没有（LLM 可能误识别）:") - for t in sorted(in_tree_not_md)[:15]: - print(f" - {t[:80]}") - - if in_md_not_tree: - print(f"\n⚠️ MD # 有但 Hierarchy 没有（可能被降级为正文）:") - for t in sorted(in_md_not_tree)[:15]: - print(f" - {t[:80]}") - - return {"matched": len(both), "only_tree": len(in_tree_not_md), "only_md": len(in_md_not_tree)} - - -def main(): - print(f"📄 读取: {CSV_PATH}") - df = pd.read_csv(CSV_PATH, encoding="utf-8-sig") - print(f" 总行数: {len(df)}") - - # 构建树 - hierarchy = build_hierarchy(df) - top_level = [k for k in hierarchy if k != "Root"] - print(f" 顶层章节数: {len(top_level)}") - - # 输出 manifest.json - manifest = { - "version": "2.0", - "job_id": SOURCE_FILE, - "source_file_name": SOURCE_FILE, - "processing_date": "2026-05-21T04:46:35Z", - "statistics": { - "total_chunks": int((df["level"] == -1).sum()), - "heading_count": int((df["level"].astype(str).str.match(r"^[1-9]")).sum()), - }, - "HIERARCHY": hierarchy, - } - - out_path = os.path.join(OUT_DIR, "manifest.json") - with open(out_path, "w", encoding="utf-8") as f: - json.dump(manifest, f, ensure_ascii=False, indent=2) - size_kb = os.path.getsize(out_path) / 1024 - print(f"\n✅ 已保存: {out_path} ({size_kb:.0f} KB)") - - # 预览顶层结构 - print(f"\n📋 顶层章节（前 20）:") - for i, k in enumerate(top_level[:20], 1): - children = hierarchy[k] - child_count = len(children) if isinstance(children, dict) else 0 - print(f" L1 [{i:2d}] {k[:60]}{'...' if len(k)>60 else ''} → {child_count} 子节点") - - # 与 MD 对比 - print(f"\n📄 读取 MD: {MD_PATH}") - md_headings = extract_md_headings(MD_PATH) - print(f" MD # 标题数: {len(md_headings)}") - compare_with_md(hierarchy, md_headings) - - -if __name__ == "__main__": - main() diff --git a/apps/worker/run_hierarchy_sjsyj.py b/apps/worker/run_hierarchy_sjsyj.py deleted file mode 100644 index a3021b7d8..000000000 --- a/apps/worker/run_hierarchy_sjsyj.py +++ /dev/null @@ -1,117 +0,0 @@ -""" -针对 MinerU 解析输出的 MD 文件，运行 pred_titles 完整流程 -输出 preds_3_llm_base.csv 和 preds_5_final_output.csv 到同目录 -""" - -import os -import sys -import json -import shutil - -# ── 配置 ── -MODEL_NAME = "qwen3.5-27b" -ENABLE_THINKING = False - -# 目标 MD 文件 -TARGET_MD = ( - "/Users/wuchengke/Desktop/temp/ontos/parse_comparison/" - "sjsyj_mineru_notable/SJSYJ-SC-2024 企业制度汇编（上册）/txt/" - "SJSYJ-SC-2024 企业制度汇编（上册）.md" -) -OUTPUT_DIR = os.path.dirname(TARGET_MD) - -# ── 工程路径注入 ── -WORKER_DIR = "/Users/wuchengke/Desktop/knowhere/knowhereapi-main/apps/worker" -SHARED_DIR = "/Users/wuchengke/Desktop/knowhere/knowhereapi-main/packages/shared-python" -sys.path.insert(0, SHARED_DIR) -sys.path.insert(0, WORKER_DIR) - -from dotenv import load_dotenv -load_dotenv(os.path.join(WORKER_DIR, ".env")) -os.environ["LOCAL_DEBUG"] = "1" - -import openai -from loguru import logger - -# ── Qwen thinking mode 关闭 patch ── -_original_create = openai.resources.chat.completions.Completions.create - -def patched_create(self, *args, **kwargs): - kwargs.setdefault("extra_body", {}) - kwargs["extra_body"]["enable_thinking"] = False - return _original_create(self, *args, **kwargs) - -openai.resources.chat.completions.Completions.create = patched_create -logger.info("🚫 Thinking mode 已关闭 (enable_thinking=false)") - -from app.services.document_parser.structure.layout_parser import pred_titles -from app.services.document_parser.structure.toc_parser import detect_tocs_in_texts -from app.services.document_parser.formats.html.parser import merge_html_tables - - -def main(): - logger.info("=" * 60) - logger.info(f"🚀 SJSYJ hierarchy 检测 — 模型: {MODEL_NAME}") - logger.info(f" 输入: {TARGET_MD}") - logger.info(f" 输出: {OUTPUT_DIR}") - logger.info("=" * 60) - - # 1. 加载 MD - with open(TARGET_MD, "r", encoding="utf-8") as f: - md_lines = f.read().splitlines() - md_lines = [line.strip() for line in md_lines if line.strip()] - md_lines = merge_html_tables(md_lines) - logger.info(f"📄 MD 加载完毕: {len(md_lines)} 行") - - # 2. TOC 检测 - toc_json_path = os.path.join(OUTPUT_DIR, "toc_hierarchies.json") - toc_hierarchies = None - if os.path.exists(toc_json_path): - os.remove(toc_json_path) - logger.info("🗑️ 已删除旧 toc_hierarchies.json") - - logger.info("🔍 检测 TOC...") - toc_hierarchies, md_lines = detect_tocs_in_texts(md_lines, model_name=MODEL_NAME) - toc_hierarchies = toc_hierarchies or [] - if toc_hierarchies: - with open(toc_json_path, "w", encoding="utf-8") as f: - json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) - logger.info(f"✅ TOC 检测完毕: {len(toc_hierarchies)} 个区域 → {toc_json_path}") - else: - logger.info(" 未检测到 TOC") - - # 3. 运行 pred_titles（完整流程，输出 CSV） - logger.info(f"🧠 运行 pred_titles (model={MODEL_NAME}, smart_parse=True)...") - heading_preds = pred_titles( - infos=md_lines, - doc_type="md", - toc_hierarchies=toc_hierarchies or [], - prompt_limt=4000, - enable_regx=True, - smart_parse=True, - model_name=MODEL_NAME, - output_dir=OUTPUT_DIR, # ← CSV 保存到这里 - layout_json_path=None, - ) - - if heading_preds.empty: - logger.warning("⚠️ 没有检测到任何有效标题") - return - - valid = heading_preds[heading_preds["level"] > 0] - logger.info(f"✅ 完成! 有效标题 {len(valid)} 个 / 总行 {len(heading_preds)}") - logger.info(f" 层级分布:\n{heading_preds['level'].value_counts().sort_index().to_string()}") - - # 汇报生成的 CSV - for csv_name in ["preds_3_llm_base.csv", "preds_4_llm_final.csv", "preds_5_final_output.csv"]: - p = os.path.join(OUTPUT_DIR, csv_name) - if os.path.exists(p): - size_kb = os.path.getsize(p) / 1024 - logger.info(f" 📄 {csv_name} → {p} ({size_kb:.0f} KB)") - - logger.info("=" * 60) - logger.info("🎉 Done!") - - -if __name__ == "__main__": - main() diff --git a/packages/shared-python/shared/models/database/document_page_plan.py b/packages/shared-python/shared/models/database/document_page_plan.py index b7f2424b7..ee2fb228f 100644 --- a/packages/shared-python/shared/models/database/document_page_plan.py +++ b/packages/shared-python/shared/models/database/document_page_plan.py @@ -20,12 +20,7 @@ class DocumentPagePlan(Base): String(36), ForeignKey("jobs.job_id", ondelete="CASCADE"), nullable=False ) page_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) - hierarchy_assist: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) shard_plan: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) - page_processing_plan: Mapped[Optional[Dict[str, Any]]] = mapped_column( - JSON, - nullable=True, - ) global_signals: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) created_at: Mapped[datetime] = mapped_column( DateTime, default=utc_now_naive, nullable=False From 4d770260edd7f9fd86cb5943a9f60687b4a40baa Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 26 May 2026 04:21:39 +0800 Subject: [PATCH 07/11] fix(doc-agent): resolve pyright type error in inspect_pages.py setdefault call --- apps/worker/app/services/document_agent/tools/inspect_pages.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 102f4ae1c..3b31d97e0 100644 --- a/apps/worker/app/services/document_agent/tools/inspect_pages.py +++ b/apps/worker/app/services/document_agent/tools/inspect_pages.py @@ -89,7 +89,7 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: ) ctx.budget.commit("visual", actual=usage.get("total_tokens", est), est=est) try: - payload = json.loads(raw) + payload: dict[str, Any] = json.loads(raw) except json.JSONDecodeError: payload = {"raw": raw} if isinstance(payload, dict): From 838167c2507830d8de9d326336a2c46da47c4f61 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 27 May 2026 10:28:57 +0800 Subject: [PATCH 08/11] feat: integrate budget-aware VLM token usage tracking and improve shard planning resilience with validation fallback logic. --- .../services/document_agent/coordinator.py | 2 +- .../document_agent/executor/react_loop.py | 15 ++++-- .../document_agent/persist/__init__.py | 7 ++- .../document_agent/persist/persist.py | 8 --- .../document_agent/planner/planner.py | 49 ++++++++++--------- .../tools/classify_page_kinds.py | 4 -- .../tools/extract_toc_with_boundaries.py | 45 ++++++++++++++--- .../tools/propose_shard_plan.py | 20 +++++--- 8 files changed, 95 insertions(+), 55 deletions(-) delete mode 100644 apps/worker/app/services/document_agent/persist/persist.py diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index eedd989b6..c32b05d46 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -1,4 +1,4 @@ -"""Plan-then-act coordinator for the document profile agent.""" +"""ReAct-style coordinator for the document profile agent.""" from __future__ import annotations 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 2d95a62ae..ea2daf32d 100644 --- a/apps/worker/app/services/document_agent/executor/react_loop.py +++ b/apps/worker/app/services/document_agent/executor/react_loop.py @@ -288,9 +288,18 @@ def _deterministic_decision(self) -> ReflexionDecision: "rationale": "Validation succeeded; finishing profile run.", }, ) + # Validation failed: fallback to single shard instead of aborting. + # Clear the invalid plan and re-propose as a single shard. + from app.services.document_agent.tools.propose_shard_plan import single_shard_plan + + self.ctx.blackboard.shard_plan = single_shard_plan( + self.ctx.blackboard.page_count + ) + self.ctx.blackboard.validation_report = None return ReflexionDecision( - action="verdict_now", - rationale="Validation failed in deterministic mode.", - verdict=AgentVerdict(status="abort", rationale="Validation failed."), + action="tool_call", + rationale="Validation failed; falling back to single shard plan.", + tool_name="validate.anatomy_map", + tool_args={}, ) diff --git a/apps/worker/app/services/document_agent/persist/__init__.py b/apps/worker/app/services/document_agent/persist/__init__.py index 56ef9f3e1..a34439602 100644 --- a/apps/worker/app/services/document_agent/persist/__init__.py +++ b/apps/worker/app/services/document_agent/persist/__init__.py @@ -1,5 +1,8 @@ -"""Deterministic persistence for document-agent outputs.""" +"""Persist anatomy map artifacts.""" -from app.services.document_agent.persist.persist import build_anatomy_map, persist_anatomy_map +from app.services.document_agent.tools.persist_anatomy_map import ( + build_anatomy_map, + persist_anatomy_map, +) __all__ = ["build_anatomy_map", "persist_anatomy_map"] diff --git a/apps/worker/app/services/document_agent/persist/persist.py b/apps/worker/app/services/document_agent/persist/persist.py deleted file mode 100644 index b83ef8244..000000000 --- a/apps/worker/app/services/document_agent/persist/persist.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Compatibility wrapper for deterministic anatomy-map persistence.""" - -from app.services.document_agent.tools.persist_anatomy_map import ( - build_anatomy_map, - persist_anatomy_map, -) - -__all__ = ["build_anatomy_map", "persist_anatomy_map"] diff --git a/apps/worker/app/services/document_agent/planner/planner.py b/apps/worker/app/services/document_agent/planner/planner.py index 5b361673b..61d63f9f5 100644 --- a/apps/worker/app/services/document_agent/planner/planner.py +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -169,6 +169,26 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: or os.environ.get("IMAGE_MODEL") ) pages = _sample_pages(self.ctx.blackboard.page_count, self.ctx.blackboard.extrema_pages) + if not model: + profile = DocumentProfile( + is_scanned=False, + category="unknown document", + rationale="No planner model configured.", + ) + decision = ReflexionDecision( + action="tool_call", + rationale=profile.rationale, + tool_name="propose.shard_plan", + tool_args={}, + ) + return profile, decision, ToolResult( + status="ok", + payload={"source": "deterministic", "sampled_pages": pages}, + latency_ms=int((time.monotonic() - start) * 1000), + warnings=["No planner model configured; using conservative profile."], + input_summary={"page_count": self.ctx.blackboard.page_count}, + output_summary={"profile": profile.to_dict(), "decision": decision.to_dict()}, + ) pngs = render_pages( self.ctx, pages, @@ -214,28 +234,8 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: ensure_ascii=False, ) prompt_tokens_est = estimate_tokens(prompt_text) + len(pngs) * 800 - if not model: - profile = DocumentProfile( - is_scanned=False, - category="unknown document", - rationale="No planner model configured.", - ) - decision = ReflexionDecision( - action="tool_call", - rationale=profile.rationale, - tool_name="propose.shard_plan", - tool_args={}, - ) - return profile, decision, ToolResult( - status="ok", - payload={"source": "deterministic", "sampled_pages": pages}, - latency_ms=int((time.monotonic() - start) * 1000), - warnings=["No planner model configured; using conservative profile."], - input_summary=payload, - output_summary={"profile": profile.to_dict(), "decision": decision.to_dict()}, - ) - if not self.ctx.budget.try_reserve("plan", prompt_tokens_est): - raise RuntimeError("Insufficient planner budget for profile planning.") + if not self.ctx.budget.try_reserve("visual", prompt_tokens_est): + raise RuntimeError("Insufficient visual budget for profile planning.") content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt_text}] for item in pngs: @@ -266,7 +266,7 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: response_format={"type": "json_object"}, ) self.ctx.budget.commit( - "plan", + "visual", actual=usage.get("total_tokens", prompt_tokens_est), est=prompt_tokens_est, ) @@ -289,6 +289,7 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: }, ) except Exception: - self.ctx.budget.refund("plan", est=prompt_tokens_est) + self.ctx.budget.refund("visual", est=prompt_tokens_est) raise + diff --git a/apps/worker/app/services/document_agent/tools/classify_page_kinds.py b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py index bd7daf67e..fb60765d7 100644 --- a/apps/worker/app/services/document_agent/tools/classify_page_kinds.py +++ b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py @@ -9,10 +9,6 @@ from app.services.document_agent.manifest import PageFeature, PageLabel, ToolContext, ToolResult -def _joined_preview(feature: PageFeature) -> str: - return "\n".join(feature.text_lines_preview).lower() - - def _label_feature(feature: PageFeature) -> PageLabel: page = feature.page if ( diff --git a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py index ac457ab41..21083eb5d 100644 --- a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py +++ b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py @@ -9,6 +9,8 @@ from pathlib import Path from typing import Any, cast +from shared.utils.token_estimate import estimate_tokens + from app.services.document_agent.manifest import ( TocAnchorPage, TocResult, @@ -65,6 +67,7 @@ def _render_single_page_worker( def _vlm_confirm_anchors( anchor_pages: list[TocAnchorPage], model: str, + budget: Any | None = None, ) -> tuple[list[TocAnchorPage], bool]: """Phase 1: send all anchor PNGs to VLM, ask which are real TOC starts.""" from shared.services.ai.openai_compatible_client_sync import get_openai_client @@ -115,6 +118,10 @@ def _vlm_confirm_anchors( ) messages = cast(Any, [{"role": "user", "content": content_parts}]) + est = estimate_tokens(str(content_parts[0]["text"])) + len(anchor_pages) * 800 + if budget and not budget.try_reserve("visual", est): + logger.warning("[extract.toc] insufficient visual budget for anchor confirmation") + return [], True try: client = get_openai_client(model=model) @@ -125,6 +132,8 @@ def _vlm_confirm_anchors( max_tokens=500, response_format={"type": "json_object"}, ) + if budget: + budget.commit("visual", actual=usage.get("total_tokens", est), est=est) data = json.loads(raw) if isinstance(data, dict): items = data.get("pages") or data.get("results") or data.get("data") or [] @@ -149,6 +158,8 @@ def _vlm_confirm_anchors( ) return confirmed, False except Exception as exc: + if budget: + budget.refund("visual", est=est) logger.warning( "[extract.toc] VLM anchor confirmation failed: {}, " "falling back to no confirmed anchors (safe degradation)", @@ -161,6 +172,7 @@ def _vlm_check_boundary_page( png_path: str, page_num: int, model: str, + budget: Any | None = None, ) -> tuple[bool, bool]: """Phase 2: check if a single page still contains TOC content.""" from shared.services.ai.openai_compatible_client_sync import get_openai_client @@ -198,18 +210,27 @@ def _vlm_check_boundary_page( }, ] + est = estimate_tokens(str(content_parts[0]["text"])) + 800 + if budget and not budget.try_reserve("visual", est): + logger.warning("[extract.toc] insufficient visual budget for boundary check p{}", page_num) + return False, True + try: client = get_openai_client(model=model) - raw, _ = client.chat_completion_with_usage( + raw, usage = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, temperature=0.1, max_tokens=200, response_format={"type": "json_object"}, ) + if budget: + budget.commit("visual", actual=usage.get("total_tokens", est), est=est) data = json.loads(raw) return bool(data.get("still_toc")), False except Exception as exc: + if budget: + budget.refund("visual", est=est) logger.warning( "[extract.toc] boundary VLM check failed for page {}: {}", page_num, exc ) @@ -228,6 +249,7 @@ def _detect_toc_range_for_anchor( output_dir: str, dpi: int, model: str, + budget: Any | None = None, ) -> tuple[int, int, list[dict[str, Any]], list[str]]: """Progressively expand from anchor_page to find the TOC end boundary. @@ -251,7 +273,7 @@ def _detect_toc_range_for_anchor( timeout=60, ) - still_toc, failed = _vlm_check_boundary_page(png_path, check_page, model) + still_toc, failed = _vlm_check_boundary_page(png_path, check_page, model, budget=budget) if failed: warnings.append(f"vlm_boundary_check_failed:p{check_page}") trace_rounds.append( @@ -312,9 +334,19 @@ def extract_toc_with_boundaries( latency_ms=int((time.monotonic() - start) * 1000), ) - model = ctx.settings.get("vlm_model") or os.environ.get( - "IMAGE_MODEL", "qwen3.5-flash" - ) + model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") + if not model: + logger.warning("[extract.toc] no VLM model configured; skipping TOC extraction") + ctx.blackboard.toc_result = TocResult( + method="none", + notes="No VLM model configured for TOC extraction", + ) + return ToolResult( + status="ok", + payload={"toc_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + warnings=["No VLM model configured; skipping TOC extraction."], + ) dpi = int(ctx.settings.get("toc_png_dpi", "144")) page_count = ctx.blackboard.page_count output_dir = str( @@ -324,7 +356,7 @@ def extract_toc_with_boundaries( os.makedirs(output_dir, exist_ok=True) # -- Phase 1: VLM confirm anchors ----------------------------------------- - confirmed, confirm_failed = _vlm_confirm_anchors(anchors, model) + confirmed, confirm_failed = _vlm_confirm_anchors(anchors, model, budget=ctx.budget) if confirm_failed: warnings.append("vlm_anchor_confirmation_failed") debug_info["phase1_confirmed"] = [a.page for a in confirmed] @@ -357,6 +389,7 @@ def extract_toc_with_boundaries( output_dir=output_dir, dpi=dpi, model=model, + budget=ctx.budget, ) warnings.extend(boundary_warnings) toc_ranges.append((toc_start, toc_end)) 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 dc5bfda52..82be670f6 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 @@ -4,7 +4,6 @@ import json import os -import re import time from typing import Any @@ -124,11 +123,12 @@ def _build_prompt( ) -def _sanitize_rationale(text: str) -> str: - # Keep rationales structural. The model may quote page preview text; those - # literals are useful as input evidence but should not become baked-in rules. - sanitized = re.sub(r"'[^']{1,40}'", "sparse separator text", text or "") - sanitized = re.sub(r'"[^"]{1,40}"', "sparse separator text", sanitized) +def _sanitize_rationale(text: str, max_length: int = 120) -> str: + # Truncate overlong rationales but preserve H1 title references + # which provide valuable semantic context for shard boundaries. + sanitized = (text or "").strip() + if len(sanitized) > max_length: + sanitized = sanitized[:max_length].rstrip() + "…" return sanitized @@ -195,6 +195,7 @@ def _parse_llm_plan( def _deterministic_guardrail_plan( *, page_count: int, + min_pages: int, max_pages: int, h1_pages: list[int], ) -> tuple[list[tuple[int, str, str, float]], str]: @@ -213,6 +214,9 @@ def _deterministic_guardrail_plan( else: cuts.append((target, "forced_max_size", "guardrail max shard size", 0.25)) previous = target + # Merge final shard into previous if it's smaller than min_pages + if cuts and (page_count - cuts[-1][0]) < min_pages: + cuts.pop() return cuts, "too_large" @@ -265,7 +269,7 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: client = get_openai_client(model=model) raw_response, usage = client.chat_completion_with_usage( - messages=prompt, + messages=[{"role": "user", "content": prompt}], model=model, temperature=0.0, max_tokens=1600, @@ -284,6 +288,7 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: ) cuts, reason = _deterministic_guardrail_plan( page_count=page_count, + min_pages=min_pages, max_pages=max_pages, h1_pages=[item["page"] for item in h1_pages], ) @@ -296,6 +301,7 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: ) cuts, reason = _deterministic_guardrail_plan( page_count=page_count, + min_pages=min_pages, max_pages=max_pages, h1_pages=[item["page"] for item in h1_pages], ) From f39b9bae6ef1789937e18b86fbbf4e3748972d35 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 27 May 2026 21:00:43 +0800 Subject: [PATCH 09/11] refactor: improve heading LLM prompt accuracy with strict TOC rules and implement PDF shard splitting and merging logic. --- .../tools/extract_toc_with_boundaries.py | 345 +++++++----------- .../tools/find_toc_anchor_pages.py | 22 ++ .../document_agent/tools/vlm_toc_extractor.py | 312 +++++++++++----- .../formats/markdown/parser.py | 33 +- .../document_parser/formats/pdf/parser.py | 255 +++++++------ .../formats/pdf/shard_merger.py | 86 +++++ .../formats/pdf/shard_splitter.py | 166 +++++++++ .../orchestration/parse_session.py | 43 ++- .../structure/heading_llm_executor.py | 36 +- .../structure/layout_parser.py | 22 +- .../shared/core/config/storage.py | 22 +- .../shared/services/ai/prompt_service.py | 77 ++-- 12 files changed, 898 insertions(+), 521 deletions(-) create mode 100644 apps/worker/app/services/document_parser/formats/pdf/shard_merger.py create mode 100644 apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py diff --git a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py index 21083eb5d..4a424a5cc 100644 --- a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py +++ b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py @@ -20,7 +20,6 @@ from app.services.document_agent.registry import register_tool from app.services.document_agent.tools.vlm_toc_extractor import ( vlm_entries_to_toc_hierarchies, - vlm_extract_toc_entries, ) from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( run_in_child_process, @@ -168,150 +167,14 @@ def _vlm_confirm_anchors( return [], True -def _vlm_check_boundary_page( - png_path: str, - page_num: int, - model: str, - budget: Any | None = None, -) -> tuple[bool, bool]: - """Phase 2: check if a single page still contains TOC content.""" - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - import base64 - - with open(png_path, "rb") as f: - img_b64 = base64.b64encode(f.read()).decode() - - content_parts: list[dict[str, Any]] = [ - { - "type": "text", - "text": ( - "You are a document structure analysis expert. " - "Below is a screenshot of one page from a PDF. " - "I am determining the boundary of a Table of Contents (TOC) region.\n\n" - "Does this page still contain TOC content?\n\n" - "TOC content characteristics:\n" - "- Entry titles paired with page numbers\n" - "- Dots (...), leader lines (.....), or spaces connecting titles to page numbers\n" - "- Systematic numbering (e.g. 1. / 1.1 / Chapter 1 / (1))\n\n" - "NOT TOC content:\n" - "- Body text paragraphs\n" - "- Data tables\n" - "- Image-heavy pages\n" - "- A single heading with no page-number listing\n\n" - "Return strict JSON (no markdown fences):\n" - '{"still_toc": true/false, "confidence": "high"/"medium"/"low", ' - '"reason": "brief reason"}' - ), - }, - { - "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{img_b64}"}, - }, - ] - - est = estimate_tokens(str(content_parts[0]["text"])) + 800 - if budget and not budget.try_reserve("visual", est): - logger.warning("[extract.toc] insufficient visual budget for boundary check p{}", page_num) - return False, True - - try: - client = get_openai_client(model=model) - raw, usage = client.chat_completion_with_usage( - messages=cast(Any, [{"role": "user", "content": content_parts}]), - model=model, - temperature=0.1, - max_tokens=200, - response_format={"type": "json_object"}, - ) - if budget: - budget.commit("visual", actual=usage.get("total_tokens", est), est=est) - data = json.loads(raw) - return bool(data.get("still_toc")), False - except Exception as exc: - if budget: - budget.refund("visual", est=est) - logger.warning( - "[extract.toc] boundary VLM check failed for page {}: {}", page_num, exc - ) - # Conservative: stop expansion on failure - return False, True - - -# -- Progressive boundary detection -------------------------------------------- - - -def _detect_toc_range_for_anchor( - *, - anchor_page: int, - pdf_path: str, - page_count: int, - output_dir: str, - dpi: int, - model: str, - budget: Any | None = None, -) -> tuple[int, int, list[dict[str, Any]], list[str]]: - """Progressively expand from anchor_page to find the TOC end boundary. - - Returns: - (start_page, end_page, trace_rounds) -- all 1-based inclusive. - """ - start_page = anchor_page - current_end = min(anchor_page + BOUNDARY_STEP_PAGES - 1, page_count) - trace_rounds: list[dict[str, Any]] = [] - warnings: list[str] = [] - - for round_idx in range(MAX_BOUNDARY_ROUNDS): - check_page = current_end - png_path = os.path.join(output_dir, f"toc_boundary_p{check_page}.png") - run_in_child_process( - _render_single_page_worker, - pdf_path, - check_page, - png_path, - dpi, - timeout=60, - ) - - still_toc, failed = _vlm_check_boundary_page(png_path, check_page, model, budget=budget) - if failed: - warnings.append(f"vlm_boundary_check_failed:p{check_page}") - trace_rounds.append( - { - "round": round_idx, - "check_page": check_page, - "window": [start_page, current_end], - "still_toc": still_toc, - } - ) - logger.info( - "[extract.toc] round {}: page {} still_toc={}", - round_idx, - check_page, - still_toc, - ) - - if not still_toc: - # Boundary page is NOT TOC; TOC ends at previous page. - current_end = max(check_page - 1, start_page) - break - - next_end = min(current_end + BOUNDARY_STEP_PAGES, page_count) - if next_end == current_end: - break - current_end = next_end - - return start_page, current_end, trace_rounds, warnings - - # -- Main tool ----------------------------------------------------------------- @register_tool( name="extract.toc_with_boundaries", description=( - "VLM-confirms TOC anchor pages, progressively detects TOC boundaries, " - "then extracts TOC entries directly from rendered pages with VLM." + "VLM-confirms TOC anchor pages, then batch-classifies and extracts " + "TOC entries from rendered page windows using VLM." ), ) def extract_toc_with_boundaries( @@ -377,89 +240,144 @@ def extract_toc_with_boundaries( debug=debug_info, ) - # -- Phase 2: progressive boundary detection ------------------------------- - toc_ranges: list[tuple[int, int]] = [] - all_trace_rounds: list[dict[str, Any]] = [] + # -- Phase 2+3 (unified): batch classify + extract --------------------------- + # Instead of separate boundary detection (Phase 2) then per-page extraction + # (Phase 3), we send batches of BOUNDARY_STEP_PAGES images to VLM in one + # call. The VLM classifies each page (TOC vs non-TOC) AND extracts entries + # from TOC pages simultaneously. If the last page in a batch is still TOC, + # we expand the window and use prior entries as continuation context. + from app.services.document_agent.tools.vlm_toc_extractor import ( + vlm_extract_toc_batch, + ) - for anchor in confirmed: - toc_start, toc_end, trace_rounds, boundary_warnings = _detect_toc_range_for_anchor( - anchor_page=anchor.page, - pdf_path=ctx.pdf_path, - page_count=page_count, - output_dir=output_dir, - dpi=dpi, - model=model, - budget=ctx.budget, - ) - warnings.extend(boundary_warnings) - toc_ranges.append((toc_start, toc_end)) - all_trace_rounds.extend(trace_rounds) - logger.info( - "[extract.toc] TOC region: pages {}-{}", toc_start, toc_end - ) + all_entries: list[dict[str, Any]] = [] + all_toc_pages: list[int] = [] + toc_hierarchies: list[dict[str, Any]] = [] + batch_meta: list[dict[str, Any]] = [] + batch_trace: list[dict[str, Any]] = [] - debug_info["phase2_ranges"] = toc_ranges - debug_info["phase2_trace_rounds"] = all_trace_rounds + for anchor in confirmed: + anchor_page = anchor.page + region_entries: list[dict[str, Any]] = [] + region_toc_pages: list[int] = [] + region_scan_end = anchor_page + + for round_idx in range(MAX_BOUNDARY_ROUNDS): + batch_start = anchor_page + round_idx * BOUNDARY_STEP_PAGES + batch_end = min( + batch_start + BOUNDARY_STEP_PAGES - 1, page_count + ) + if batch_start > page_count: + break + + batch_pages = list(range(batch_start, batch_end + 1)) + logger.info( + "[extract.toc] batch round {}: pages {}-{} for anchor {}", + round_idx, batch_start, batch_end, anchor_page, + ) + + # Render all pages in this batch + page_pngs: list[tuple[int, str]] = [] + for page_num in batch_pages: + png_path = os.path.join(output_dir, f"toc_page_{page_num}.png") + run_in_child_process( + _render_single_page_worker, + ctx.pdf_path, + page_num, + png_path, + dpi, + timeout=60, + ) + page_pngs.append((page_num, png_path)) + + # Send batch to VLM — classify + extract in one call + batch_result = vlm_extract_toc_batch( + page_pngs=page_pngs, + model=model, + previous_entries=region_entries if region_entries else None, + ) + batch_meta.append(batch_result.meta) + + # Collect results + region_entries.extend(batch_result.all_entries) + region_toc_pages.extend(batch_result.toc_pages) + region_scan_end = batch_end + + batch_trace.append({ + "anchor": anchor_page, + "round": round_idx, + "batch_pages": batch_pages, + "toc_pages": batch_result.toc_pages, + "non_toc_pages": batch_result.non_toc_pages, + "entries_found": len(batch_result.all_entries), + }) + + # Determine if we need to continue expanding + # If the last page in the batch is NOT TOC, boundary found + last_page_is_toc = ( + batch_result.page_results + and batch_result.page_results[-1].is_toc + ) + if not last_page_is_toc: + logger.info( + "[extract.toc] boundary found: last page {} is not TOC", + batch_end, + ) + break + + # Last page is still TOC — continue expanding + if batch_end >= page_count: + break + logger.info( + "[extract.toc] last page {} still TOC, expanding window", + batch_end, + ) + + all_entries.extend(region_entries) + all_toc_pages.extend(region_toc_pages) + + if region_entries: + region_hierarchies = vlm_entries_to_toc_hierarchies( + region_entries, + toc_page_nums=region_toc_pages, + scan_end_page=region_scan_end, + page_count=page_count, + ) + toc_hierarchies.extend(region_hierarchies) + else: + logger.warning( + "[extract.toc] anchor {} produced no TOC entries", + anchor_page, + ) - # -- Phase 3: VLM entry extraction ----------------------------------------- - all_toc_pages = sorted( - {p for s, e in toc_ranges for p in range(s, e + 1)} - ) - all_entries: list[dict[str, Any]] = [] - per_page_meta: list[dict[str, Any]] = [] - rendered_pages: list[dict[str, Any]] = [] - for page_num in all_toc_pages: - png_path = os.path.join(output_dir, f"toc_page_{page_num}.png") - render_result = run_in_child_process( - _render_single_page_worker, - ctx.pdf_path, - page_num, - png_path, - dpi, - timeout=60, - ) - rendered_pages.append(render_result) - entries, meta = vlm_extract_toc_entries( - png_path=str(render_result.get("png_path") or png_path), - page_num=page_num, - model=model, - previous_entries=all_entries, + if not all_entries: + raise RuntimeError( + "VLM TOC extractor returned no entries for confirmed TOC pages" ) - all_entries.extend(entries) - per_page_meta.append(meta) - if not all_entries: - raise RuntimeError("VLM TOC extractor returned no entries for confirmed TOC pages") + debug_info["batch_trace"] = batch_trace + debug_info["batch_meta"] = batch_meta + debug_info["vlm_entry_count"] = len(all_entries) - scan_end_page = max( - (round_info.get("check_page", 0) for round_info in all_trace_rounds), - default=max(all_toc_pages), - ) - toc_hierarchies = vlm_entries_to_toc_hierarchies( - all_entries, - toc_page_nums=all_toc_pages, - scan_end_page=int(scan_end_page), - page_count=page_count, - ) - debug_info["phase3_vlm_entry_count"] = len(all_entries) - debug_info["phase3_vlm_per_page_meta"] = per_page_meta - debug_info["phase3_rendered_pages"] = rendered_pages + all_toc_pages_sorted = sorted(set(all_toc_pages)) + toc_region_count = len(toc_hierarchies) ctx.blackboard.toc_result = TocResult( - toc_pages=all_toc_pages, - method="vlm_progressive", + toc_pages=all_toc_pages_sorted, + method="vlm_batch", notes=( f"VLM confirmed {len(confirmed)} TOC starts, " - f"expanded to {len(toc_ranges)} ranges: {toc_ranges}" + f"batch classify+extract found {toc_region_count} regions, " + f"toc_pages={all_toc_pages_sorted}" ), ) ctx.blackboard.toc_hierarchies = toc_hierarchies if toc_hierarchies else None ctx.blackboard.global_signals["vlm_toc_entries"] = { "model": model, - "toc_pages": all_toc_pages, + "toc_pages": all_toc_pages_sorted, "total_entries": len(all_entries), "entries": all_entries, - "per_page_meta": per_page_meta, + "batch_meta": batch_meta, } # Persist toc_hierarchies to disk for inspection / downstream reuse @@ -472,11 +390,18 @@ def extract_toc_with_boundaries( except Exception as exc: logger.warning("[extract.toc] failed to write toc_hierarchies: {}", exc) + # Build toc_ranges from confirmed TOC pages for summary + toc_ranges_out: list[list[int]] = [] + if toc_hierarchies: + for hier in toc_hierarchies: + toc_ranges_out.append(hier.get("toc_range", [])) + toc_summary: dict[str, Any] = { - "toc_ranges": toc_ranges, - "toc_page_count": len(all_toc_pages), + "toc_ranges": toc_ranges_out, + "toc_page_count": len(all_toc_pages_sorted), "toc_entry_count": len(all_entries), - "toc_source": "vlm", + "toc_region_count": toc_region_count, + "toc_source": "vlm_batch", } if toc_hierarchies: for i, hier in enumerate(toc_hierarchies): @@ -488,10 +413,12 @@ def extract_toc_with_boundaries( status="ok", payload={ "toc_count": len(toc_hierarchies) if toc_hierarchies else 0, - "toc_page_count": len(all_toc_pages), + "toc_page_count": len(all_toc_pages_sorted), + "toc_region_count": toc_region_count, }, latency_ms=int((time.monotonic() - start) * 1000), output_summary=toc_summary, warnings=warnings, debug=debug_info, ) + diff --git a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py index f6cfdae17..db03c90eb 100644 --- a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py +++ b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py @@ -136,14 +136,36 @@ def find_toc_anchor_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult raw_hit_pages: set[int] = set() for feature in ctx.blackboard.page_features: + page_matched = False for line_idx, raw_line in enumerate(feature.text_lines_preview): norm_line = _normalize_for_toc(raw_line) for keyword in TOC_KEYWORDS: if keyword in norm_line: keyword_matches.append((feature.page, raw_line.strip(), line_idx)) raw_hit_pages.add(feature.page) + page_matched = True break # one match per line is enough + # Fallback: check if a TOC keyword spans across adjacent lines. + # PyMuPDF sometimes splits large headings across lines, e.g. + # "目" + "录" or "Table of" + "Contents". Join the first few + # preview lines (where a page title would appear) and re-check + # with the same keywords and normalisation. + if not page_matched and feature.text_lines_preview: + head = feature.text_lines_preview[:10] + joined_head = _normalize_for_toc("".join(head)) + for keyword in TOC_KEYWORDS: + if keyword in joined_head: + keyword_matches.append((feature.page, keyword, 0)) + raw_hit_pages.add(feature.page) + logger.debug( + "[find.toc_anchor_pages] cross-line keyword '{}' " + "detected on page {} (head lines joined)", + keyword, + feature.page, + ) + break + # Apply recurring element fingerprint filter if keyword_matches: anchor_pages = _filter_recurring_elements(keyword_matches, total_pages) diff --git a/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py index 17821930e..7c1d62479 100644 --- a/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py +++ b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py @@ -5,57 +5,109 @@ import base64 import json import time +from dataclasses import dataclass from typing import Any, cast -VLM_TOC_EXTRACT_PROMPT = """\ -You are analyzing a Table of Contents (TOC) page from a document. - -Your task is to extract every TOC entry visible on this page. - -Each entry consists of three parts: -1. title — the section or chapter name, copied verbatim from the page. - - EXCLUDE any trailing dots, dashes, or leader lines that connect the title to its page number. - - If one entry's title wraps across multiple printed lines, combine them into a single string. - - Include any numbering prefix that is part of the title text. -2. page_number — the page reference at the right side of the entry. - - Use an integer when the reference is a plain number. - - Use a string when the reference is non-numeric, such as iv, F-1, or A-3. - - Use null when no page reference is visible for that entry. -3. level — the hierarchy depth of the entry, determined by visual formatting cues: - - level 1: top-level entries with no indentation, or the largest / boldest text. - - level 2: sub-entries indented under a level-1 entry, or in a noticeably smaller font. - - level 3+: deeper indentation, if present. - - Category headers or group labels that are visually distinct and do NOT have a page number should be treated as level 1. - -Additional rules: -- Extract ALL entries, even if the page only shows a partial continuation of the TOC. -- Do NOT include the TOC page's own heading, such as TABLE OF CONTENTS, 目录, or 目 录. -- Do NOT include column headers, such as a standalone Page or 页码 label. -- Preserve the original language and wording of each title exactly. -- If this screenshot is not actually a TOC page, return {"entries": []}. - -Return strict JSON with no markdown fences: -{"entries": [{"title": "...", "page_number": ..., "level": ...}, ...]} +# --------------------------------------------------------------------------- +# Batch-mode prompt: send a window of candidate pages in one VLM call. +# The VLM first classifies each page, then extracts entries only from TOC pages. +# --------------------------------------------------------------------------- + +VLM_TOC_BATCH_PROMPT = """\ +You will receive {page_count} consecutive page screenshots from a document. +Some of these pages may be Table of Contents (TOC) pages, while others may be +regular body text, section dividers, blank pages, or other non-TOC content. + +**Your task has two parts:** + +### Part 1: Classify each page +For each page, decide whether it is a TOC page or not. + +A page IS a TOC page when it shows a STRUCTURED LISTING of document sections, +recognizable by MOST of these visual patterns: +- Multiple entry lines, each pairing a section/chapter TITLE with a PAGE NUMBER +- Leader characters (dots "......", dashes "------", or whitespace) connecting + titles on the left to page numbers aligned on the right +- Systematic numbering in the titles (1. / 1.1 / Chapter 1 / 一、 / 第一章, etc.) +- An explicit heading such as "Table of Contents", "Contents", "目录", or "目次" + (may appear only on the first page of a multi-page TOC) + +A page is NOT a TOC page when: +- It contains narrative paragraphs or body text, even if the text has numbered + headings (e.g. "1.0.1 为建立并落实..." followed by explanatory sentences) +- It is a section divider / title page with only a single heading and no listing +- It is blank or nearly blank +- It shows data tables, charts, or images rather than a contents listing +- It has numbered definitions or terms with explanations (e.g. "2.0.3 风险 risk") + — these are glossary/body content, NOT a TOC + +The KEY distinction: TOC entries are SHORT titles pointing to page numbers. +Body text has EXPLANATORY content after the heading. If a numbered item is +followed by sentences of explanation, it is body text, not a TOC entry. + +### Part 2: Extract entries from TOC pages only +For each page you classify as TOC, extract every entry with: +- title: the section/chapter name, verbatim, without trailing dots or leaders. + Combine wrapped lines into one string. Include numbering prefixes. +- page_number: integer for plain numbers, string for non-numeric (iv, F-1), + null when no page reference is visible. +- level: hierarchy depth from visual cues (1=top-level, 2=indented sub-entry, 3+=deeper). + Category headers or group labels without page numbers → level 1. + +Do NOT include the TOC heading itself ("Table of Contents", "目录", etc.) or +column labels ("Page", "页码"). + +Return strict JSON (no markdown fences): +{{ + "pages": [ + {{ + "page": , + "is_toc": true/false, + "entries": [{{"title": "...", "page_number": ..., "level": ...}}, ...] + }}, + ... + ] +}} + +For non-TOC pages, set "entries" to an empty array []. """ -VLM_TOC_CONTINUATION_CONTEXT = """\ +VLM_TOC_BATCH_CONTINUATION = """\ ---- IMPORTANT: Continuation Context --- -This is a CONTINUATION page of a multi-page Table of Contents. -The previous page(s) already extracted the following entries: +--- Continuation Context --- +Previous batch(es) already confirmed TOC pages and extracted these entries: {previous_summary} -The LAST active category/section before this page was: - Level {last_l1_level}: "{last_l1_title}" +Last active section: Level {last_l1_level}: "{last_l1_title}" -Entries on THIS page that visually continue as sub-items under that -category must keep their correct subordinate level. +Use this to maintain hierarchy consistency for any TOC pages in this batch. """ -def _build_continuation_context(previous_entries: list[dict[str, Any]]) -> str: +@dataclass +class BatchPageResult: + """Result for a single page within a batch VLM call.""" + + page: int + is_toc: bool + entries: list[dict[str, Any]] + + +@dataclass +class BatchTocResult: + """Result from a batch VLM TOC extraction call.""" + + page_results: list[BatchPageResult] + toc_pages: list[int] # pages classified as TOC + non_toc_pages: list[int] # pages classified as non-TOC + all_entries: list[dict[str, Any]] # entries from TOC pages only + meta: dict[str, Any] + + +def _build_batch_continuation(previous_entries: list[dict[str, Any]]) -> str: + """Build continuation context for batch mode.""" if not previous_entries: return "" @@ -69,7 +121,9 @@ def _build_continuation_context(previous_entries: list[dict[str, Any]]) -> str: summary_lines.append(f" L{level}: {title}{suffix}") if len(previous_entries) > 8: - summary_lines.insert(0, f" ... ({len(previous_entries) - 8} earlier entries omitted)") + summary_lines.insert( + 0, f" ... ({len(previous_entries) - 8} earlier entries omitted)" + ) previous_summary = "\n".join(summary_lines) last_l1 = None @@ -80,85 +134,148 @@ def _build_continuation_context(previous_entries: list[dict[str, Any]]) -> str: if last_l1 is None: return ( - "\n\n--- IMPORTANT: Continuation Context ---\n" - f"This is a CONTINUATION page. Previous entries:\n{previous_summary}\n" + "\n\n--- Continuation Context ---\n" + f"Previous batch extracted entries:\n{previous_summary}\n" ) - return VLM_TOC_CONTINUATION_CONTEXT.format( + return VLM_TOC_BATCH_CONTINUATION.format( previous_summary=previous_summary, last_l1_level=last_l1.get("level", 1), last_l1_title=last_l1.get("title", "?"), ) -def vlm_extract_toc_entries( +def vlm_extract_toc_batch( *, - png_path: str, - page_num: int, + page_pngs: list[tuple[int, str]], model: str, previous_entries: list[dict[str, Any]] | None = None, -) -> tuple[list[dict[str, Any]], dict[str, Any]]: - """Extract visible TOC entries from one rendered page screenshot.""" +) -> BatchTocResult: + """Extract TOC entries from a batch of page images in a single VLM call. + + Args: + page_pngs: list of (page_number, png_path) pairs, in page order. + model: VLM model name. + previous_entries: entries from prior batches, for continuation context. + + Returns: + BatchTocResult with per-page classification and extracted entries. + """ + from loguru import logger from shared.services.ai.openai_compatible_client_sync import get_openai_client - with open(png_path, "rb") as f: - img_b64 = base64.b64encode(f.read()).decode() + if not page_pngs: + return BatchTocResult( + page_results=[], toc_pages=[], non_toc_pages=[], + all_entries=[], meta={}, + ) + + prompt_text = VLM_TOC_BATCH_PROMPT.format(page_count=len(page_pngs)) + prompt_text += _build_batch_continuation(previous_entries or []) - prompt_text = VLM_TOC_EXTRACT_PROMPT + _build_continuation_context( - previous_entries or [] - ) content_parts: list[dict[str, Any]] = [ {"type": "text", "text": prompt_text}, - { - "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{img_b64}"}, - }, ] + for page_num, png_path in page_pngs: + with open(png_path, "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + content_parts.append( + {"type": "text", "text": f"\n--- Page {page_num} ---"} + ) + content_parts.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + } + ) + start = time.monotonic() client = get_openai_client(model=model) raw, usage = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, temperature=0.1, - max_tokens=4096, + max_tokens=8192, response_format={"type": "json_object"}, ) elapsed_ms = int((time.monotonic() - start) * 1000) + data = json.loads(raw) + raw_pages: list[dict[str, Any]] = [] if isinstance(data, dict): - raw_entries = data.get("entries", []) + raw_pages = data.get("pages", []) elif isinstance(data, list): - raw_entries = data - else: - raw_entries = [] - - entries: list[dict[str, Any]] = [] - for item in raw_entries: - if not isinstance(item, dict): - continue - title = str(item.get("title") or "").strip() - if not title: - continue - try: - level = int(item.get("level") or 1) - except (TypeError, ValueError): - level = 1 - entries.append( - { - "title": title, - "page_number": item.get("page_number"), - "level": level, - } + raw_pages = data + + # Build lookup from VLM response + page_lookup: dict[int, dict[str, Any]] = {} + for item in raw_pages: + if isinstance(item, dict) and "page" in item: + page_lookup[int(item["page"])] = item + + # Process results for each page in the original order + page_results: list[BatchPageResult] = [] + toc_pages: list[int] = [] + non_toc_pages: list[int] = [] + all_entries: list[dict[str, Any]] = [] + + for page_num, _png_path in page_pngs: + vlm_page = page_lookup.get(page_num, {}) + is_toc = bool(vlm_page.get("is_toc", False)) + raw_entries = vlm_page.get("entries", []) + + entries: list[dict[str, Any]] = [] + if is_toc: + for entry_item in raw_entries: + if not isinstance(entry_item, dict): + continue + title = str(entry_item.get("title") or "").strip() + if not title: + continue + try: + level = int(entry_item.get("level") or 1) + except (TypeError, ValueError): + level = 1 + entries.append( + { + "title": title, + "page_number": entry_item.get("page_number"), + "level": level, + } + ) + toc_pages.append(page_num) + else: + non_toc_pages.append(page_num) + + page_results.append( + BatchPageResult(page=page_num, is_toc=is_toc, entries=entries) ) + all_entries.extend(entries) + + logger.info( + "[vlm_toc_batch] {} pages: toc={} non_toc={} entries={} elapsed={}ms", + len(page_pngs), + toc_pages, + non_toc_pages, + len(all_entries), + elapsed_ms, + ) + + return BatchTocResult( + page_results=page_results, + toc_pages=toc_pages, + non_toc_pages=non_toc_pages, + all_entries=all_entries, + meta={ + "pages_sent": [p for p, _ in page_pngs], + "model": model, + "elapsed_ms": elapsed_ms, + "usage": dict(usage), + "raw_response_length": len(raw), + "has_continuation_context": bool(previous_entries), + }, + ) - return entries, { - "page": page_num, - "model": model, - "elapsed_ms": elapsed_ms, - "usage": dict(usage), - "raw_response_length": len(raw), - "has_continuation_context": bool(previous_entries), - } def build_toc_tree(entries: list[dict[str, Any]]) -> dict[str, Any]: @@ -186,13 +303,21 @@ def build_toc_tree(entries: list[dict[str, Any]]) -> dict[str, Any]: def build_toc_with_level_md(entries: list[dict[str, Any]]) -> str: + """Build a compact Markdown table of TOC entries (heading + level only). + + Accepts both raw VLM entries (key='title') and stored toc_with_level + dicts (key='heading'). Call this dynamically at consumption time rather + than persisting the MD string. + """ if not entries: return "" - lines = ["| id | heading | level |", "|----|---------|-------|"] - for index, entry in enumerate(entries, 1): - heading = str(entry.get("title") or "").strip().replace("|", "\\|") + lines = ["| heading | level |", "|---------|-------|"] + for entry in entries: + heading = str( + entry.get("heading") or entry.get("title") or "" + ).strip().replace("|", "\\|") level = entry.get("level", 1) - lines.append(f"| {index:<2} | {heading:<60} | {level:<5} |") + lines.append(f"| {heading} | {level} |") return "\n".join(lines) @@ -207,10 +332,9 @@ def vlm_entries_to_toc_hierarchies( return [] toc_with_level = [] - for index, entry in enumerate(entries, 1): + for entry in entries: toc_with_level.append( { - "id": index, "heading": str(entry.get("title") or "").strip(), "level": entry.get("level", 1), "page_number": entry.get("page_number"), @@ -231,7 +355,7 @@ def vlm_entries_to_toc_hierarchies( "scan_range": [start_page, scan_end_page], "source": "vlm", "toc_with_level": toc_with_level, - "toc_with_level_md": build_toc_with_level_md(entries), "toc_tree": build_toc_tree(entries), } ] + diff --git a/apps/worker/app/services/document_parser/formats/markdown/parser.py b/apps/worker/app/services/document_parser/formats/markdown/parser.py index a33be14e9..39b4e7483 100755 --- a/apps/worker/app/services/document_parser/formats/markdown/parser.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parser.py @@ -213,6 +213,7 @@ def parse_md( md_lines=None, base_llm_paras=None, relative_root=None, + toc_hierarchies=None, ): if md_lines is None and file_path is not None: from app.services.common.file_loading import is_remote, load_file_bytes @@ -242,14 +243,32 @@ def parse_md( else (settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL) ) - with stage_timer( - "md.detect_toc", line_count=len(md_lines), model_name=toc_model_name - ): - toc_hierarchies, md_lines = detect_tocs_in_texts( - md_lines, - model_name=toc_model_name, - hierarchy_model_name=hierarchy_model_name, + if toc_hierarchies is not None: + # Pre-detected TOC from upstream (e.g. DOC_AGENT VLM-based extraction). + # Skip row-based detection entirely — TOC pages have already been + # physically stripped from the PDF, so no TOC rows exist in md_lines. + logger.info( + f"📌 Using pre-detected TOC hierarchies " + f"({len(toc_hierarchies)} regions), " + f"skipping detect_tocs_in_texts" ) + else: + # DEPRECATED: Row-based TOC detection via detect_tocs_in_texts. + # This path is being replaced by upstream DOC_AGENT VLM-based TOC + # detection (extract_toc_with_boundaries). It is kept as a fallback + # for: + # 1. Standard-path PDFs (≤MAX_PDF_PAGE_LIMIT pages, no DOC_AGENT) + # 2. Non-PDF sources (pure markdown files) + # Once DOC_AGENT TOC detection covers all paths, this block should + # be removed. + with stage_timer( + "md.detect_toc", line_count=len(md_lines), model_name=toc_model_name + ): + toc_hierarchies, md_lines = detect_tocs_in_texts( + md_lines, + model_name=toc_model_name, + hierarchy_model_name=hierarchy_model_name, + ) # Save toc_hierarchies.json to output_dir (will be included in final zip package) if toc_hierarchies: diff --git a/apps/worker/app/services/document_parser/formats/pdf/parser.py b/apps/worker/app/services/document_parser/formats/pdf/parser.py index 6a8dfaf4a..e40850a6a 100755 --- a/apps/worker/app/services/document_parser/formats/pdf/parser.py +++ b/apps/worker/app/services/document_parser/formats/pdf/parser.py @@ -5,10 +5,11 @@ from app.services.document_parser.formats.markdown.parser import parse_md from app.services.document_parser.providers.mineru.pdf_service import parse_via_full -from app.services.document_parser.formats.pdf.pymupdf_subprocess import worker from app.services.document_parser.support.stage_profiler import stage_timer from loguru import logger +from shared.core.config import settings + def _inject_page_markers(output_dir: str) -> None: """Inject markers into full.md using layout.json page info. @@ -109,104 +110,6 @@ def _inject_page_markers(output_dir: str) -> None: logger.info(f"Injected {len(insertions)} page markers into full.md") -def _inject_page_markers_pymupdf(pdf_path: str, output_dir: str) -> None: - """Inject markers into full.md for pymupdf4llm fast path. - - Must run inside the same process that holds the PyMuPDF import. - """ - import pymupdf - - md_path = os.path.join(output_dir, "full.md") - if not os.path.exists(md_path): - return - - try: - doc = pymupdf.open(pdf_path) - except Exception: - return - - with open(md_path, "r", encoding="utf-8") as f: - md_lines = f.readlines() - - anchors = [] - for page_idx in range(len(doc)): - page = doc[page_idx] - page_num = page_idx + 1 - blocks = page.get_text("blocks") - for block in blocks: - if block[6] == 0: - text = block[4].strip().split("\n")[0].strip() - if text and len(text) >= 3: - anchors.append((text, page_num)) - break - - doc.close() - - if not anchors: - return - - insertions = [] - used_lines = set() - - for anchor_text, page_num in anchors: - anchor_norm = re.sub(r"\s+", " ", anchor_text).strip() - search_key = anchor_norm[:50] - for i, line in enumerate(md_lines): - if i in used_lines: - continue - line_norm = re.sub(r"^#+\s*", "", line.strip()) - line_norm = re.sub(r"\s+", " ", line_norm).strip() - if search_key in line_norm: - insertions.append((i, page_num)) - used_lines.add(i) - break - - if not insertions: - return - - insertions.sort(key=lambda x: x[0], reverse=True) - for line_idx, page_num in insertions: - md_lines.insert(line_idx, f"\n") - - with open(md_path, "w", encoding="utf-8") as f: - f.writelines(md_lines) - - -# ─── Child-process workers (top-level for pickling) ───────────────── - - -@worker -def _fast_path_worker(queue, pdf_path, output_dir, image_dir): - """Child process: pymupdf4llm extraction + page marker injection.""" - import pymupdf - import pymupdf4llm - - doc = pymupdf.open(pdf_path) - try: - md_text = pymupdf4llm.to_markdown( - doc, - write_images=True, - image_path=image_dir, - image_format="png", - ) - finally: - doc.close() - - full_md_path = os.path.join(output_dir, "full.md") - with open(full_md_path, "w", encoding="utf-8") as f: - f.write(md_text) - - _inject_page_markers_pymupdf(pdf_path, output_dir) - - img_count = len([n for n in os.listdir(image_dir) if n.endswith(".png")]) - queue.put( - { - "ok": True, - "md_chars": len(md_text), - "image_count": img_count, - } - ) - def parse_pdfs( pdf_path, @@ -220,7 +123,7 @@ def parse_pdfs( route = profile.route if profile else "standard" base_llm_paras.update({"doc_name": filename}) - # ── Atlas routing: bypass MinerU entirely, use PyMuPDF for per-page chunking ── + # ── Atlas routing: bypass MinerU entirely ── if profile and profile.doc_category == "atlas": logger.info(f"📐 Atlas detected, bypassing MinerU for {filename}") from app.services.document_parser.formats.atlas.parser import parse_atlas @@ -229,39 +132,21 @@ def parse_pdfs( pdf_path, output_dir, base_llm_paras, relative_root, profile=profile ) - # TODO: Re-enable fast path after thorough debugging. - # Conservative strategy: until the fast path (pymupdf4llm) is fully validated, - # all non-atlas PDFs are forced to MinerU (standard route) regardless of what - # DocProfiler recommends. The routing logic below is intentionally bypassed. - # - # Original fast-path block (keep for reference, do NOT delete): - # if route == "fast": - # logger.info(f"⚡ Fast path: extracting with pymupdf4llm for {filename}") - # - # os.makedirs(output_dir, exist_ok=True) - # image_dir = os.path.join(output_dir, "images") - # os.makedirs(image_dir, exist_ok=True) - # - # with stage_timer("pdf.extract.fast", filename=filename): - # result = run_in_child_process( - # _fast_path_worker, pdf_path, output_dir, image_dir, - # ) - # logger.info( - # f"⚡ Fast path done: {result['md_chars']} chars, " - # f"{result['image_count']} images" - # ) - # else: - # with stage_timer("pdf.extract.standard", filename=filename): - # parse_via_full(pdf_path, filename, output_dir, s3_key=s3_key) - # _inject_page_markers(output_dir) + # ── Oversized PDF: doc_agent → shard → parallel MinerU → merge → parse_md ── + if profile and profile.page_count > settings.MAX_PDF_PAGE_LIMIT: + logger.info( + f"📄 Oversized PDF: {profile.page_count} pages > " + f"{settings.MAX_PDF_PAGE_LIMIT} limit, entering shard pipeline" + ) + return _parse_oversized_pdf( + pdf_path, filename, output_dir, base_llm_paras, + profile=profile, relative_root=relative_root, s3_key=s3_key, + ) - logger.info( - f"🛡️ Conservative mode: forcing MinerU (standard) for {filename} [route={route}]" - ) + # ── Standard single-pass MinerU ── + logger.info(f"📄 Standard MinerU parse for {filename} [route={route}]") with stage_timer("pdf.extract.standard", filename=filename): parse_via_full(pdf_path, filename, output_dir, s3_key=s3_key) - - # Inject page markers from MinerU layout.json _inject_page_markers(output_dir) logger.info("✅ PDF parsing step 1 complete: text extracted") @@ -274,3 +159,113 @@ def parse_pdfs( base_llm_paras=base_llm_paras, relative_root=relative_root, ) + + +def _parse_oversized_pdf( + pdf_path, filename, output_dir, base_llm_paras, + profile=None, relative_root=None, s3_key=None, +): + """Handle PDFs exceeding MinerU's page limit via doc_agent shard-and-stitch.""" + from concurrent.futures import ThreadPoolExecutor, as_completed + + from app.services.document_parser.formats.pdf.shard_merger import ( + merge_shard_outputs, + ) + from app.services.document_parser.formats.pdf.shard_splitter import ( + bin_pack_shards, + run_doc_agent, + split_pdf, + ) + + job_id = base_llm_paras.get("doc_name", filename) + + # 1. Run doc_agent to get full anatomy map (shard plan + TOC info) + with stage_timer("pdf.doc_agent", filename=filename): + anatomy = run_doc_agent(pdf_path, job_id=job_id, output_dir=output_dir) + + agent_shards = anatomy.shard_plan.shards + + # 2. Extract TOC info from anatomy for page exclusion and heading constraint + toc_pages: set[int] = set() + toc_hierarchies = None + if anatomy.toc_result and anatomy.toc_result.toc_pages: + toc_pages = set(anatomy.toc_result.toc_pages) + toc_hierarchies = anatomy.toc_hierarchies + logger.info( + f"📌 DOC_AGENT TOC detected: {len(toc_pages)} pages to exclude " + f"({sorted(toc_pages)}), " + f"{len(toc_hierarchies) if toc_hierarchies else 0} hierarchy regions" + ) + + # 3. Bin-pack agent shards to maximize MinerU page limit + merged_shards = bin_pack_shards(agent_shards, max_pages=settings.MAX_PDF_PAGE_LIMIT) + logger.info( + f"📦 Bin-packed {len(agent_shards)} agent shards → " + f"{len(merged_shards)} MinerU shards" + ) + for ms in merged_shards: + logger.info( + f" shard_{ms.shard_index}: pages {ms.page_start}-{ms.page_end} " + f"({ms.page_count} pages)" + ) + + # 4. Physically split PDF (exclude TOC pages if detected) + work_dir = os.path.join(output_dir, "_shards") + os.makedirs(work_dir, exist_ok=True) + with stage_timer("pdf.split", filename=filename): + shard_pdf_paths, _page_remap = split_pdf( + pdf_path, merged_shards, work_dir, + exclude_pages=toc_pages if toc_pages else None, + ) + + # 5. Parse each shard via MinerU (parallel) + shard_output_dirs: list[str | None] = [None] * len(shard_pdf_paths) + concurrency = settings.MINERU_SHARD_CONCURRENCY + + def _parse_single_shard(shard_idx, shard_pdf): + shard_out = os.path.join(work_dir, f"shard_{shard_idx}_output") + os.makedirs(shard_out, exist_ok=True) + shard_filename = ( + f"{os.path.splitext(filename)[0]}_shard{shard_idx}.pdf" + ) + logger.info( + f" 🔄 MinerU shard_{shard_idx}: parsing" + ) + parse_via_full(shard_pdf, shard_filename, shard_out, s3_key=None) + return shard_out + + with stage_timer( + "pdf.mineru_parallel", filename=filename, shard_count=len(shard_pdf_paths) + ): + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = { + executor.submit(_parse_single_shard, i, shard_pdf_path): i + for i, shard_pdf_path in enumerate(shard_pdf_paths) + } + for future in as_completed(futures): + idx = futures[future] + shard_output_dirs[idx] = future.result() + + # 6. Merge all shard outputs into main output_dir + with stage_timer("pdf.merge_shards", filename=filename): + merge_shard_outputs(shard_output_dirs, merged_shards, output_dir) + + # 7. Inject page markers (uses merged layout.json with corrected page_idx) + # Note: page markers may be inaccurate when TOC pages are excluded, but + # this only affects chunk metadata (page_nums), not heading hierarchy. + _inject_page_markers(output_dir) + + logger.info("✅ Oversized PDF shard-and-stitch complete, entering parse_md") + + # 8. Standard parse_md — pass DOC_AGENT TOC hierarchies to skip + # row-based TOC detection and enable hard-constraint heading assignment + with stage_timer("pdf.parse_md", filename=filename): + return parse_md( + output_dir, + source_type="md", + file_path=os.path.join(output_dir, "full.md"), + base_llm_paras=base_llm_paras, + relative_root=relative_root, + toc_hierarchies=toc_hierarchies, + ) + diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py new file mode 100644 index 000000000..f1b78e623 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py @@ -0,0 +1,86 @@ +"""Merge MinerU outputs from multiple shards into a single unified output.""" + +from __future__ import annotations + +import json +import os +import shutil +from typing import TYPE_CHECKING + +from loguru import logger + +if TYPE_CHECKING: + from app.services.document_parser.formats.pdf.shard_splitter import MergedShard + + +def merge_shard_outputs( + shard_output_dirs: list[str], + shards: list[MergedShard], + target_output_dir: str, +) -> None: + """Merge MinerU outputs from all shards into target_output_dir. + + After this call target_output_dir contains: + - full.md: concatenated markdown (UUID image refs, no conflicts) + - layout.json: merged page array with corrected page_idx + - images/: all images from all shards + """ + _merge_full_md(shard_output_dirs, target_output_dir) + _merge_layout_json(shard_output_dirs, shards, target_output_dir) + _merge_images(shard_output_dirs, target_output_dir) + + +def _merge_full_md(shard_dirs: list[str], target_dir: str) -> None: + target_path = os.path.join(target_dir, "full.md") + with open(target_path, "w", encoding="utf-8") as out: + for i, shard_dir in enumerate(shard_dirs): + md_path = os.path.join(shard_dir, "full.md") + if not os.path.exists(md_path): + logger.warning(f"shard {i}: full.md not found at {md_path}") + continue + with open(md_path, "r", encoding="utf-8") as f: + content = f.read() + if i > 0: + out.write("\n\n") + out.write(content) + logger.info(f"Merged {len(shard_dirs)} full.md files → {target_path}") + + +def _merge_layout_json( + shard_dirs: list[str], + shards: list[MergedShard], + target_dir: str, +) -> None: + merged_pages: list[dict] = [] + for shard, shard_dir in zip(shards, shard_dirs): + layout_path = os.path.join(shard_dir, "layout.json") + if not os.path.exists(layout_path): + logger.warning(f"shard {shard.shard_index}: layout.json not found") + continue + with open(layout_path, "r", encoding="utf-8") as f: + layout_data = json.load(f) + for page in layout_data.get("pdf_info", []): + page["page_idx"] = page.get("page_idx", 0) + shard.page_offset + merged_pages.append(page) + + target_path = os.path.join(target_dir, "layout.json") + with open(target_path, "w", encoding="utf-8") as f: + json.dump({"pdf_info": merged_pages}, f, ensure_ascii=False) + logger.info(f"Merged layout.json: {len(merged_pages)} pages → {target_path}") + + +def _merge_images(shard_dirs: list[str], target_dir: str) -> None: + target_img_dir = os.path.join(target_dir, "images") + os.makedirs(target_img_dir, exist_ok=True) + total = 0 + for shard_dir in shard_dirs: + img_dir = os.path.join(shard_dir, "images") + if not os.path.isdir(img_dir): + continue + for fname in os.listdir(img_dir): + src = os.path.join(img_dir, fname) + dst = os.path.join(target_img_dir, fname) + if os.path.isfile(src): + shutil.copy2(src, dst) + total += 1 + logger.info(f"Merged {total} images → {target_img_dir}") diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py new file mode 100644 index 000000000..ebf599145 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py @@ -0,0 +1,166 @@ +"""PDF shard splitting: doc_agent integration + bin-packing + physical split.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import pymupdf +from loguru import logger + +if TYPE_CHECKING: + from app.services.document_agent.manifest import PageAnatomyMap, Shard + + +@dataclass +class MergedShard: + """A contiguous page range within MinerU's per-request page limit.""" + + shard_index: int + page_start: int # 1-based inclusive + page_end: int # 1-based inclusive + + @property + def page_count(self) -> int: + return self.page_end - self.page_start + 1 + + @property + def page_offset(self) -> int: + """Offset to add to MinerU's 0-based page_idx to get original page_idx.""" + return self.page_start - 1 + + +def run_doc_agent( + pdf_path: str, job_id: str, output_dir: str +) -> "PageAnatomyMap": + """Run doc_agent ProfileCoordinator and return the full anatomy map. + + Returns the complete PageAnatomyMap so callers can access TOC info + (toc_result.toc_pages, toc_hierarchies) in addition to the shard plan. + + Raises RuntimeError if the agent fails or produces no shards. + """ + from app.services.document_agent.coordinator import ProfileCoordinator + + agent_output_dir = os.path.join(output_dir, "_doc_agent") + os.makedirs(agent_output_dir, exist_ok=True) + + coordinator = ProfileCoordinator( + pdf_path=pdf_path, + job_id=job_id, + output_dir=agent_output_dir, + ) + anatomy = coordinator.run() + + if not anatomy.shard_plan.enabled or not anatomy.shard_plan.shards: + raise RuntimeError( + f"Doc agent did not produce a valid shard plan for {job_id}" + ) + + shards = anatomy.shard_plan.shards + logger.info( + f"📋 Doc agent: {len(shards)} shards via {anatomy.shard_plan.reason}" + ) + return anatomy + + +def bin_pack_shards( + agent_shards: list["Shard"], + max_pages: int, +) -> list[MergedShard]: + """Greedy left-to-right bin-packing: merge adjacent agent shards up to max_pages.""" + if not agent_shards: + return [] + + merged: list[MergedShard] = [] + cur_start = agent_shards[0].page_start + cur_end = agent_shards[0].page_end + + for shard in agent_shards[1:]: + if shard.page_end - cur_start + 1 <= max_pages: + cur_end = shard.page_end + else: + merged.append( + MergedShard(len(merged), page_start=cur_start, page_end=cur_end) + ) + cur_start = shard.page_start + cur_end = shard.page_end + + merged.append( + MergedShard(len(merged), page_start=cur_start, page_end=cur_end) + ) + return merged + + +def split_pdf( + pdf_path: str, + shards: list[MergedShard], + work_dir: str, + exclude_pages: set[int] | None = None, +) -> tuple[list[str], dict[int, int] | None]: + """Physically split PDF into sub-PDFs using PyMuPDF. + + Args: + pdf_path: Path to the source PDF. + shards: Merged shard ranges to extract. + work_dir: Directory for temporary shard PDFs. + exclude_pages: Optional set of 1-based page numbers to strip + (e.g. TOC pages detected by DOC_AGENT). + + Returns: + (shard_paths, page_remap) + - shard_paths: one temp PDF path per shard. + - page_remap: when pages are excluded, maps each shard's local + 0-based page index to the original 1-based page number. + ``None`` when no pages are excluded. + """ + doc = pymupdf.open(pdf_path) + paths: list[str] = [] + page_remap: dict[int, int] | None = None + + if exclude_pages: + page_remap = {} + logger.info( + f"📌 Excluding {len(exclude_pages)} pages from PDF: " + f"{sorted(exclude_pages)}" + ) + + try: + global_new_idx = 0 # running counter across all shards + for shard in shards: + sub_doc = pymupdf.open() + shard_included = 0 + for page_num in range(shard.page_start, shard.page_end + 1): + if exclude_pages and page_num in exclude_pages: + continue + sub_doc.insert_pdf( + doc, + from_page=page_num - 1, + to_page=page_num - 1, + ) + if page_remap is not None: + page_remap[global_new_idx] = page_num + global_new_idx += 1 + shard_included += 1 + + shard_path = os.path.join(work_dir, f"shard_{shard.shard_index}.pdf") + if shard_included > 0: + sub_doc.save(shard_path) + paths.append(shard_path) + else: + logger.warning( + f" ⚠️ shard_{shard.shard_index}: all pages excluded, skipping" + ) + sub_doc.close() + + excluded_in_shard = shard.page_count - shard_included + logger.info( + f" ✂️ shard_{shard.shard_index}: " + f"pages {shard.page_start}-{shard.page_end} " + f"({shard_included} included" + f"{f', {excluded_in_shard} excluded' if excluded_in_shard else ''})" + ) + finally: + doc.close() + return paths, page_remap diff --git a/apps/worker/app/services/document_parser/orchestration/parse_session.py b/apps/worker/app/services/document_parser/orchestration/parse_session.py index 7e1141a84..b90323553 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_session.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -109,20 +109,37 @@ def build_parse_session(parse_input: ParseInput) -> ParseSession: f"ℹ️ VLM rejected atlas for {parse_input.filename}, routing as generic" ) - pdf_page_limit = settings.MAX_PDF_PAGE_LIMIT - if profile.file_type == "pdf" and profile.page_count > pdf_page_limit: - raise ValidationException( - user_message=( - f"Document too large: {profile.page_count} pages exceeds the {pdf_page_limit}-page limit. " - "Please split the document and upload in smaller batches." - ), - violations=[ - { + if profile.file_type == "pdf" and profile.page_count > settings.MAX_PDF_PAGE_LIMIT: + if profile.page_count > settings.OVERSIZED_PDF_SOFT_LIMIT: + raise ValidationException( + user_message=( + f"This document has {profile.page_count} pages. Processing ultra-long " + f"documents (over {settings.OVERSIZED_PDF_SOFT_LIMIT} pages) requires " + "dedicated resources. Please contact our support team for assistance." + ), + violations=[{ "field": "page_count", - "description": f"PDF has {profile.page_count} pages, limit is {pdf_page_limit}", - } - ], - ) + "description": ( + f"PDF has {profile.page_count} pages, " + f"soft limit is {settings.OVERSIZED_PDF_SOFT_LIMIT}" + ), + }], + ) + if not settings.OVERSIZED_PDF_SHARD_ENABLED: + raise ValidationException( + user_message=( + f"Document has {profile.page_count} pages, exceeding the " + f"{settings.MAX_PDF_PAGE_LIMIT}-page limit. Please split the " + "document into smaller parts and upload them separately." + ), + violations=[{ + "field": "page_count", + "description": ( + f"PDF has {profile.page_count} pages, " + f"limit is {settings.MAX_PDF_PAGE_LIMIT}" + ), + }], + ) if profile.doc_category == "atlas": filename, internal_output_filename, relative_root, full_output_dir = ( diff --git a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py index 6f314416b..9b984f796 100644 --- a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py +++ b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py @@ -345,9 +345,16 @@ def get_code_status(row: pd.Series) -> str: def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: - """Collapse consecutive body rows into placeholder rows before LLM chunking.""" + """Collapse consecutive body rows into placeholder rows before LLM chunking. + + The output DataFrame has columns [id, heading, reason]. The ``level`` + column is intentionally NOT forwarded to the LLM — preliminary estimates + were found to mislead the model more often than they helped. The naive- + stage body-text detection (level == -1) is still used here to decide which + rows become placeholders vs candidates. + """ if df is None or len(df) == 0: - return pd.DataFrame(columns=["id", "heading", "level", "reason"]) + return pd.DataFrame(columns=["id", "heading", "reason"]) rows: list[dict[str, Any]] = [] index = 0 @@ -376,7 +383,6 @@ def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: { "id": f"{start_id}-{end_id}", "heading": f"[{run_length} BODY LINES]", - "level": "-", "reason": PLACEHOLDER_REASON, } ) @@ -387,17 +393,12 @@ def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: { "id": int(row["id"]), "heading": str(row["heading"]), - "level": ( - int(lvl_int) - if lvl_int is not None and lvl_int != -2 - else "Not Sure" - ), "reason": str(row.get("reason", "") or ""), } ) index += 1 - return pd.DataFrame(rows, columns=["id", "heading", "level", "reason"]) + return pd.DataFrame(rows, columns=["id", "heading", "reason"]) def split_heading_table( @@ -413,7 +414,8 @@ def split_heading_table( current_rows: list[list[Any]] = [] current_len = 0 for _, row in working_df.iterrows(): - row_filtered = row.drop(labels=["reason"], errors="ignore") + # Drop internal-only columns before measuring token length + row_filtered = row.drop(labels=["reason", "level"], errors="ignore") row_len = sum(count_cn_en(str(value)) for value in row_filtered.values) if current_len + row_len > threshold and current_rows: @@ -501,7 +503,7 @@ def execute_llm_heading_hierarchy( model_name=model_name, ): logger.debug("smart parse => interpreting hierarchy patterns...") - df4llm = basic_df.drop(columns=["reason"]).copy() + df4llm = basic_df.drop(columns=["reason", "level"], errors="ignore").copy() df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm) logger.debug(f"DataFrame transformation completed, rows: {len(df4llm)}") @@ -544,9 +546,15 @@ def level_for(row_id: Any) -> Any: figure_mask_base = base_preds["heading"].eq("Figure/Image") exclude_mask_base = placeholder_mask_base | figure_mask_base base_preds_for_mapping = base_preds[~exclude_mask_base].copy() - base_origin_for_mapping = basic_df.loc[ - ~exclude_mask_base.values, "level" - ].tolist() + # Read origin levels from raw_preds (naive-stage estimates), + # not from compact output which no longer carries level. + base_candidate_ids = base_preds_for_mapping["id"].tolist() + raw_level_lookup = dict( + zip(raw_preds["id"].tolist(), raw_preds["level"].tolist()) + ) + base_origin_for_mapping = [ + raw_level_lookup.get(cid, -1) for cid in base_candidate_ids + ] base_preds_for_mapping, lvl_mapping = build_level_mapping( base_preds_for_mapping, base_origin_for_mapping, mode="freq" diff --git a/apps/worker/app/services/document_parser/structure/layout_parser.py b/apps/worker/app/services/document_parser/structure/layout_parser.py index 47f0650af..52541e793 100755 --- a/apps/worker/app/services/document_parser/structure/layout_parser.py +++ b/apps/worker/app/services/document_parser/structure/layout_parser.py @@ -121,20 +121,16 @@ def format_toc_context_for_llm(toc_context) -> str: formatted_blocks.append("- No TOC entries available") continue - for entry in toc_entries: - if not isinstance(entry, dict): - continue - - heading = str(entry.get("heading", "")).strip().replace("\n", " ") - if not heading: - continue + # Dynamically generate MD table from JSON list (more token-efficient) + from app.services.document_agent.tools.vlm_toc_extractor import ( + build_toc_with_level_md, + ) - level = entry.get("level") - line_id = entry.get("id") - if isinstance(level, int): - formatted_blocks.append(f"- level {level} | id {line_id} | {heading}") - else: - formatted_blocks.append(f"- id {line_id} | {heading}") + md_table = build_toc_with_level_md(toc_entries) + if md_table: + formatted_blocks.append(md_table) + else: + formatted_blocks.append("- No TOC entries available") return "\n".join(formatted_blocks) diff --git a/packages/shared-python/shared/core/config/storage.py b/packages/shared-python/shared/core/config/storage.py index 3b3afbe61..b4eeda866 100644 --- a/packages/shared-python/shared/core/config/storage.py +++ b/packages/shared-python/shared/core/config/storage.py @@ -63,9 +63,27 @@ class StorageConfig(BaseModel): default=104857600, description="Maximum file size in bytes" ) MAX_PDF_PAGE_LIMIT: int = Field( - default=600, + default=200, ge=1, - description="Maximum allowed PDF page count before parsing is rejected", + description="Maximum PDF page count for single-pass MinerU parsing. " + "Documents exceeding this trigger the shard pipeline when enabled.", + ) + OVERSIZED_PDF_SHARD_ENABLED: bool = Field( + default=False, + description="Enable doc_agent shard pipeline for PDFs exceeding MAX_PDF_PAGE_LIMIT. " + "When False, oversized PDFs are rejected.", + ) + OVERSIZED_PDF_SOFT_LIMIT: int = Field( + default=1500, + ge=1, + description="Soft page limit for oversized PDF shard pipeline. " + "Documents exceeding this are rejected with a contact-support message.", + ) + MINERU_SHARD_CONCURRENCY: int = Field( + default=3, + ge=1, + le=10, + description="Maximum concurrent MinerU API calls for shard parsing.", ) SUPPORTED_EXTENSIONS: str = Field( default=".doc,.docx,.pdf,.txt,.xls,.xlsx,.pptx,.jpg,.jpeg,.png,.md", diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index 339e00d7d..e587294e7 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -339,13 +339,11 @@ def build_prompt(task, texts, query, **kwargs): # """ elif task == "eval-headings": - # COMPACT-input variant. Input is pre-compressed by `_compact_for_llm` + # COMPACT-input variant. Input is pre-compressed by `compact_for_llm` # so that consecutive body-text rows are folded into a single - # ``[N BODY LINES]`` placeholder row. The LLM therefore sees only: - # * heading CANDIDATES (integer id, real heading text), and - # * PLACEHOLDER rows (id with "-" or a single collapsed id, heading - # "[N BODY LINES]", level "-") that carry positional / section-bulk - # signals. + # ``[N BODY LINES]`` placeholder row. The LLM sees only two columns + # (``id`` and ``heading``) — no preliminary level estimate is provided, + # so the model assigns levels purely from structural/semantic analysis. temperature = 0 top_p = 0.01 max_depth = kwargs["paras"]["max_depth"] @@ -354,38 +352,44 @@ def build_prompt(task, texts, query, **kwargs): if toc_context: toc_section = f""" - ***Important Reference: Table of Contents (TOC)*** - The following is the document's table of contents with predefined levels. - Use it as a prior when assigning levels to CANDIDATE rows: + ***CONFIRMED Structure: Table of Contents (TOC)*** ''' {toc_context} ''' - - If a candidate's heading matches a TOC entry, use the TOC's predefined level. - - If a candidate appears to be a sub-section of a TOC entry, assign a deeper level. - - If a candidate does NOT appear in the TOC, it can ONLY be either body text - (level = -1) or a sub-section deeper than the nearest TOC heading above it. + RULES for using the TOC: + 1. MUST TRUST the TOC levels as ground truth. If a candidate heading matches or + closely corresponds to a TOC entry, you MUST assign it the same level as + the TOC entry. Do NOT override or re-interpret the TOC's level assignment. + 2. Candidates that appear between two TOC entries should be assigned a level + DEEPER than the TOC entry they fall under (they are sub-sections not listed + in the TOC). + 3. A candidate that does NOT correspond to any TOC entry can only be: + - Body text (level = -1), OR + - A sub-section with a level deeper than its nearest TOC heading above it. + 4. The TOC provides the SKELETON of the document. Your job is to fill in the + gaps for candidates not covered by the TOC, while strictly preserving the + TOC's structure. """ else: toc_section = "" prompt = f""" You are a document structure auditing expert. The input you receive is a - COMPACT skeleton of a document. Body-text lines have already been collapsed for - you so that every row is one of two kinds: + COMPACT skeleton of a document. Body-text lines have already been collapsed + for you so that every row is one of two kinds: - 1) HEADING CANDIDATE — ``id`` is an integer. ``heading`` is the candidate text. - ``level`` is a preliminary estimate: a positive integer (1 = shallowest, deeper = larger) - or the string "Not Sure" (undetermined). These rows — and ONLY these — are the ones you must evaluate. + 1) HEADING CANDIDATE — ``id`` is an integer, ``heading`` is the candidate + text. These — and ONLY these — are the rows you can evaluate. - 2) PLACEHOLDER — ``id`` is ALWAYS a hyphenated range "start-end" (for a - single-line it is "N-N", e.g. "56-56"); ``heading`` is "[N BODY LINES]" - where N is the number of body lines folded here; ``level`` is the literal "-". + 2) PLACEHOLDER — ``id`` is ALWAYS a range "start-end" (for a single-line + it is "N-N", e.g. "56-56"); ``heading`` is "[N BODY LINES]" + where N is the number of body lines folded here. Placeholders are positional markers that tell you how many body lines sit between adjacent candidates. Use them as context ONLY. - Data to be adjusted: + Data to be evaluated: ''' {texts} ''' @@ -394,10 +398,8 @@ def build_prompt(task, texts, query, **kwargs): ***Hard rules about placeholders*** - Placeholders are NEVER candidates. Do not output them. - - Every ``id`` in your output MUST be a single integer; never emit an id - containing a hyphen ("-"). Never emit the level string "-". - - Use N in ``[N BODY LINES]`` as a "section bulk" signal when applying - the rules below (Rule 6 in particular). + - Every ``id`` in your output MUST be a single integer; never emit an id containing a hyphen. + - Use N in ``[N BODY LINES]`` as a "section bulk" signal when applying rules below (Rule 2 in particular). ***Process in TWO steps:*** @@ -407,15 +409,14 @@ def build_prompt(task, texts, query, **kwargs): - Decimal numbering: "1.", "1.1", "1.1.1" → depth increases with dot count - Enumeration styles: "一、" "（一）" "1、" "①" "1 " → shallower to deeper with increasing numbers - Chapter/section keywords: "Chapter X", "Part X", "第X章", "第X节" - - Upper case / lower case differences in candidate headings - Clear semantic granularities or groups of themes + - Upper case / lower case differences Rank these patterns from shallowest to deepest to form a pattern → level mapping. Placeholder rows MUST NOT influence this scan. **STEP 2 — Assign a level to every candidate (rules in priority order)** - A candidate whose preliminary ``level`` is "Not Sure" or any positive - integer is **always** open to revision. Pure body text has already been - folded into placeholders, but a candidate **can still be** demoted to level = -1. + Your task is to determine each candidate's heading level from scratch + based on its text, context, and the patterns discovered in STEP 1. Rule 0 — Global consistency: Candidates sharing the same structural pattern or semantic granularity SHOULD receive the @@ -423,7 +424,7 @@ def build_prompt(task, texts, query, **kwargs): shares one level; every "X.Y.Z" shares a different, deeper level.) Rule 1 — Parent-child continuity and no level skipping: - A heading, compared to candidates before it, may stay at the same level, + A heading, compared to candidates before it, may stay at the same level, or go ONE level deeper than its nearest valid ancestor heading. However, jumps such as level 1 → level 3 are **always invalid**. @@ -431,19 +432,17 @@ def build_prompt(task, texts, query, **kwargs): A candidate WITHOUT any structural/numbering marker can still be a heading, but ONLY when ALL of the following hold: a) The text is short and title-like — no sentence-ending punctuation. - b) It is NOT a broken fragment that continues into the next row + b) It is NOT a broken fragment that continues into the next row. c) In the input sequence it is IMMEDIATELY followed by a placeholder ``[N BODY LINES]``, or by another candidate with finer granularity. This is the "section bulk" signal — the row introduces a body block or a subsection group. - When Rule-2 is satisfied, pick a level consistent with Rule 1 + When Rule-2 is satisfied, pick a level consistent with Rule 1. Rule 3 — Body text demotion (candidate → -1): - Demote a candidate to level = -1 when it clearly does NOT serve as a section title. - In compact input, the strongest demotion cues are: - - Two CANDIDATE rows appear adjacent with NO placeholder between them - - The text contains equations and math symbols such as + = - × ÷. - - The text is exactly "Figure/Image", demote it to level = -1. - - The text is an isolated broken phrase, fragment, data value, or caption-like snippet (e.g. "Table 3-2", "Figure 4" + Demote a candidate to level = -1 when it clearly does NOT serve as a + section title. Strongest demotion cues are: + - Two CANDIDATE rows appear adjacent with NO placeholder between them. + - The text is an isolated broken phrase, fragment, data value, or caption-like snippet (e.g. "Table 3-2", "Figure 4"). Rule 4 — Normalise to start at level 1: The shallowest (the most coarse granularity) heading found MUST be assigned level 1. From ccbe4ba61b280f7593c44dd8e339eebc0d36a434 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 28 May 2026 11:59:50 +0800 Subject: [PATCH 10/11] refactor: remove page numbers from assets and optimize heading detection logic for markdown parsing --- .../tests/contract/test_documents_contract.py | 12 +- .../document_parser/assets/inline_asset.py | 5 +- .../document_parser/formats/atlas/parser.py | 1 - .../formats/markdown/image_asset.py | 7 - .../formats/markdown/parse_state.py | 31 +- .../formats/markdown/parser.py | 157 +++--- .../formats/markdown/table_asset.py | 4 - .../document_parser/formats/pdf/parser.py | 222 ++++----- .../formats/pdf/shard_merger.py | 68 +-- .../structure/heading_candidates.py | 36 +- .../structure/heading_llm_executor.py | 466 ++---------------- .../structure/layout_parser.py | 95 ++-- .../tables/table_asset_writer.py | 2 - 13 files changed, 309 insertions(+), 797 deletions(-) diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index d0b401839..6016aaeb6 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -622,7 +622,7 @@ async def test_should_list_current_document_chunks_by_document_id( "chunk_type": "text", "content": "First chunk content", "source_chunk_path": "Chapter 1/Intro", - "metadata": {"summary": "Intro", "page_nums": [1]}, + "metadata": {"summary": "Intro", "page_nums": []}, }, { "id": second_chunk_id, @@ -631,7 +631,7 @@ async def test_should_list_current_document_chunks_by_document_id( "content": "| A | B |", "source_chunk_path": "Chapter 1/Table", "file_path": "tables/table-1.html", - "metadata": {"summary": "Table", "page_nums": [2]}, + "metadata": {"summary": "Table", "page_nums": []}, }, ], ) @@ -666,7 +666,7 @@ async def test_should_list_current_document_chunks_by_document_id( "source_chunk_path": "Chapter 1/Intro", "file_path": None, "sort_order": 0, - "metadata": {"summary": "Intro", "page_nums": [1]}, + "metadata": {"summary": "Intro", "page_nums": []}, "created_at": chunks[0]["created_at"], } ] @@ -720,7 +720,7 @@ async def test_should_return_one_document_chunk_by_document_chunk_id( "content": "Figure summary", "source_chunk_path": "Chapter 1/Figure", "file_path": "images/figure-1.png", - "metadata": {"summary": "Figure", "page_nums": [3]}, + "metadata": {"summary": "Figure", "page_nums": []}, } ], ) @@ -745,7 +745,7 @@ async def test_should_return_one_document_chunk_by_document_chunk_id( assert chunk["section_path"] == "Chapter 1" assert chunk["source_chunk_path"] == "Chapter 1/Figure" assert chunk["file_path"] == "images/figure-1.png" - assert chunk["metadata"] == {"summary": "Figure", "page_nums": [3]} + assert chunk["metadata"] == {"summary": "Figure", "page_nums": []} assert chunk["created_at"] @@ -768,7 +768,7 @@ async def test_should_return_not_found_when_requesting_a_missing_document_chunk( "chunk_type": "text", "content": "First chunk content", "source_chunk_path": "Chapter 1/Intro", - "metadata": {"summary": "Intro", "page_nums": [1]}, + "metadata": {"summary": "Intro", "page_nums": []}, } ], ) diff --git a/apps/worker/app/services/document_parser/assets/inline_asset.py b/apps/worker/app/services/document_parser/assets/inline_asset.py index 9e792ee1d..1d67eae09 100644 --- a/apps/worker/app/services/document_parser/assets/inline_asset.py +++ b/apps/worker/app/services/document_parser/assets/inline_asset.py @@ -10,7 +10,6 @@ def build_image_asset_row( summary: str, know_id: str, addtime: str, - page_nums: str = "", ) -> ParsedRow: return ParsedRow( content=content, @@ -22,7 +21,6 @@ def build_image_asset_row( tokens="", connectto="", addtime=addtime, - page_nums=page_nums, ) @@ -34,7 +32,6 @@ def build_table_asset_row( keywords: str, know_id: str, addtime: str, - page_nums: str = "", ) -> ParsedRow: return ParsedRow( content=content, @@ -46,5 +43,5 @@ def build_table_asset_row( tokens="", connectto="", addtime=addtime, - page_nums=page_nums, ) + diff --git a/apps/worker/app/services/document_parser/formats/atlas/parser.py b/apps/worker/app/services/document_parser/formats/atlas/parser.py index a3210d67e..ddac6a6d0 100644 --- a/apps/worker/app/services/document_parser/formats/atlas/parser.py +++ b/apps/worker/app/services/document_parser/formats/atlas/parser.py @@ -444,7 +444,6 @@ def _vlm_task(page_num, img_name): know_id=know_id, addtime=time_stamp, tokens=tokens, - page_nums=str(page_num), ) ) diff --git a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py index b71ace46a..ed5f51abf 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py +++ b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py @@ -39,7 +39,6 @@ class MarkdownImageAssetRequest: last_context: str image_summary: str | None timestamp: str - current_page_number: int seen_images: dict[str, dict[str, str]] summary_image: bool row_index: int @@ -65,7 +64,6 @@ def build_markdown_image_asset( source_path=source_path, cache_entry=request.seen_images[image_binary_hash], timestamp=request.timestamp, - current_page_number=request.current_page_number, ) relative_image_path = f"images/{request.image_name}{image_suffix}" @@ -91,7 +89,6 @@ def build_markdown_image_asset( summary=image_summary_field, know_id=image_know_id, timestamp=request.timestamp, - current_page_number=request.current_page_number, ) cache_entry = { "relative_img_path": relative_image_path, @@ -166,7 +163,6 @@ def _build_duplicate_image_asset( source_path: Path, cache_entry: dict[str, str], timestamp: str, - current_page_number: int, ) -> MarkdownImageAsset: row_values = _build_image_row_values( content=cache_entry["img_content"], @@ -174,7 +170,6 @@ def _build_duplicate_image_asset( summary=cache_entry["img_summary_field"], know_id=cache_entry["temp_uid"], timestamp=timestamp, - current_page_number=current_page_number, ) try: source_path.unlink() @@ -205,7 +200,6 @@ def _build_image_row_values( summary: str, know_id: str, timestamp: str, - current_page_number: int, ) -> ParserRowValues: image_row = build_image_asset_row( content=content, @@ -213,7 +207,6 @@ def _build_image_row_values( summary=summary, know_id=know_id, addtime=timestamp, - page_nums=str(current_page_number) if current_page_number > 0 else "", ) return cast(ParserRowValues, image_row.to_list()) diff --git a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py index b7c6da508..95e5088e8 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py @@ -1,6 +1,6 @@ from __future__ import annotations -import re + from collections.abc import Callable from dataclasses import dataclass, field from typing import Any @@ -35,8 +35,7 @@ class MarkdownParseState: inner_paths: list[str] = field(default_factory=list) error_line_numbers: list[int] = field(default_factory=list) table_lines: list[str] = field(default_factory=list) - current_page_number: int = 0 - chunk_pages: set[int] = field(default_factory=set) + base_level: int | None = None path: str = "" path_counter: dict[str, int] = field(default_factory=dict) @@ -50,45 +49,38 @@ def __post_init__(self) -> None: self.path = self.relative_root def record_page_marker(self, line: str) -> bool: + """Detect and skip HTML comment lines (page markers, slide markers). + + Page number tracking has been removed; PAGE MEMORY will provide + accurate page numbers in a future release. + """ if "" not in line: return False if "page" not in line and "Slide number" not in line: return False - - page_match = re.search(r"page\s+(\d+)", line) - if page_match: - self.current_page_number = int(page_match.group(1)) - else: - self.current_page_number += 1 - self.chunk_pages.add(self.current_page_number) return True def flush_current_content(self) -> None: - page_numbers = self._format_chunk_pages() self.rows = self.row_updater( self.rows, self.content_items, self.path, self.llm_parameters, self.timestamp, - page_numbers, + "", 1500, True, ) self.content_items = [] - self.chunk_pages = set() - if self.current_page_number > 0: - self.chunk_pages.add(self.current_page_number) def flush_placeholder_chunk(self) -> None: - page_numbers = self._format_chunk_pages() self.rows = self.row_updater( self.rows, [], self.path, self.llm_parameters, self.timestamp, - page_numbers, + "", 1500, True, ) @@ -135,8 +127,6 @@ def append_content_item(self, item: str) -> None: def append_plain_text(self, text: str) -> None: self.content_items.append(text.strip()) - if self.current_page_number > 0: - self.chunk_pages.add(self.current_page_number) def append_row(self, row: ParserRowValues) -> None: self.rows.append(row) @@ -181,5 +171,4 @@ def to_dataframe(self) -> pd.DataFrame: ) return process_dup_paths_df(rows_builder.to_dataframe()) - def _format_chunk_pages(self) -> str: - return ",".join(str(page) for page in sorted(self.chunk_pages)) + diff --git a/apps/worker/app/services/document_parser/formats/markdown/parser.py b/apps/worker/app/services/document_parser/formats/markdown/parser.py index 39b4e7483..614bab23a 100755 --- a/apps/worker/app/services/document_parser/formats/markdown/parser.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parser.py @@ -214,69 +214,97 @@ def parse_md( base_llm_paras=None, relative_root=None, toc_hierarchies=None, + lines_with_heading=None, ): - if md_lines is None and file_path is not None: - from app.services.common.file_loading import is_remote, load_file_bytes + if lines_with_heading is not None: + # ── Phase A bypass ── + # Caller (e.g. oversized PDF shard-first path) already ran per-shard + # heading prediction and passed in the merged lines_with_heading. + # Skip TOC detection and heading prediction entirely. + logger.info( + f"📌 Using pre-identified headings ({len(lines_with_heading)} lines), " + f"skipping TOC detection and heading prediction" + ) + else: + # ── Phase A: TOC detection + heading prediction ── + if md_lines is None and file_path is not None: + from app.services.common.file_loading import is_remote, load_file_bytes + + if is_remote(file_path): + file_bytes = load_file_bytes(file_path) + md_content = file_bytes.decode("utf-8") + md_lines = md_content.splitlines() + else: + with open(file_path, "r", encoding="utf-8") as file: + md_lines = file.readlines() - if is_remote(file_path): - file_bytes = load_file_bytes(file_path) - md_content = file_bytes.decode("utf-8") - md_lines = md_content.splitlines() - else: - with open(file_path, "r", encoding="utf-8") as file: - md_lines = file.readlines() + md_lines = [line.strip() for line in md_lines if line.strip() != ""] - md_lines = [line.strip() for line in md_lines if line.strip() != ""] + # Preprocess: merge multi-line HTML tables into single lines + md_lines = merge_html_tables(md_lines) - # Preprocess: merge multi-line HTML tables into single lines - md_lines = merge_html_tables(md_lines) + # Detect TOC using async LLM-based detection + toc_model_name = ( + base_llm_paras.get("model_name", settings.NORMOL_MODEL) + if base_llm_paras + else settings.NORMOL_MODEL + ) + hierarchy_model_name = ( + (base_llm_paras.get("hierarchy_model_name") or toc_model_name) + if base_llm_paras + else (settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL) + ) - # Detect TOC using async LLM-based detection - toc_model_name = ( - base_llm_paras.get("model_name", settings.NORMOL_MODEL) - if base_llm_paras - else settings.NORMOL_MODEL - ) - hierarchy_model_name = ( - (base_llm_paras.get("hierarchy_model_name") or toc_model_name) - if base_llm_paras - else (settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL) - ) + if toc_hierarchies is not None: + # Pre-detected TOC from upstream (e.g. DOC_AGENT VLM-based extraction). + # Skip row-based detection entirely — TOC pages have already been + # physically stripped from the PDF, so no TOC rows exist in md_lines. + logger.info( + f"📌 Using pre-detected TOC hierarchies " + f"({len(toc_hierarchies)} regions), " + f"skipping detect_tocs_in_texts" + ) + else: + with stage_timer( + "md.detect_toc", line_count=len(md_lines), model_name=toc_model_name + ): + toc_hierarchies, md_lines = detect_tocs_in_texts( + md_lines, + model_name=toc_model_name, + hierarchy_model_name=hierarchy_model_name, + ) - if toc_hierarchies is not None: - # Pre-detected TOC from upstream (e.g. DOC_AGENT VLM-based extraction). - # Skip row-based detection entirely — TOC pages have already been - # physically stripped from the PDF, so no TOC rows exist in md_lines. - logger.info( - f"📌 Using pre-detected TOC hierarchies " - f"({len(toc_hierarchies)} regions), " - f"skipping detect_tocs_in_texts" - ) - else: - # DEPRECATED: Row-based TOC detection via detect_tocs_in_texts. - # This path is being replaced by upstream DOC_AGENT VLM-based TOC - # detection (extract_toc_with_boundaries). It is kept as a fallback - # for: - # 1. Standard-path PDFs (≤MAX_PDF_PAGE_LIMIT pages, no DOC_AGENT) - # 2. Non-PDF sources (pure markdown files) - # Once DOC_AGENT TOC detection covers all paths, this block should - # be removed. + # Save toc_hierarchies.json to output_dir (will be included in final zip package) + if toc_hierarchies: + toc_json_path = os.path.join(output_dir, "toc_hierarchies.json") + with open(toc_json_path, "w", encoding="utf-8") as f: + json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) + logger.info(f"Saved TOC hierarchies to {toc_json_path}") + + # Find layout.json path + layout_json_path = os.path.join(output_dir, "layout.json") + if not os.path.exists(layout_json_path): + layout_json_path = None + logger.debug("layout.json not found, META features will not be added") + + # estimate hierarchies with toc_hierarchies context with stage_timer( - "md.detect_toc", line_count=len(md_lines), model_name=toc_model_name + "md.predict_headings", + line_count=len(md_lines), + smart_parse=base_llm_paras["smart_title_parse"], + model_name=hierarchy_model_name, ): - toc_hierarchies, md_lines = detect_tocs_in_texts( + lines_with_heading = eval_md_headings( md_lines, - model_name=toc_model_name, - hierarchy_model_name=hierarchy_model_name, + source_type, + toc_hierarchies=toc_hierarchies, + smart_parse=base_llm_paras["smart_title_parse"], + model_name=hierarchy_model_name, + output_dir=output_dir, + layout_json_path=layout_json_path, ) - # Save toc_hierarchies.json to output_dir (will be included in final zip package) - if toc_hierarchies: - toc_json_path = os.path.join(output_dir, "toc_hierarchies.json") - with open(toc_json_path, "w", encoding="utf-8") as f: - json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) - logger.info(f"Saved TOC hierarchies to {toc_json_path}") - + # ── Phase B: MarkdownParseState traversal ── # Clean old artifacts to prevent accumulation across debug runs. # In production each job uses a fresh workspace so rmtree never triggers. tb_dir = os.path.join(output_dir, "tables") @@ -301,29 +329,6 @@ def parse_md( row_updater=update_df_list, ) - # Find layout.json path - layout_json_path = os.path.join(output_dir, "layout.json") - if not os.path.exists(layout_json_path): - layout_json_path = None - logger.debug("layout.json not found, META features will not be added") - - # estimate hierarchies with toc_hierarchies context - with stage_timer( - "md.predict_headings", - line_count=len(md_lines), - smart_parse=base_llm_paras["smart_title_parse"], - model_name=hierarchy_model_name, - ): - lines_with_heading = eval_md_headings( - md_lines, - source_type, - toc_hierarchies=toc_hierarchies, - smart_parse=base_llm_paras["smart_title_parse"], - model_name=hierarchy_model_name, - output_dir=output_dir, - layout_json_path=layout_json_path, - ) - logger.debug("Parsing md data... total_lines={}", len(lines_with_heading)) for i, line in enumerate(lines_with_heading): if parser_state.record_page_marker(line): @@ -366,7 +371,6 @@ def parse_md( last_context=last_context, image_summary=img_summary, timestamp=parser_state.timestamp, - current_page_number=parser_state.current_page_number, seen_images=parser_state.seen_images, summary_image=bool(base_llm_paras["summary_image"]), row_index=len(parser_state.rows), @@ -426,7 +430,6 @@ def parse_md( table_dir=tb_dir, table_count=parser_state.table_count, timestamp=parser_state.timestamp, - current_page_number=parser_state.current_page_number, summary_table=bool(base_llm_paras["summary_table"]), row_index=len(parser_state.rows), ) diff --git a/apps/worker/app/services/document_parser/formats/markdown/table_asset.py b/apps/worker/app/services/document_parser/formats/markdown/table_asset.py index 19f795f4f..a489e036d 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/table_asset.py +++ b/apps/worker/app/services/document_parser/formats/markdown/table_asset.py @@ -32,7 +32,6 @@ class MarkdownTableAssetRequest: table_dir: str table_count: int timestamp: str - current_page_number: int summary_table: bool row_index: int @@ -62,9 +61,6 @@ def build_markdown_table_asset( keywords="", know_id=gen_str_codes((request.table_html + str(request.table_count))), addtime=request.timestamp, - page_nums=str(request.current_page_number) - if request.current_page_number > 0 - else "", ) deferred_task = None diff --git a/apps/worker/app/services/document_parser/formats/pdf/parser.py b/apps/worker/app/services/document_parser/formats/pdf/parser.py index e40850a6a..c5baf58f3 100755 --- a/apps/worker/app/services/document_parser/formats/pdf/parser.py +++ b/apps/worker/app/services/document_parser/formats/pdf/parser.py @@ -1,7 +1,5 @@ # pyright: reportArgumentType=false -import json import os -import re from app.services.document_parser.formats.markdown.parser import parse_md from app.services.document_parser.providers.mineru.pdf_service import parse_via_full @@ -11,106 +9,6 @@ from shared.core.config import settings -def _inject_page_markers(output_dir: str) -> None: - """Inject markers into full.md using layout.json page info. - - Reads layout.json to find the first text content of each page, - then searches for that text in full.md and inserts a page marker above it. - - If layout.json is not available (e.g. fast path without MinerU), - this function does nothing gracefully. - """ - layout_path = os.path.join(output_dir, "layout.json") - md_path = os.path.join(output_dir, "full.md") - - if not os.path.exists(layout_path) or not os.path.exists(md_path): - logger.debug("layout.json or full.md not found, skipping page marker injection") - return - - try: - with open(layout_path, "r", encoding="utf-8") as f: - layout_data = json.load(f) - except (json.JSONDecodeError, IOError) as e: - logger.warning(f"Failed to read layout.json: {e}") - return - - pdf_info = layout_data.get("pdf_info", []) - if not pdf_info: - return - - with open(md_path, "r", encoding="utf-8") as f: - md_lines = f.readlines() - - # Build anchor map: {normalized_text: page_number (1-based)} - # Use the first text span of each page's first para_block as anchor - anchors = [] # list of (anchor_text, page_num) - for page in pdf_info: - page_idx = page.get("page_idx", 0) - page_num = page_idx + 1 # 1-based page number - - # Find the first non-empty text content in this page - anchor_text = None - for block in page.get("para_blocks", []): - for line in block.get("lines", []): - for span in line.get("spans", []): - content = span.get("content", "").strip() - if content and len(content) >= 3: # skip very short anchors - anchor_text = content - break - if anchor_text: - break - if anchor_text: - break - - if anchor_text: - anchors.append((anchor_text, page_num)) - - if not anchors: - logger.debug( - "No anchor texts found in layout.json, skipping page marker injection" - ) - return - - # Match anchors against md_lines and insert markers - # Process from end to start so line indices don't shift - insertions = [] # list of (line_index, page_num) - used_lines = set() - - for anchor_text, page_num in anchors: - # Normalize anchor for matching - anchor_norm = re.sub(r"\s+", " ", anchor_text).strip() - if len(anchor_norm) < 3: - continue - - # Search for anchor in md_lines (use first 50 chars for substring match) - search_key = anchor_norm[:50] - for i, line in enumerate(md_lines): - if i in used_lines: - continue - line_norm = re.sub(r"^#+\s*", "", line.strip()) - line_norm = re.sub(r"\s+", " ", line_norm).strip() - if search_key in line_norm: - insertions.append((i, page_num)) - used_lines.add(i) - break - - if not insertions: - logger.debug("No page marker matches found, skipping injection") - return - - # Sort by line index descending to insert from bottom to top - insertions.sort(key=lambda x: x[0], reverse=True) - for line_idx, page_num in insertions: - md_lines.insert(line_idx, f"\n") - - # Write back - with open(md_path, "w", encoding="utf-8") as f: - f.writelines(md_lines) - - logger.info(f"Injected {len(insertions)} page markers into full.md") - - - def parse_pdfs( pdf_path, filename, @@ -147,7 +45,6 @@ def parse_pdfs( logger.info(f"📄 Standard MinerU parse for {filename} [route={route}]") with stage_timer("pdf.extract.standard", filename=filename): parse_via_full(pdf_path, filename, output_dir, s3_key=s3_key) - _inject_page_markers(output_dir) logger.info("✅ PDF parsing step 1 complete: text extracted") @@ -165,12 +62,25 @@ def _parse_oversized_pdf( pdf_path, filename, output_dir, base_llm_paras, profile=None, relative_root=None, s3_key=None, ): - """Handle PDFs exceeding MinerU's page limit via doc_agent shard-and-stitch.""" + """Handle PDFs exceeding MinerU's page limit via shard-first hierarchy. + + Pipeline: + 1. DOC_AGENT → shard plan + TOC + 2. bin_pack → merged shards + 3. split_pdf (exclude TOC pages) + 4. MinerU per shard (parallel) + 5. **Per-shard heading prediction** (parallel) ← NEW + 6. Merge lines_with_heading + images + 7. parse_md Phase B (skip TOC detection + heading prediction) + """ from concurrent.futures import ThreadPoolExecutor, as_completed + from dataclasses import dataclass - from app.services.document_parser.formats.pdf.shard_merger import ( - merge_shard_outputs, + from app.services.document_parser.formats.markdown.parser import ( + eval_md_headings, + merge_html_tables, ) + from app.services.document_parser.formats.pdf.shard_merger import merge_images from app.services.document_parser.formats.pdf.shard_splitter import ( bin_pack_shards, run_doc_agent, @@ -246,26 +156,102 @@ def _parse_single_shard(shard_idx, shard_pdf): idx = futures[future] shard_output_dirs[idx] = future.result() - # 6. Merge all shard outputs into main output_dir - with stage_timer("pdf.merge_shards", filename=filename): - merge_shard_outputs(shard_output_dirs, merged_shards, output_dir) + # 6. Per-shard heading prediction (parallel) + @dataclass + class ShardHeadingResult: + shard_index: int + lines_with_heading: list[str] + heading_count: int + + smart_parse = base_llm_paras.get("smart_title_parse", True) + hierarchy_model_name = ( + base_llm_paras.get("hierarchy_model_name") + or base_llm_paras.get("model_name", settings.NORMOL_MODEL) + ) + + def _predict_shard_headings(shard_idx: int, shard_out_dir: str) -> ShardHeadingResult: + """Run full heading prediction pipeline on a single shard's full.md.""" + md_path = os.path.join(shard_out_dir, "full.md") + if not os.path.exists(md_path): + logger.warning(f"shard_{shard_idx}: full.md not found, returning empty") + return ShardHeadingResult(shard_index=shard_idx, lines_with_heading=[], heading_count=0) + + with open(md_path, "r", encoding="utf-8") as f: + md_lines = f.readlines() + md_lines = [line.strip() for line in md_lines if line.strip() != ""] + md_lines = merge_html_tables(md_lines) + + # TOC context: first TOC shared by all shards; subsequent TOCs assigned + # by page boundary. For simplicity, all TOCs are passed since pred_titles + # only matches headings actually present in this shard's content. + shard_toc = toc_hierarchies + + lines_with_heading = eval_md_headings( + md_lines, + source_type="md", + toc_hierarchies=shard_toc, + smart_parse=smart_parse, + model_name=hierarchy_model_name, + output_dir=shard_out_dir, + layout_json_path=( + os.path.join(shard_out_dir, "layout.json") + if os.path.exists(os.path.join(shard_out_dir, "layout.json")) + else None + ), + ) + + heading_count = sum(1 for line in lines_with_heading if line.startswith("#")) + logger.info( + f" ✅ shard_{shard_idx}: {heading_count} headings identified " + f"from {len(lines_with_heading)} lines" + ) + return ShardHeadingResult( + shard_index=shard_idx, + lines_with_heading=lines_with_heading, + heading_count=heading_count, + ) - # 7. Inject page markers (uses merged layout.json with corrected page_idx) - # Note: page markers may be inaccurate when TOC pages are excluded, but - # this only affects chunk metadata (page_nums), not heading hierarchy. - _inject_page_markers(output_dir) + shard_heading_results: list[ShardHeadingResult | None] = [None] * len(shard_output_dirs) - logger.info("✅ Oversized PDF shard-and-stitch complete, entering parse_md") + with stage_timer( + "pdf.shard_headings", filename=filename, shard_count=len(shard_output_dirs) + ): + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = { + executor.submit(_predict_shard_headings, i, shard_dir): i + for i, shard_dir in enumerate(shard_output_dirs) + if shard_dir is not None + } + for future in as_completed(futures): + idx = futures[future] + shard_heading_results[idx] = future.result() + + # 7. Merge: concatenate lines_with_heading (in shard order) + merge images + all_lines_with_heading: list[str] = [] + total_headings = 0 + for result in shard_heading_results: + if result is not None: + all_lines_with_heading.extend(result.lines_with_heading) + total_headings += result.heading_count + + logger.info( + f"📎 Merged {len(shard_heading_results)} shards: " + f"{len(all_lines_with_heading)} lines, {total_headings} headings" + ) - # 8. Standard parse_md — pass DOC_AGENT TOC hierarchies to skip - # row-based TOC detection and enable hard-constraint heading assignment + with stage_timer("pdf.merge_images", filename=filename): + merge_images(shard_output_dirs, output_dir) + + logger.info("✅ Shard-first hierarchy complete, entering parse_md Phase B") + + # 8. parse_md Phase B only (skip TOC detection + heading prediction) with stage_timer("pdf.parse_md", filename=filename): return parse_md( output_dir, source_type="md", - file_path=os.path.join(output_dir, "full.md"), base_llm_paras=base_llm_paras, relative_root=relative_root, - toc_hierarchies=toc_hierarchies, + lines_with_heading=all_lines_with_heading, ) + diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py index f1b78e623..55afa3ca0 100644 --- a/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py +++ b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py @@ -1,79 +1,21 @@ -"""Merge MinerU outputs from multiple shards into a single unified output.""" +"""Merge MinerU image outputs from multiple shards into a single unified output.""" from __future__ import annotations -import json import os import shutil -from typing import TYPE_CHECKING from loguru import logger -if TYPE_CHECKING: - from app.services.document_parser.formats.pdf.shard_splitter import MergedShard - -def merge_shard_outputs( - shard_output_dirs: list[str], - shards: list[MergedShard], - target_output_dir: str, -) -> None: - """Merge MinerU outputs from all shards into target_output_dir. - - After this call target_output_dir contains: - - full.md: concatenated markdown (UUID image refs, no conflicts) - - layout.json: merged page array with corrected page_idx - - images/: all images from all shards - """ - _merge_full_md(shard_output_dirs, target_output_dir) - _merge_layout_json(shard_output_dirs, shards, target_output_dir) - _merge_images(shard_output_dirs, target_output_dir) - - -def _merge_full_md(shard_dirs: list[str], target_dir: str) -> None: - target_path = os.path.join(target_dir, "full.md") - with open(target_path, "w", encoding="utf-8") as out: - for i, shard_dir in enumerate(shard_dirs): - md_path = os.path.join(shard_dir, "full.md") - if not os.path.exists(md_path): - logger.warning(f"shard {i}: full.md not found at {md_path}") - continue - with open(md_path, "r", encoding="utf-8") as f: - content = f.read() - if i > 0: - out.write("\n\n") - out.write(content) - logger.info(f"Merged {len(shard_dirs)} full.md files → {target_path}") - - -def _merge_layout_json( - shard_dirs: list[str], - shards: list[MergedShard], - target_dir: str, -) -> None: - merged_pages: list[dict] = [] - for shard, shard_dir in zip(shards, shard_dirs): - layout_path = os.path.join(shard_dir, "layout.json") - if not os.path.exists(layout_path): - logger.warning(f"shard {shard.shard_index}: layout.json not found") - continue - with open(layout_path, "r", encoding="utf-8") as f: - layout_data = json.load(f) - for page in layout_data.get("pdf_info", []): - page["page_idx"] = page.get("page_idx", 0) + shard.page_offset - merged_pages.append(page) - - target_path = os.path.join(target_dir, "layout.json") - with open(target_path, "w", encoding="utf-8") as f: - json.dump({"pdf_info": merged_pages}, f, ensure_ascii=False) - logger.info(f"Merged layout.json: {len(merged_pages)} pages → {target_path}") - - -def _merge_images(shard_dirs: list[str], target_dir: str) -> None: +def merge_images(shard_dirs: list[str], target_dir: str) -> None: + """Copy all images from shard images/ dirs into target_dir/images/.""" target_img_dir = os.path.join(target_dir, "images") os.makedirs(target_img_dir, exist_ok=True) total = 0 for shard_dir in shard_dirs: + if shard_dir is None: + continue img_dir = os.path.join(shard_dir, "images") if not os.path.isdir(img_dir): continue diff --git a/apps/worker/app/services/document_parser/structure/heading_candidates.py b/apps/worker/app/services/document_parser/structure/heading_candidates.py index f9d060349..b790a6cf7 100644 --- a/apps/worker/app/services/document_parser/structure/heading_candidates.py +++ b/apps/worker/app/services/document_parser/structure/heading_candidates.py @@ -15,11 +15,13 @@ HEADING_COLUMNS = Index(["id", "heading", "level", "reason"]) -def get_max_lvl(code_str: str): - match = re.search(r"\[([^]]+)]", code_str) - if not match: - return "Sure" +def get_max_lvl(code_str: str) -> int: + """Extract the maximum hierarchy depth from a POS trigger code string. + ``code_str`` is always ``str(pos_code)`` where *pos_code* is a list of + integers, so the ``[…]`` bracket match is guaranteed. + """ + match = re.search(r"\[([^]]+)]", code_str) nums = [int(item.strip()) for item in match.group(1).split(",")] max_value = int(max(nums)) return max_value if max_value > 1 else -2 @@ -86,7 +88,7 @@ def judge_by_conditions( return pos_triggered_code -def remove_by_conditions(text, include_punc: bool = False): +def remove_by_conditions(text): neg_conditions = [ r"^\d{3,}", r"(?i)(^https?://\S+|^www\.\S+|^P\.S|^\b\d{0,2}\s*(?:a\.m|p\.m)\b)", @@ -101,7 +103,7 @@ def remove_by_conditions(text, include_punc: bool = False): r"|left|right|begin|end|overline|underline|hat|vec|tilde)\b" r")" ), - r"^0\.\d+[\u4e00-\u9fa5A-Za-z\S]*", + r"^0\.\d+\S*", r"^\d*\.\d+$", r"[。！；].+", ( @@ -113,16 +115,15 @@ def remove_by_conditions(text, include_punc: bool = False): r"|Hz|kHz|MHz|GHz" r"|mol|mL|dL|dB|Nm|kN|MN|kW|MW|GW|hp|rpm|cc|cal|kcal)\b" ), + r"[.,;，。；]$", ] neg_triggered_code = [] for regex in neg_conditions: neg_triggered_code.append(1 if re.search(regex, text) else 0) - - if include_punc: - neg_triggered_code.append(1 if re.search(r"[.,;，。；]$", text) else 0) - else: - neg_triggered_code.append(0) + + MAX_HEADING_TOKENS = 10 + neg_triggered_code.append(1 if count_cn_en(text) > MAX_HEADING_TOKENS else 0) return neg_triggered_code @@ -136,10 +137,13 @@ def md_heading_match(line, as_is: bool = True): return (line, level) if as_is else (line.lstrip("#").strip(), level) +# Pre-compute zero-filled code arrays so non-heading lines get correct-width reason strings. +_ZERO_POS_CODE = judge_by_conditions("") +_ZERO_NEG_CODE = remove_by_conditions("") + + def filter_markdown_headings( md_lines: list[str], - num_pos: int = 17, - num_neg: int = 7, layout_json_path: str | None = None, ) -> pd.DataFrame: meta_ctx = None @@ -159,9 +163,7 @@ def filter_markdown_headings( if _is_non_heading_markdown_line(line): est_level = -1 - zero_pos_code = [0] * num_pos - zero_neg_code = [0] * num_neg - reason = f"POS {zero_pos_code} NEG {zero_neg_code}" + reason = f"POS {_ZERO_POS_CODE} NEG {_ZERO_NEG_CODE}" if meta_ctx: reason += " META [0, 0, 0]" line = "Figure/Image" @@ -327,7 +329,7 @@ def _is_bold_docx_paragraph(paragraph: Any): def _judge_negative_headings(df: pd.DataFrame) -> pd.DataFrame: for index, row in df.iterrows(): - neg_code = remove_by_conditions(row["heading"], include_punc=True) + neg_code = remove_by_conditions(row["heading"]) if any(value > 0 for value in neg_code): current_code = str(df.loc[index, "reason"]) diff --git a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py index 9b984f796..bea5e8540 100644 --- a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py +++ b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py @@ -2,8 +2,6 @@ from __future__ import annotations import os -import re -from collections import Counter from collections.abc import Callable from typing import Any @@ -13,8 +11,6 @@ from app.services.document_parser.support.text_helpers import count_cn_en, truncate_text_by_tokens from loguru import logger -from shared.core.exceptions.domain_exceptions import WorkerHandlingException - PLACEHOLDER_REASON = "__PLACEHOLDER__" HierarchyJudge = Callable[..., list[dict[str, Any]]] @@ -22,328 +18,6 @@ SaveIntermediateCsv = Callable[[pd.DataFrame, str | None, str], None] -def build_level_mapping( - df: pd.DataFrame, origin_lvls: list[int], mode: str = "max" -) -> tuple[pd.DataFrame, dict[str, dict[str, Any]]]: - mapped_df = df.copy() - mapped_df["origin_level"] = origin_lvls - - mapping = mapped_df.groupby("reason")["level"].apply(list).to_dict() - - processed_mapping: dict[str, dict[str, Any]] = {} - for reason, lvls in mapping.items(): - positive_lvls = [lvl for lvl in lvls if lvl > -1] - counts = Counter(lvls) - - if not positive_lvls: - mapped_lvl = -1 - elif mode == "max": - mapped_lvl = max(positive_lvls) - elif mode == "freq": - mapped_lvl = counts.most_common(1)[0][0] - else: - raise WorkerHandlingException( - internal_message=f"wrong input mode: {mode}. Must be 'max' or 'freq'" - ) - - processed_mapping[reason] = { - "lvls": lvls, - "positive_lvls": positive_lvls, - "freqs": dict(counts), - "mapped_lvl": mapped_lvl, - } - return mapped_df, processed_mapping - - -def execute_level_mapping( - df: pd.DataFrame, mapping: dict[str, dict[str, Any]] -) -> pd.DataFrame: - def map_row(row: pd.Series) -> int: - reason = row["reason"] - if reason in mapping: - return int(mapping[reason]["mapped_lvl"]) - return int(row["level"]) - - mapped_df = df.copy() - origin_est_lvls = mapped_df["level"].tolist() - mapped_df["level"] = mapped_df.apply(map_row, axis=1) - mapped_df["origin_level"] = origin_est_lvls - return mapped_df - - -def extract_non_neg_code(reason_str: str) -> str: - """Extract the non-NEG code from a heading reason string.""" - if not reason_str or not isinstance(reason_str, str): - return "" - neg_match = re.search(r"\s*NEG\s*\[[^\]]*\]", reason_str) - if neg_match: - before_neg = reason_str[: neg_match.start()] - after_neg = reason_str[neg_match.end() :] - return (before_neg + after_neg).strip() - return reason_str.strip() - - -def build_non_neg_mapping(lvl_mapping: dict[str, dict[str, Any]]) -> dict[str, int]: - non_neg_levels: dict[str, list[int]] = {} - for reason, info in lvl_mapping.items(): - non_neg_code = extract_non_neg_code(reason) - mapped_lvl = int(info.get("mapped_lvl", -1)) - if non_neg_code: - non_neg_levels.setdefault(non_neg_code, []).append(mapped_lvl) - - non_neg_mapping: dict[str, int] = {} - for non_neg_code, levels in non_neg_levels.items(): - positive_levels = [lvl for lvl in levels if lvl > -1] - if positive_levels: - level_counts = Counter(positive_levels) - non_neg_mapping[non_neg_code] = level_counts.most_common(1)[0][0] - else: - non_neg_mapping[non_neg_code] = -1 - - return non_neg_mapping - - -def handle_unseen_codes( - df: pd.DataFrame, - level_dfs: list[pd.DataFrame], - lvl_mapping: dict[str, dict[str, Any]], - output_dir: str | None = None, - window_half_size: int = 10, - strategy: str = "double_mapping", -) -> dict[str, dict[str, Any]]: - """Extend first-chunk reason mapping to reason codes only seen in later chunks.""" - - def extract_reason_signature(reason: str) -> str: - return reason.strip() if reason else "" - - def has_neg_signal(reason_str: str) -> bool: - if not reason_str or not isinstance(reason_str, str): - return False - neg_match = re.search(r"NEG\s*\[([^\]]*)\]", reason_str) - if not neg_match: - return False - neg_content = neg_match.group(1) - try: - nums = [int(x.strip()) for x in neg_content.split(",") if x.strip()] - return any(x >= 1 for x in nums) - except Exception: - return False - - def build_context_window( - target_idx: int, known_codes_set: set[str], total_rows: int, half_size: int = 10 - ) -> dict[str, Any]: - min_start = max(0, target_idx - half_size) - min_end = min(total_rows - 1, target_idx + half_size) - - start_idx = min_start - end_idx = min_end - - found_known_above = False - found_known_below = False - known_positions: list[int] = [] - - for index in range(start_idx, target_idx): - reason = df.iloc[index].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_above = True - known_positions.append(index) - - for index in range(target_idx + 1, end_idx + 1): - reason = df.iloc[index].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_below = True - known_positions.append(index) - - if not found_known_above and min_start > 0: - search_idx = min_start - 1 - while search_idx >= 0: - reason = df.iloc[search_idx].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_above = True - known_positions.append(search_idx) - start_idx = search_idx - break - search_idx -= 1 - - if not found_known_below and min_end < total_rows - 1: - search_idx = min_end + 1 - while search_idx < total_rows: - reason = df.iloc[search_idx].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_below = True - known_positions.append(search_idx) - end_idx = search_idx - break - search_idx += 1 - - return { - "start": start_idx, - "end": end_idx, - "found_known": found_known_above or found_known_below, - "known_positions": known_positions, - } - - non_neg_mapping = build_non_neg_mapping(lvl_mapping) - known_codes = set(lvl_mapping.keys()) - - all_codes_in_full: dict[str, dict[str, Any]] = {} - for seg_idx, seg_df in enumerate(level_dfs): - for _, row in seg_df.iterrows(): - reason = row.get("reason", "") - sig = extract_reason_signature(reason) - if not sig or sig == PLACEHOLDER_REASON: - continue - if sig not in all_codes_in_full: - all_codes_in_full[sig] = { - "first_seg": seg_idx, - "first_id": row.get("id", 0), - "reason": reason, - } - - unseen_codes: dict[str, dict[str, Any]] = {} - unseen_neg_filtered: dict[str, dict[str, Any]] = {} - for sig, info in all_codes_in_full.items(): - if sig in known_codes: - continue - if has_neg_signal(info["reason"]): - unseen_neg_filtered[sig] = info - else: - unseen_codes[sig] = info - - logger.info( - f"Unseen codes total: {len(unseen_codes) + len(unseen_neg_filtered)}, " - f"NEG filtered: {len(unseen_neg_filtered)}, to process: {len(unseen_codes)}" - ) - - for sig in unseen_neg_filtered: - lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NEG_FILTERED"} - - if unseen_codes: - if strategy == "double_mapping": - fallback_success = 0 - fallback_failed = 0 - failed_codes = [] - for sig in unseen_codes: - non_neg_code = extract_non_neg_code(sig) - if non_neg_code in non_neg_mapping: - mapped_level = non_neg_mapping[non_neg_code] - lvl_mapping[sig] = { - "mapped_lvl": mapped_level, - "note": f"NON_NEG_FALLBACK from '{non_neg_code}'", - } - fallback_success += 1 - else: - lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NO_MATCH_FALLBACK"} - fallback_failed += 1 - failed_codes.append( - f"'{non_neg_code}' (from '{sig[:60]}...')" - if len(sig) > 60 - else f"'{non_neg_code}' (from '{sig}')" - ) - - logger.info( - f"Double mapping result: success={fallback_success}, failed={fallback_failed}" - ) - if failed_codes: - logger.warning( - f"Failed codes (non_neg not in mapping): {failed_codes[:5]}" - f"{'...' if len(failed_codes) > 5 else ''}" - ) - - elif strategy == "window_llm" and output_dir: - total_rows = len(df) - windows: list[dict[str, Any]] = [] - for sig, info in unseen_codes.items(): - first_id = info["first_id"] - first_seg = info["first_seg"] - df_indices = df.index[df["id"] == first_id].tolist() - if df_indices: - first_df_idx = df_indices[0] - window_info = build_context_window( - first_df_idx, known_codes, total_rows, window_half_size - ) - windows.append( - { - "code": sig, - "first_id": first_id, - "first_seg": first_seg, - "start": window_info["start"], - "end": window_info["end"], - "found_known": window_info["found_known"], - } - ) - - sorted_windows = sorted(windows, key=lambda window: window["start"]) - merged_windows: list[dict[str, Any]] = [] - current_window: dict[str, Any] | None = None - - for window in sorted_windows: - if current_window is None: - current_window = { - "start": window["start"], - "end": window["end"], - "codes": [window["code"]], - "segments": [window["first_seg"]], - } - elif window["start"] <= current_window["end"]: - current_window["end"] = max(current_window["end"], window["end"]) - current_window["codes"].append(window["code"]) - current_window["segments"].append(window["first_seg"]) - else: - merged_windows.append(current_window) - current_window = { - "start": window["start"], - "end": window["end"], - "codes": [window["code"]], - "segments": [window["first_seg"]], - } - - if current_window: - merged_windows.append(current_window) - - windows_dir = os.path.join(output_dir, "merged_windows") - os.makedirs(windows_dir, exist_ok=True) - - unseen_codes_set = set(unseen_codes.keys()) - unseen_neg_set = set(unseen_neg_filtered.keys()) - - for index, merged_window in enumerate(merged_windows): - window_df = df.iloc[ - merged_window["start"] : merged_window["end"] + 1 - ].copy() - - def get_code_status(row: pd.Series) -> str: - reason = row.get("reason", "") - sig = extract_reason_signature(reason) - if not sig: - return "" - if sig in unseen_codes_set: - return "UNSEEN_TARGET" - if sig in unseen_neg_set: - return "NEG_TO_NEGATIVE_ONE" - if sig in known_codes: - return "KNOWN" - return "" - - window_df["code_status"] = window_df.apply(get_code_status, axis=1) - window_path = os.path.join( - windows_dir, - f"window_{index + 1:02d}_rows_" - f"{merged_window['start']}-{merged_window['end']}.csv", - ) - window_df.to_csv(window_path, index=False, encoding="utf-8-sig") - - logger.debug( - f"Window LLM: {len(merged_windows)} windows created in {windows_dir}" - ) - - return lvl_mapping - - def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: """Collapse consecutive body rows into placeholder rows before LLM chunking. @@ -480,118 +154,68 @@ def execute_llm_heading_hierarchy( f"threshold={prompt_limt} | max_start={max_len}" ) - basic_idx = 0 - for idx, chunk in enumerate(level_dfs): - if (chunk["reason"].astype(str) != PLACEHOLDER_REASON).any(): - basic_idx = idx - break - basic_df = level_dfs[basic_idx] - if basic_idx != 0: - logger.info( - f"smart parse => promoted chunk {basic_idx} as basic_df " - f"(chunks 0..{basic_idx - 1} contain only placeholders)" - ) - full_preds: pd.DataFrame | None = None try: with stage_timer( "heading.hierarchy_llm", chunk_count=len(level_dfs), - base_chunk_rows=len(basic_df), compact_enabled=compact_enabled, source_row_count=len(raw_preds), model_name=model_name, ): - logger.debug("smart parse => interpreting hierarchy patterns...") - df4llm = basic_df.drop(columns=["reason", "level"], errors="ignore").copy() - df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm) - logger.debug(f"DataFrame transformation completed, rows: {len(df4llm)}") - - layout_res = hierarchy_judge( - df4llm, model_name, max_depth, toc_hierarchies, task="eval-headings" - ) - - layout_level_by_id: dict[Any, Any] = {} - if isinstance(layout_res, list): - for item in layout_res: - if isinstance(item, dict) and "id" in item and "level" in item: - layout_level_by_id[item["id"]] = item["level"] - - def level_for(row_id: Any) -> Any: - if row_id in layout_level_by_id: - return layout_level_by_id[row_id] - try: - return layout_level_by_id.get(int(row_id), -1) - except (TypeError, ValueError): - return -1 - - base_preds = ( - basic_df[["id", "heading", "reason"]].copy().reset_index(drop=True) - ) - base_preds.insert(2, "level", base_preds["id"].map(level_for)) - save_intermediate_csv( - base_preds, output_dir, f"preds_3_llm_base{csv_suffix}" - ) - + # ── Per-chunk independent LLM calls ── llm_levels: dict[int, Any] = {} - for _, row in base_preds.iterrows(): - row_id = row["id"] - if isinstance(row_id, bool): + + for chunk_idx, chunk_df in enumerate(level_dfs): + # Skip chunks that contain only placeholders + non_placeholder_mask = chunk_df["reason"].astype(str) != PLACEHOLDER_REASON + if not non_placeholder_mask.any(): + logger.debug( + f"smart parse => chunk {chunk_idx}: all placeholders, skipping" + ) continue - if isinstance(row_id, int): - llm_levels[row_id] = row["level"] - if len(level_dfs) > 1: - placeholder_mask_base = base_preds["reason"].eq(PLACEHOLDER_REASON) - figure_mask_base = base_preds["heading"].eq("Figure/Image") - exclude_mask_base = placeholder_mask_base | figure_mask_base - base_preds_for_mapping = base_preds[~exclude_mask_base].copy() - # Read origin levels from raw_preds (naive-stage estimates), - # not from compact output which no longer carries level. - base_candidate_ids = base_preds_for_mapping["id"].tolist() - raw_level_lookup = dict( - zip(raw_preds["id"].tolist(), raw_preds["level"].tolist()) - ) - base_origin_for_mapping = [ - raw_level_lookup.get(cid, -1) for cid in base_candidate_ids - ] + df4llm = chunk_df.drop(columns=["reason", "level"], errors="ignore").copy() + df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm) - base_preds_for_mapping, lvl_mapping = build_level_mapping( - base_preds_for_mapping, base_origin_for_mapping, mode="freq" + logger.info( + f"smart parse => chunk {chunk_idx}/{len(level_dfs)}: " + f"sending {len(df4llm)} rows to LLM" ) - logger.debug( - f"mapping development finished: {len(lvl_mapping)} rules " - f"(placeholders and Figure/Image excluded)" + chunk_result = hierarchy_judge( + df4llm, model_name, max_depth, toc_hierarchies, task="eval-headings" ) - logger.debug( - f"mapping dataframe to levels across {len(level_dfs)} chunks..." - ) - lvl_mapping = handle_unseen_codes( - preds_for_llm, level_dfs, lvl_mapping, output_dir + if isinstance(chunk_result, list): + for item in chunk_result: + if isinstance(item, dict) and "id" in item and "level" in item: + try: + llm_levels[int(item["id"])] = item["level"] + except (TypeError, ValueError): + pass + + # Save per-chunk intermediate CSV + chunk_preds = ( + chunk_df[["id", "heading", "reason"]].copy().reset_index(drop=True) ) - - for level_df in level_dfs: - placeholder_mask_chunk = level_df["reason"].eq(PLACEHOLDER_REASON) - figure_mask_chunk = level_df["heading"].eq("Figure/Image") - exclude_mask_chunk = placeholder_mask_chunk | figure_mask_chunk - non_excluded = level_df[~exclude_mask_chunk].copy() - if not non_excluded.empty: - non_excluded = execute_level_mapping(non_excluded, lvl_mapping) - for _, row in non_excluded.iterrows(): - row_id = row["id"] - if isinstance(row_id, bool): - continue - if isinstance(row_id, int): - llm_levels[row_id] = row["level"] - logger.info( - f"multi-chunk mapping produced {len(llm_levels)} id->level entries" + chunk_preds.insert( + 2, "level", + chunk_preds["id"].map( + lambda rid: llm_levels.get( + int(rid) if not isinstance(rid, str) or rid.isdigit() else -1, -1 + ) + ), ) - else: - logger.info( - "single chunk - skipping reason-code mapping, using LLM output directly" + save_intermediate_csv( + chunk_preds, output_dir, + f"preds_llm{csv_suffix}_{chunk_idx}" ) + logger.info( + f"smart parse => per-chunk LLM produced {len(llm_levels)} " + f"id->level entries across {len(level_dfs)} chunks" + ) + full_preds = raw_preds.copy()[["id", "heading", "level", "reason"]] def resolve_level(row_id: Any) -> int: @@ -604,11 +228,8 @@ def resolve_level(row_id: Any) -> int: return int(level) except (TypeError, ValueError): return -1 - + full_preds["level"] = full_preds["id"].map(resolve_level).astype(int) - save_intermediate_csv( - full_preds, output_dir, f"preds_4_llm_final{csv_suffix}" - ) except Exception as exc: logger.warning( @@ -616,3 +237,4 @@ def resolve_level(row_id: Any) -> int: ) full_preds = fallback_hierarchy(raw_preds.copy()) return full_preds + diff --git a/apps/worker/app/services/document_parser/structure/layout_parser.py b/apps/worker/app/services/document_parser/structure/layout_parser.py index 52541e793..9416d11e7 100755 --- a/apps/worker/app/services/document_parser/structure/layout_parser.py +++ b/apps/worker/app/services/document_parser/structure/layout_parser.py @@ -9,8 +9,6 @@ postprocess_headings, ) from app.services.document_parser.structure.heading_llm_executor import ( - build_level_mapping, - execute_level_mapping, execute_llm_heading_hierarchy, ) from app.services.document_parser.structure.heading_tree import ( @@ -321,12 +319,32 @@ def _compute_zone_boundaries(toc_hierarchies, coordinate_mode="post_removal"): def _resolve_first_toc_boundary(toc_hierarchies=None, first_toc_ele_num=None): - """Resolve the earliest available first-TOC boundary across coordinate sources.""" + """Resolve the earliest available first-TOC boundary across coordinate sources. + + Page-based TOC boundaries (``toc_range_unit == "page"``, produced by + DOC_AGENT/VLM for PDF/PPT) are in *page numbers*, NOT line/element IDs. + They must NOT be used for pre-TOC row removal because ``raw_preds["id"]`` + are line indices that restart from 0 in each shard. For these documents + the DOC_AGENT has already handled shard splitting around TOC pages. + """ toc_range_start = None + toc_unit = None if toc_hierarchies: first_range = toc_hierarchies[0].get("toc_range") if first_range and len(first_range) == 2: toc_range_start = first_range[0] + toc_unit = toc_hierarchies[0].get("toc_range_unit") + + # Page-based coordinates cannot be compared against line/element IDs. + if toc_unit == "page": + if first_toc_ele_num is not None: + # DOCX fallback: element-based boundary is safe to use. + return first_toc_ele_num + logger.debug( + "📌 Skipping pre-TOC removal: TOC uses page-based coordinates " + "(DOC_AGENT already handled shard boundaries)" + ) + return None candidates = [ value for value in (toc_range_start, first_toc_ele_num) if value is not None @@ -571,53 +589,23 @@ def _process_single_zone(zone_idx, zone_start, zone_end, zone_toc): f"📌 Spliced {len(pre_toc_rows)} pre-TOC lines back into predictions" ) - # Save heading_preds as preds_5 - save_intermediate_csv(heading_preds, output_dir, "preds_5_final_output") + # Save final heading predictions + save_intermediate_csv(heading_preds, output_dir, "preds_final") return heading_preds def est_hierarchies_naive(raw_preds, proceed_smart=True, output_dir=None): - """Detect hierarchies by non-LLM - - Args: - raw_preds: raw data - proceed_smart: whether to proceed with smart parsing - output_dir: output directory, used to save intermediate results CSV - """ - logger.debug("🚀 non-llm parsing => recursive processing") - save_preds = raw_preds.copy() - - heading_preds = postprocess_headings(raw_preds, task="collapse") - save_preds.insert( - save_preds.columns.get_loc("level") + 1, - "lvl_cola", - heading_preds["level"].tolist(), - ) + """Pre-LLM heading filtering via regex and negative pattern checks. - heading_preds = postprocess_headings(heading_preds, task="judge_negs") - save_preds.insert( - save_preds.columns.get_loc("lvl_cola") + 1, - "lvl_neg", - heading_preds["level"].tolist(), - ) - save_preds["reason"] = heading_preds["reason"] - - # mapping based on freq - if not proceed_smart: - heading_preds["level"] = heading_preds["level"].map( - lambda x: -1 if x == -2 else x - ) - heading_preds, lvl_mapping = build_level_mapping( - heading_preds, heading_preds["level"].tolist(), mode="freq" - ) - heading_preds = execute_level_mapping(heading_preds, lvl_mapping) - heading_preds.drop("origin_level", axis=1, inplace=True) - save_preds.insert( - save_preds.columns.get_loc("lvl_neg") + 1, - "lvl_map", - heading_preds["level"].tolist(), - ) + This stage determines the initial candidate/body split that + ``compact_for_llm`` relies on (level > -1 → candidate, -1 → body text). + The estimated level itself is NOT forwarded to the LLM. + The ``proceed_smart`` and ``output_dir`` parameters are retained for API + compatibility but have no effect. + """ + logger.debug("🚀 non-llm parsing => judge_negs filtering") + heading_preds = postprocess_headings(raw_preds, task="judge_negs") return heading_preds @@ -631,22 +619,19 @@ def est_hierarchies_llm( output_dir=None, csv_suffix="", ): - """LLM-based hierarchy detection — first chunk via LLM, remaining chunks via reason-code mapping. + """LLM-based hierarchy detection — all chunks evaluated independently. When ``KB_LAYOUT_LLM_COMPACT_INPUT`` is enabled (default), consecutive ``level == -1`` rows in ``raw_preds`` are folded into a single placeholder - row (``[N BODY LINES]``) before chunking. This shrinks the prompt, makes - most documents fit into a single chunk (skipping the lossy reason-code - mapping), and preserves the positional signal for the LLM. + row (``[N BODY LINES]``) before chunking. This shrinks the prompt and + preserves the positional signal for the LLM. Strategy: - 1. (Optional) Compact raw_preds so consecutive body rows become placeholders. - 2. Send only the first chunk to LLM for hierarchy prediction. - 3. Collect ``{id -> level}`` from the LLM response (int ids only). - 4. For multi-chunk docs, extend that mapping via reason-code mapping on - chunks 1..N (placeholders excluded). - 5. Expand the id->level mapping back onto the ORIGINAL ``raw_preds``; - any row not present in the mapping defaults to ``level = -1``. + 1. Compact raw_preds so consecutive body rows become placeholders. + 2. Split into chunks and send each independently to LLM. + 3. Collect ``{id -> level}`` from LLM responses (int ids only). + 4. Map levels back onto the ORIGINAL ``raw_preds``; + any row not in the map defaults to ``level = -1``. Args: raw_preds: raw data diff --git a/apps/worker/app/services/document_parser/tables/table_asset_writer.py b/apps/worker/app/services/document_parser/tables/table_asset_writer.py index 679ee028e..c8e5196f5 100644 --- a/apps/worker/app/services/document_parser/tables/table_asset_writer.py +++ b/apps/worker/app/services/document_parser/tables/table_asset_writer.py @@ -15,7 +15,6 @@ class TableAssetInput: keywords: str know_id: str addtime: str - page_nums: str = "" content: str | None = None tokens: str = "" length: int | None = None @@ -39,7 +38,6 @@ def write_table_asset(table_input: TableAssetInput) -> ParsedRow: tokens=table_input.tokens, connectto="", addtime=table_input.addtime, - page_nums=table_input.page_nums, length=table_input.length, ) From f34958a5be25415289f63b8a2aad8e7e82e6b604 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 28 May 2026 12:31:09 +0800 Subject: [PATCH 11/11] fix: resolve heading detection NEG condition inconsistencies and clean up stale pipeline code - Remove hardcoded num_pos/num_neg params; derive zero-filled arrays dynamically - Remove dead 'Sure' return branch from get_max_lvl, add -> int type annotation - Restore include_punc=False default in remove_by_conditions to defer punc checking to judge_negs second pass (reduces false heading filtering) - Remove dead 'collapse' task code and rename functions for clarity - Simplify redundant regex character class in NEG rule 3 - Normalize -2 -> 1 in est_hierarchies_naive for valid LLM-failure fallback Closes #112 --- .../structure/heading_candidates.py | 63 ++++++++++--------- .../structure/layout_parser.py | 18 ++++-- 2 files changed, 47 insertions(+), 34 deletions(-) diff --git a/apps/worker/app/services/document_parser/structure/heading_candidates.py b/apps/worker/app/services/document_parser/structure/heading_candidates.py index b790a6cf7..0a7596aed 100644 --- a/apps/worker/app/services/document_parser/structure/heading_candidates.py +++ b/apps/worker/app/services/document_parser/structure/heading_candidates.py @@ -88,7 +88,14 @@ def judge_by_conditions( return pos_triggered_code -def remove_by_conditions(text): +def remove_by_conditions(text, *, include_punc: bool = False): + """Evaluate negative (non-heading) conditions against *text*. + + ``include_punc`` controls the end-of-line punctuation rule. It is + intentionally **off** during initial scanning so that lines remain heading candidates. + The punctuation check is enabled only during the ``judge_negs`` second pass (after merges + may have altered heading text). + """ neg_conditions = [ r"^\d{3,}", r"(?i)(^https?://\S+|^www\.\S+|^P\.S|^\b\d{0,2}\s*(?:a\.m|p\.m)\b)", @@ -115,13 +122,18 @@ def remove_by_conditions(text): r"|Hz|kHz|MHz|GHz" r"|mol|mL|dL|dB|Nm|kN|MN|kW|MW|GW|hp|rpm|cc|cal|kcal)\b" ), - r"[.,;，。；]$", ] neg_triggered_code = [] for regex in neg_conditions: neg_triggered_code.append(1 if re.search(regex, text) else 0) - + + # End-of-line punctuation — only checked during judge_negs second pass. + if include_punc: + neg_triggered_code.append(1 if re.search(r"[.,;，。；]$", text) else 0) + else: + neg_triggered_code.append(0) + MAX_HEADING_TOKENS = 10 neg_triggered_code.append(1 if count_cn_en(text) > MAX_HEADING_TOKENS else 0) @@ -244,8 +256,8 @@ def postprocess_headings(df: pd.DataFrame, task: str, max_depth: int = -1) -> pd if task == "merge_continuous": return _merge_continuous_non_headings(df) - if task == "merge_short" or task == "collapse": - return _collapse_heading_groups(df, task) + if task == "merge_short": + return _merge_short_heading_groups(df) return df @@ -329,7 +341,7 @@ def _is_bold_docx_paragraph(paragraph: Any): def _judge_negative_headings(df: pd.DataFrame) -> pd.DataFrame: for index, row in df.iterrows(): - neg_code = remove_by_conditions(row["heading"]) + neg_code = remove_by_conditions(row["heading"], include_punc=True) if any(value > 0 for value in neg_code): current_code = str(df.loc[index, "reason"]) @@ -379,7 +391,7 @@ def _merge_continuous_non_headings(df: pd.DataFrame) -> pd.DataFrame: return pd.DataFrame(denoised_rows, columns=HEADING_COLUMNS) -def _collapse_heading_groups(df: pd.DataFrame, task: str) -> pd.DataFrame: +def _merge_short_heading_groups(df: pd.DataFrame) -> pd.DataFrame: group_to_indices = defaultdict(list) for index, row in df.iterrows(): level = row["level"] @@ -389,24 +401,22 @@ def _collapse_heading_groups(df: pd.DataFrame, task: str) -> pd.DataFrame: checked_pairs = set() for _, indices in group_to_indices.items(): - _collapse_recursive(df, task, indices, merge_threshold=3, checked_pairs=checked_pairs) - - if task == "merge_short": - drop_between = df.index[ - df["reason"].astype(str).str.startswith("Merged into", na=False) - ].tolist() - if drop_between: - logger.debug( - f"🛠️ Delete rows labeled as merged into, total {len(drop_between)} rows" - ) - df.drop(drop_between, inplace=True) - df.reset_index(drop=True, inplace=True) + _merge_short_recursive(df, indices, merge_threshold=3, checked_pairs=checked_pairs) + + drop_between = df.index[ + df["reason"].astype(str).str.startswith("Merged into", na=False) + ].tolist() + if drop_between: + logger.debug( + f"🛠️ Delete rows labeled as merged into, total {len(drop_between)} rows" + ) + df.drop(drop_between, inplace=True) + df.reset_index(drop=True, inplace=True) return df -def _collapse_recursive( +def _merge_short_recursive( df: pd.DataFrame, - task: str, indices: list[int], merge_threshold: int = 3, checked_pairs: set[tuple[int, int]] | None = None, @@ -425,18 +435,11 @@ def _collapse_recursive( between = df.loc[index + 1 : next_index - 1] current_text = df.at[index, "heading"].strip() - next_text = df.at[next_index, "heading"].strip() - if task == "merge_short" and len(between) > 0: + if len(between) > 0: _merge_short_between_headings( df, between, index, next_index, current_text, merge_threshold ) - elif task == "collapse" and len(between) == 0: - logger.debug( - f"⚠️ Empty between i={current_text[:15]}, j={next_text[:15]} => set i.level=-1, j.level=Not Sure" - ) - df.at[index, "level"] = -2 - df.at[next_index, "level"] = -2 sub_between = between[between["level"] != -1] code_to_sub = defaultdict(list) @@ -447,7 +450,7 @@ def _collapse_recursive( code_to_sub[(level, reason)].append(row_index) for _, sub_indices in code_to_sub.items(): - _collapse_recursive(df, task, sub_indices, merge_threshold, checked_pairs) + _merge_short_recursive(df, sub_indices, merge_threshold, checked_pairs) def _merge_short_between_headings( diff --git a/apps/worker/app/services/document_parser/structure/layout_parser.py b/apps/worker/app/services/document_parser/structure/layout_parser.py index 9416d11e7..78c9f56ea 100755 --- a/apps/worker/app/services/document_parser/structure/layout_parser.py +++ b/apps/worker/app/services/document_parser/structure/layout_parser.py @@ -595,17 +595,27 @@ def _process_single_zone(zone_idx, zone_start, zone_end, zone_toc): def est_hierarchies_naive(raw_preds, proceed_smart=True, output_dir=None): - """Pre-LLM heading filtering via regex and negative pattern checks. + """Regex-only heading hierarchy estimation (LLM fallback). - This stage determines the initial candidate/body split that - ``compact_for_llm`` relies on (level > -1 → candidate, -1 → body text). - The estimated level itself is NOT forwarded to the LLM. + When used as the primary pipeline step (before LLM), this function + determines the initial candidate/body split that ``compact_for_llm`` + relies on (level > -1 → candidate, -1 → body text). + + When used as a **fallback** after LLM failure, the returned levels + must be usable for tree construction. Single-level POS matches + (``get_max_lvl`` returns -2 for patterns like) + are normalized to level 1 so the output forms a valid hierarchy. The ``proceed_smart`` and ``output_dir`` parameters are retained for API compatibility but have no effect. """ logger.debug("🚀 non-llm parsing => judge_negs filtering") heading_preds = postprocess_headings(raw_preds, task="judge_negs") + + # legitimate heading candidates; default them to top-level. + heading_preds["level"] = heading_preds["level"].map( + lambda x: 1 if x == -2 else x + ) return heading_preds