diff --git a/apps/worker/app/services/document_agent/executor/prompts.py b/apps/worker/app/services/document_agent/executor/prompts.py index 8c225d2cc..06fe0458a 100644 --- a/apps/worker/app/services/document_agent/executor/prompts.py +++ b/apps/worker/app/services/document_agent/executor/prompts.py @@ -1,15 +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." + "You are the executor of a document profiling agent. Decide the next tool " + "call from the blackboard facts and available tools. Return strict JSON with " + "keys: action (must be tool_call), rationale, tool_name, tool_args. " + "Use inspect.pages when more visual evidence is needed, grep.text when " + "native-PDF text evidence is needed, propose.shard_plan when evidence is " + "sufficient to shard, validate.anatomy_map after a shard plan exists, and " + "the verdict tool to finish: verdict(status=success) only after validation " + "succeeds, or verdict(status=abort, rationale=...) only when the document " + "cannot be profiled. Do not invent other finish actions." ) __all__ = ["REFLEXION_INSTRUCTIONS"] diff --git a/apps/worker/app/services/document_agent/executor/react_loop.py b/apps/worker/app/services/document_agent/executor/react_loop.py index a890f702f..9917ac4fa 100644 --- a/apps/worker/app/services/document_agent/executor/react_loop.py +++ b/apps/worker/app/services/document_agent/executor/react_loop.py @@ -56,27 +56,49 @@ def _compact_blackboard(ctx: ToolContext) -> dict[str, Any]: } +def _coerce_legacy_finish(data: dict[str, Any]) -> ReflexionDecision: + """Map obsolete ``action=verdict_now`` into a real tool call. + + Finish belongs exclusively to the ``verdict`` tool. A bare ``verdict_now`` + without an explicit status is treated as ready-to-shard, never as abort. + """ + rationale = str(data.get("rationale") or "") + raw_verdict = data.get("verdict") + if isinstance(raw_verdict, dict): + status = str(raw_verdict.get("status") or "").strip().lower() + if status in {"success", "abort"}: + return ReflexionDecision( + action="tool_call", + rationale=rationale, + tool_name="verdict", + tool_args={ + "status": status, + "rationale": str( + raw_verdict.get("rationale") or rationale or status + ), + }, + ) + return ReflexionDecision( + action="tool_call", + rationale=rationale or "Legacy verdict_now without status; propose shard plan.", + tool_name="propose.shard_plan", + tool_args={}, + ) + + 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 = str(data.get("action") or "tool_call").strip().lower() + if action == "verdict_now": + return _coerce_legacy_finish(data if isinstance(data, dict) else {}) + if action != "tool_call": 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] + action="tool_call", rationale=str(data.get("rationale") or ""), tool_name=data.get("tool_name"), tool_args=dict(data.get("tool_args") or {}), - verdict=verdict, + verdict=None, ) @@ -98,6 +120,15 @@ 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) + if decision.action != "tool_call": + # Defensive: only tool_call is a legal executor step. + decision = ReflexionDecision( + action="tool_call", + rationale=decision.rationale + or "Non-tool executor action coerced to propose.shard_plan.", + tool_name="propose.shard_plan", + tool_args={}, + ) self.ctx.blackboard.global_signals.setdefault("reflexion_decisions", []).append( decision.to_dict() ) @@ -111,14 +142,9 @@ def run(self) -> ExecutorResult: 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 ( + tool_name, tool_args = self._resolve_tool_call(decision) + if tool_name == "verdict" and str(tool_args.get("status") or "") == "success": + if not ( self.ctx.blackboard.validation_report and self.ctx.blackboard.validation_report.get("valid") is True ): @@ -131,11 +157,6 @@ def run(self) -> ExecutorResult: 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( @@ -205,6 +226,22 @@ def _is_deterministic_mode(self) -> bool: 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 + if decision.action != "tool_call": + decision = ReflexionDecision( + action="tool_call", + rationale=decision.rationale + or "Initial non-tool decision coerced to propose.shard_plan.", + tool_name="propose.shard_plan", + tool_args={}, + ) + elif not decision.tool_name: + decision = ReflexionDecision( + action="tool_call", + rationale=decision.rationale + or "Initial decision missing tool; propose shard plan.", + tool_name="propose.shard_plan", + tool_args={}, + ) 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: @@ -224,9 +261,13 @@ def _next_decision(self, round_index: int) -> tuple[ReflexionDecision, ToolResul est = estimate_tokens(prompt) if not self.ctx.budget.try_reserve("plan", est): decision = ReflexionDecision( - action="verdict_now", + action="tool_call", rationale="Planner budget exhausted.", - verdict=AgentVerdict(status="abort", rationale="Planner budget exhausted."), + tool_name="verdict", + tool_args={ + "status": "abort", + "rationale": "Planner budget exhausted.", + }, ) return decision, ToolResult( status="ok", @@ -303,4 +344,3 @@ def _deterministic_decision(self) -> ReflexionDecision: tool_name="validate.anatomy_map", tool_args={}, ) - diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index d078f9d58..d422a55b8 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -10,7 +10,10 @@ PageKind = Literal["normal", "landscape"] TocFailureKind = Literal["none", "confirm_failed", "rejected_all", "degraded"] -ReflexionAction = Literal["tool_call", "verdict_now"] +# Executor loop steps are always tool calls. Profile success/abort is owned +# exclusively by the ``verdict`` tool (AgentVerdict.status), not by a separate +# ReflexionAction shortcut. +ReflexionAction = Literal["tool_call"] VerdictStatus = Literal["success", "abort"] @@ -28,8 +31,6 @@ class PageFeature: height: float has_asset: bool is_blank_like: bool - # PDF-space boxes for detected assets; None when none were extracted. - asset_bboxes: list[dict[str, Any]] | None = None def to_dict(self) -> dict[str, Any]: return asdict(self) diff --git a/apps/worker/app/services/document_agent/persist/__init__.py b/apps/worker/app/services/document_agent/persist/__init__.py index a34439602..26c0d806a 100644 --- a/apps/worker/app/services/document_agent/persist/__init__.py +++ b/apps/worker/app/services/document_agent/persist/__init__.py @@ -1,8 +1,13 @@ """Persist anatomy map artifacts.""" from app.services.document_agent.tools.persist_anatomy_map import ( + DOC_PROFILE_FILENAME, build_anatomy_map, persist_anatomy_map, ) -__all__ = ["build_anatomy_map", "persist_anatomy_map"] +__all__ = [ + "DOC_PROFILE_FILENAME", + "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 7e5e1ce01..acf656003 100644 --- a/apps/worker/app/services/document_agent/planner/planner.py +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -156,7 +156,11 @@ def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDec header_y=header_y, footer_y=footer_y, ) - next_action = str(data.get("next_action") or "ready_to_shard") + next_action = str(data.get("next_action") or "ready_to_shard").strip().lower() + # Legacy models may still emit verdict_now; that is not a planner finish + # signal — fall through to ready_to_shard so the executor owns success/abort. + if next_action == "verdict_now": + next_action = "ready_to_shard" tool_name: str | None = None tool_args: dict[str, Any] = {} if next_action == "inspect_more": @@ -171,12 +175,6 @@ def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDec 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", diff --git a/apps/worker/app/services/document_agent/planner/prompts.py b/apps/worker/app/services/document_agent/planner/prompts.py index 4ca93ef8e..2c6c48145 100644 --- a/apps/worker/app/services/document_agent/planner/prompts.py +++ b/apps/worker/app/services/document_agent/planner/prompts.py @@ -17,10 +17,12 @@ "footer_y is the highest footer line you observe (smallest y) when any " "footer is present, otherwise null. When both are set, require " "header_y < footer_y. " - "next_action must be one of inspect_more, grep_text, ready_to_shard, " - "verdict_now. Use inspect_more only when extra page screenshots are needed. " + "next_action must be one of inspect_more, grep_text, ready_to_shard. " + "Use inspect_more only when 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." + "structure. Use ready_to_shard when evidence is sufficient to propose shards. " + "Do not finish or abort the profile run from next_action; the executor owns " + "success/abort via the verdict tool. Do not output a fixed step plan." ) __all__ = ["PLANNER_INSTRUCTIONS"] diff --git a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py index 5815a4120..f33717875 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -651,8 +651,9 @@ def _entries_to_tree(entries: list[dict[str, Any]]) -> list[TitleNode]: stack: list[tuple[int, TitleNode]] = [] for entry in entries: - raw_title = str(entry.get("heading") or "").strip() - title = clean_toc_title(raw_title) or normalize_heading_text(raw_title) + # Keep original TOC heading (incl. numbering). Prefix stripping belongs + # only in text/compact match helpers used for null-page parents. + title = normalize_heading_text(str(entry.get("heading") or "")) level = _safe_int(entry.get("level")) or 1 if not title or len(title) < 2: continue 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 a2de3d295..1480eb4d8 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 @@ -6,6 +6,7 @@ import json import os import time +from dataclasses import dataclass, field from pathlib import Path from typing import Any, cast @@ -20,6 +21,7 @@ ) from app.services.document_agent.registry import register_tool from app.services.document_agent.tools.vlm_toc_extractor import ( + TOC_VLM_MAX_TOKENS, vlm_entries_to_toc_hierarchies, ) from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( @@ -34,81 +36,150 @@ MAX_BOUNDARY_ROUNDS = 6 MAX_TOC_PAGES = BOUNDARY_STEP_PAGES * MAX_BOUNDARY_ROUNDS # 30 +_CONFIRM_STAGE = "toc_confirm" +_CONFIRM_TOKENS_PER_PAGE = 800 + +_CONFIRM_PROMPT = ( + "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 strict JSON (no markdown fences):\n" + '{"pages": [{"page": , "is_toc_start": true/false, ' + '"reason": "brief reason"}]}' +) + # -- 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 +def _render_expand_window_worker( + queue, + pdf_path: str, + pages: list[int], + output_dir: str, + dpi: int, + anchor_page: int, ) -> None: + """Render one Phase2 expand window in a single child process. + + Opens the PDF once and writes ``toc_a{anchor}_p{page}.png`` for each page. + """ import pymupdf # type: ignore[import] + results: list[dict[str, Any]] = [] + doc = None 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) + mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) + for page_num in pages: + idx = page_num - 1 + if not (0 <= idx < doc.page_count): + continue + pix = doc[idx].get_pixmap(matrix=mat) + png_path = os.path.join( + output_dir, + f"toc_a{anchor_page}_p{page_num}.png", + ) + pix.save(png_path) + results.append({"page": page_num, "png_path": png_path}) finally: try: - doc.close() + if doc is not None: + doc.close() except Exception: pass gc.collect() - queue.put({"ok": True, "png_path": output_path}) + queue.put({"ok": True, "results": results}) # -- VLM helpers --------------------------------------------------------------- -def _vlm_confirm_anchors( - anchor_pages: list[TocAnchorPage], - model: str, - budget: Any | None = None, -) -> tuple[list[TocAnchorPage], bool, list[TocEvidence]]: - """Phase 1: send all anchor PNGs to VLM, ask which are real TOC starts.""" - from shared.services.ai.llm_overrides import get_vision_client +def _iter_chunks(items: list[Any], size: int) -> list[list[Any]]: + if size <= 0: + raise ValueError("chunk size must be positive") + return [items[i : i + size] for i in range(0, len(items), size)] + + +def _parse_confirm_items(raw: str) -> list[dict[str, Any]]: + 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 = [] + return [item for item in items if isinstance(item, dict)] + + +def _evidence_from_confirm_items( + items: list[dict[str, Any]], +) -> tuple[set[int], dict[int, TocEvidence]]: + confirmed_pages: set[int] = set() + evidence_by_page: dict[int, TocEvidence] = {} + for item in items: + if "page" not in item: + continue + page = int(item["page"]) + is_toc_start = bool(item.get("is_toc_start")) + if is_toc_start: + confirmed_pages.add(page) + raw_confidence = item.get("confidence") + try: + confidence = ( + float(raw_confidence) + if raw_confidence is not None + else (0.95 if is_toc_start else 0.05) + ) + except (TypeError, ValueError): + confidence = 0.95 if is_toc_start else 0.05 + evidence_by_page[page] = TocEvidence( + page_index=page, + source="vlm", + confidence=max(0.0, min(1.0, confidence)), + reason=str(item.get("reason") or ""), + ) + return confirmed_pages, evidence_by_page - if not anchor_pages: - return [], False, [] +def _confirm_anchor_chunk( + chunk: list[TocAnchorPage], + *, + model: str, + budget: Any | None, +) -> tuple[set[int], dict[int, TocEvidence], bool]: + """Confirm one BOUNDARY_STEP_PAGES-sized anchor chunk. + + Returns ``(confirmed_pages, evidence_by_page, failed)``. + """ import base64 - # Build multi-image message + from shared.services.ai.llm_overrides import get_vision_client + + if not chunk: + return set(), {}, False + 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"}]' - ), - } + {"type": "text", "text": _CONFIRM_PROMPT}, ] - - for anchor in anchor_pages: + for anchor in chunk: 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} ---", - } + {"type": "text", "text": f"\n--- Page {anchor.page} ---"} ) content_parts.append( { @@ -118,20 +189,22 @@ def _vlm_confirm_anchors( ) messages = cast(Any, [{"role": "user", "content": content_parts}]) - est = estimate_tokens(str(content_parts[0]["text"])) + len(anchor_pages) * 800 - stage = "toc_confirm" - if budget and not budget.try_reserve("visual", est, stage=stage): - logger.warning("[extract.toc] insufficient visual budget for anchor confirmation") - return [], True, [] + est = estimate_tokens(_CONFIRM_PROMPT) + len(chunk) * _CONFIRM_TOKENS_PER_PAGE + if budget and not budget.try_reserve("visual", est, stage=_CONFIRM_STAGE): + logger.warning( + "[extract.toc] insufficient visual budget for confirm chunk pages={}", + [a.page for a in chunk], + ) + return set(), {}, True try: client, resolved_model = get_vision_client(requested_model=model) - model = resolved_model or model + resolved = resolved_model or model raw, usage = client.chat_completion_with_usage( messages=messages, - model=model, + model=resolved, temperature=0.1, - max_tokens=500, + max_tokens=TOC_VLM_MAX_TOKENS, response_format={"type": "json_object"}, usage_task="document_agent.toc_anchor_confirm", ) @@ -140,72 +213,361 @@ def _vlm_confirm_anchors( "visual", actual=usage.get("total_tokens", est), est=est, - stage=stage, + stage=_CONFIRM_STAGE, ) - 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 + confirmed_pages, evidence_by_page = _evidence_from_confirm_items( + _parse_confirm_items(raw) + ) + return confirmed_pages, evidence_by_page, False + except Exception as exc: + if budget: + budget.refund("visual", est=est, stage=_CONFIRM_STAGE) + logger.warning( + "[extract.toc] VLM confirm chunk failed pages={}: {}", + [a.page for a in chunk], + exc, + ) + return set(), {}, True + + +def _vlm_confirm_anchors( + anchor_pages: list[TocAnchorPage], + model: str, + budget: Any | None = None, +) -> tuple[list[TocAnchorPage], bool, list[TocEvidence]]: + """Phase 1: confirm TOC starts in BOUNDARY_STEP_PAGES batches (concurrent).""" + if not anchor_pages: + return [], False, [] + + from gevent.pool import Pool as GeventPool + + chunks = _iter_chunks(anchor_pages, BOUNDARY_STEP_PAGES) + logger.info( + "[extract.toc] Phase 1 confirm: {} anchors → {} chunks (size={}, concurrency={})", + len(anchor_pages), + len(chunks), + BOUNDARY_STEP_PAGES, + min(BOUNDARY_STEP_PAGES, len(chunks)), + ) + + pool = GeventPool(size=min(BOUNDARY_STEP_PAGES, len(chunks))) + jobs = [ + pool.spawn(_confirm_anchor_chunk, chunk, model=model, budget=budget) + for chunk in chunks + ] + pool.join() + + confirmed_pages: set[int] = set() + evidence_by_page: dict[int, TocEvidence] = {} + chunk_failures = 0 + chunk_successes = 0 + for job in jobs: + try: + chunk_confirmed, chunk_evidence, failed = job.get() + except Exception as exc: + chunk_failures += 1 + logger.warning("[extract.toc] confirm greenlet failed: {}", exc) + continue + if failed: + chunk_failures += 1 + continue + chunk_successes += 1 + confirmed_pages.update(chunk_confirmed) + evidence_by_page.update(chunk_evidence) + + confirm_failed = chunk_successes == 0 and chunk_failures > 0 + 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] + evidence = [ + evidence_by_page.get( + a.page, + TocEvidence( + page_index=a.page, + source="vlm", + confidence=0.05, + reason=( + "confirm batch failed for this candidate" + if confirm_failed + else "VLM response omitted this candidate page" + ), + ), + ) + for a in anchor_pages + ] + logger.info( + "[extract.toc] Phase 1 done: confirmed={} rejected={} " + "chunk_ok={} chunk_fail={} confirm_failed={}", + len(confirmed), + rejected, + chunk_successes, + chunk_failures, + confirm_failed, + ) + return confirmed, confirm_failed, evidence + + +# -- Phase 2: expand + extract per confirmed start ----------------------------- + + +@dataclass +class _TocRegionResult: + anchor_page: int + entries: list[dict[str, Any]] = field(default_factory=list) + toc_pages: list[int] = field(default_factory=list) + hierarchies: list[dict[str, Any]] = field(default_factory=list) + batch_meta: list[dict[str, Any]] = field(default_factory=list) + batch_trace: list[dict[str, Any]] = field(default_factory=list) + error: str | None = None + + +def _render_toc_page_batch( + *, + pdf_path: str, + output_dir: str, + dpi: int, + anchor_page: int, + batch_pages: list[int], + render_lock: Any, + reuse_png_by_page: dict[int, str] | None = None, +) -> list[tuple[int, str]]: + """Render one expand window under the Phase2 global render lock. + + One child process renders the whole window (PDF opened once). Across + anchors, PyMuPDF stays serial to avoid gevent ThreadPool deadlocks. + VLM calls happen outside this lock. + """ + reuse = reuse_png_by_page or {} + planned: dict[int, str] = {} + pages_to_render: list[int] = [] + for page_num in batch_pages: + existing = reuse.get(page_num) + if existing and os.path.isfile(existing): + planned[page_num] = existing else: - items = [] + pages_to_render.append(page_num) + + with render_lock: + if pages_to_render: + result = run_in_child_process( + _render_expand_window_worker, + pdf_path, + pages_to_render, + output_dir, + dpi, + anchor_page, + timeout=120, + ) + for item in result.get("results") or []: + planned[int(item["page"])] = str(item["png_path"]) + + page_pngs: list[tuple[int, str]] = [] + missing: list[int] = [] + for page_num in batch_pages: + png_path = planned.get(page_num) + if not png_path: + missing.append(page_num) + continue + page_pngs.append((page_num, png_path)) + if missing: + raise RuntimeError( + f"TOC expand render missing pages {missing} for anchor {anchor_page}" + ) + return page_pngs - confirmed_pages: set[int] = set() - evidence_by_page: dict[int, TocEvidence] = {} - for item in items: - if not isinstance(item, dict) or "page" not in item: - continue - page = int(item["page"]) - is_toc_start = bool(item.get("is_toc_start")) - if is_toc_start: - confirmed_pages.add(page) - raw_confidence = item.get("confidence") - try: - confidence = ( - float(raw_confidence) - if raw_confidence is not None - else (0.95 if is_toc_start else 0.05) + +def _extract_region_for_anchor( + anchor: TocAnchorPage, + *, + pdf_path: str, + page_count: int, + output_dir: str, + dpi: int, + model: str, + render_lock: Any, +) -> _TocRegionResult: + """Expand + extract for one confirmed TOC start. + + Rounds within a start stay serial (continuation context). Different + starts run concurrently for VLM, but page renders share ``render_lock``. + """ + from app.services.document_agent.tools.vlm_toc_extractor import ( + vlm_extract_toc_batch, + ) + + anchor_page = anchor.page + region_entries: list[dict[str, Any]] = [] + region_toc_pages: list[int] = [] + region_scan_end = anchor_page + batch_meta: list[dict[str, Any]] = [] + batch_trace: list[dict[str, Any]] = [] + + try: + 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, + ) + + reuse_png_by_page: dict[int, str] = {} + if ( + round_idx == 0 + and anchor.png_path + and os.path.isfile(anchor.png_path) + ): + # Phase1 already rendered the confirmed start page. + reuse_png_by_page[anchor_page] = anchor.png_path + + page_pngs = _render_toc_page_batch( + pdf_path=pdf_path, + output_dir=output_dir, + dpi=dpi, + anchor_page=anchor_page, + batch_pages=batch_pages, + render_lock=render_lock, + reuse_png_by_page=reuse_png_by_page, + ) + + 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) + 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), + } + ) + + 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, ) - except (TypeError, ValueError): - confidence = 0.95 if is_toc_start else 0.05 - evidence_by_page[page] = TocEvidence( - page_index=page, - source="vlm", - confidence=max(0.0, min(1.0, confidence)), - reason=str(item.get("reason") or ""), + break + if batch_end >= page_count: + break + logger.info( + "[extract.toc] last page {} still TOC, expanding window", + batch_end, ) - 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] - evidence = [ - evidence_by_page.get( - a.page, - TocEvidence( - page_index=a.page, - source="vlm", - confidence=0.05, - reason="VLM response omitted this candidate page", - ), + hierarchies: list[dict[str, Any]] = [] + if region_entries: + hierarchies = vlm_entries_to_toc_hierarchies( + region_entries, + toc_page_nums=region_toc_pages, + scan_end_page=region_scan_end, + page_count=page_count, ) - for a in anchor_pages - ] - logger.info( - "[extract.toc] VLM confirmed {} TOC starts, rejected pages: {}", - len(confirmed), - rejected, + return _TocRegionResult( + anchor_page=anchor_page, + entries=region_entries, + toc_pages=region_toc_pages, + hierarchies=hierarchies, + batch_meta=batch_meta, + batch_trace=batch_trace, ) - return confirmed, False, evidence except Exception as exc: - if budget: - budget.refund("visual", est=est, stage=stage) logger.warning( - "[extract.toc] VLM anchor confirmation failed: {}, " - "falling back to no confirmed anchors (safe degradation)", + "[extract.toc] anchor {} region extract failed: {}", + anchor_page, exc, ) - return [], True, [] + return _TocRegionResult( + anchor_page=anchor_page, + entries=region_entries, + toc_pages=region_toc_pages, + batch_meta=batch_meta, + batch_trace=batch_trace, + error=str(exc), + ) + + +def _extract_regions_for_confirmed_anchors( + confirmed: list[TocAnchorPage], + *, + pdf_path: str, + page_count: int, + output_dir: str, + dpi: int, + model: str, +) -> list[_TocRegionResult]: + """Phase 2: concurrent VLM per start; serial PyMuPDF renders across starts.""" + if not confirmed: + return [] + + from gevent.lock import Semaphore + from gevent.pool import Pool as GeventPool + + pool_size = min(BOUNDARY_STEP_PAGES, len(confirmed)) + render_lock = Semaphore(1) + logger.info( + "[extract.toc] Phase 2 extract: {} confirmed starts, " + "vlm_concurrency={}, render=serial", + len(confirmed), + pool_size, + ) + pool = GeventPool(size=pool_size) + jobs = [ + pool.spawn( + _extract_region_for_anchor, + anchor, + pdf_path=pdf_path, + page_count=page_count, + output_dir=output_dir, + dpi=dpi, + model=model, + render_lock=render_lock, + ) + for anchor in confirmed + ] + pool.join() + + by_anchor: dict[int, _TocRegionResult] = {} + for job in jobs: + try: + result = job.get() + except Exception as exc: + logger.warning("[extract.toc] region greenlet failed: {}", exc) + continue + by_anchor[result.anchor_page] = result + + # Preserve document page order when merging regions. + ordered: list[_TocRegionResult] = [] + for anchor in confirmed: + result = by_anchor.get(anchor.page) + if result is None: + ordered.append( + _TocRegionResult( + anchor_page=anchor.page, + error="region greenlet failed", + ) + ) + else: + ordered.append(result) + return ordered # -- Main tool ----------------------------------------------------------------- @@ -261,7 +623,7 @@ def extract_toc_with_boundaries( ) os.makedirs(output_dir, exist_ok=True) - # -- Phase 1: VLM confirm anchors ----------------------------------------- + # -- Phase 1: VLM confirm anchors (batched + concurrent) ------------------- confirmed, confirm_failed, confirm_evidence = _vlm_confirm_anchors( anchors, model, budget=ctx.budget ) @@ -305,14 +667,14 @@ def extract_toc_with_boundaries( debug=debug_info, ) - # -- 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, + # -- Phase 2: per-confirmed-start expand + extract (concurrent across starts) + region_results = _extract_regions_for_confirmed_anchors( + confirmed, + pdf_path=ctx.pdf_path, + page_count=page_count, + output_dir=output_dir, + dpi=dpi, + model=model, ) all_entries: list[dict[str, Any]] = [] @@ -321,98 +683,27 @@ def extract_toc_with_boundaries( batch_meta: list[dict[str, Any]] = [] batch_trace: list[dict[str, Any]] = [] - 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 + for region in region_results: + if region.error: + warnings.append( + f"toc_region_failed:anchor={region.anchor_page}:{region.error}" ) - 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, + logger.warning( + "[extract.toc] anchor {} region failed: {}", + region.anchor_page, + region.error, ) - toc_hierarchies.extend(region_hierarchies) - else: + continue + all_entries.extend(region.entries) + all_toc_pages.extend(region.toc_pages) + batch_meta.extend(region.batch_meta) + batch_trace.extend(region.batch_trace) + if region.hierarchies: + toc_hierarchies.extend(region.hierarchies) + elif not region.entries: logger.warning( "[extract.toc] anchor {} produced no TOC entries", - anchor_page, + region.anchor_page, ) if not all_entries: 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 1e17a2202..f5865d856 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 @@ -10,6 +10,10 @@ from app.services.document_agent.manifest import PageAnatomyMap, ToolContext, ToolResult +# Debug / package-root artifact name. Keep exported so page_memory debug scripts +# can import it from ``document_agent.persist``. +DOC_PROFILE_FILENAME = "doc_profile.json" + def _artifact_dir(ctx: ToolContext) -> Path: if ctx.output_dir: 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 be92b853c..2c3de3598 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 @@ -203,7 +203,7 @@ def _probe_visual_assets( header_y: float | None, footer_y: float | None, ) -> dict[str, Any]: - """Collect counts + coarse asset bboxes from images / tables / drawings.""" + """Collect counts + has_asset gate from images / tables / drawings.""" image_area = 0.0 bboxes: list[dict[str, Any]] = [] seen_image_rects: set[tuple[float, float, float, float]] = set() @@ -276,7 +276,8 @@ def _probe_visual_assets( "image_count": image_count, "table_count": table_count, "drawings_count": drawings_count, - "asset_bboxes": bboxes or None, + # Geometry is only used to derive the gate; do not persist bboxes. + "has_asset": bool(bboxes), } @@ -364,7 +365,6 @@ def probe_page_features(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: height=float(item.get("height") or 0.0), has_asset=False, is_blank_like=bool(item.get("is_blank_like")), - asset_bboxes=None, ) for item in (result.get("features") or []) ] @@ -412,7 +412,7 @@ def probe_page_assets(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: updated: list[PageFeature] = [] for feature in ctx.blackboard.page_features: visual = by_page.get(feature.page) or {} - has_asset = visual.get("asset_bboxes") is not None + has_asset = bool(visual.get("has_asset")) updated.append( PageFeature( page=feature.page, @@ -427,11 +427,6 @@ def probe_page_assets(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: height=feature.height, has_asset=has_asset, is_blank_like=feature.raw_text_length < 50 and not has_asset, - asset_bboxes=( - list(visual["asset_bboxes"]) - if isinstance(visual.get("asset_bboxes"), list) - else None - ), ) ) ctx.blackboard.page_features = sorted(updated, key=lambda f: f.page) 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 9c6c87c95..cc331b290 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 @@ -8,6 +8,9 @@ from dataclasses import dataclass from typing import Any, cast +# Shared completion budget for TOC VLM calls (confirm batches + extract batches). +TOC_VLM_MAX_TOKENS = 8192 + # --------------------------------------------------------------------------- # Batch-mode prompt: send a window of candidate pages in one VLM call. @@ -196,7 +199,7 @@ def vlm_extract_toc_batch( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, temperature=0.1, - max_tokens=8192, + max_tokens=TOC_VLM_MAX_TOKENS, response_format={"type": "json_object"}, usage_task="document_agent.vlm_toc_batch", ) 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 0af790288..07fe1c6d0 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 @@ -263,6 +263,16 @@ def demote_consecutive_same_level( Note: a candidate the LLM already demoted to -1 naturally breaks the run, because -1 never equals a positive level — it acts like a body separator. + + TODO(heading-demote): Consider an iterative fixpoint pass that re-runs this + demotion while treating already-demoted ids as transparent (skipped for + adjacency only; final level stays -1). That would clean leftover singleton + TOC lines (e.g. ``2.4 foo....11``) after their same-level siblings were + demoted, i.e. pure outline regions with no body placeholders. Do NOT ship + without guarding the false-positive where empty subsections demote first + (``2.4 / 2.4.1 / 2.4.2 / 2.5 / 2.5.1 / [BODY]``) and a later pass then + wrongly merges real sibling parents ``2.4`` and ``2.5`` into one same-level + run. Deferred until we have TOC-page vs empty-subsection regression cases. """ demoted: set[int] = set() run: list[tuple[int, int]] = [] # (row_id, level) for a no-body candidate run diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 43f664029..ce835659d 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -28,9 +28,6 @@ locate_title_compact_strict, resolve_hierarchy_page_ranges, ) -from app.services.document_parser.structure.body_boundary import ( - clean_toc_title, -) from loguru import logger from shared.services.chunks.path_segments import ( append_document_path, @@ -298,7 +295,8 @@ def _range_to_skeleton( ) -> SectionSkeleton: start_page = _clamp_page(item.start_page, page_count) end_page = _clamp_page(item.end_page, page_count) - path_titles = [clean_toc_title(title) or title for title in item.path_titles] + # Keep original TOC titles (incl. numbering) in section_path / HIERARCHY. + path_titles = [str(title).strip() for title in item.path_titles if str(title).strip()] section_path = join_document_path([filename, *path_titles]) parent_path = ( join_document_path([filename, *path_titles[:-1]]) diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index 839b61654..ae11c7750 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -58,7 +58,6 @@ def _page_feature(page: int = 1) -> PageFeature: height=72.0, has_asset=False, is_blank_like=False, - asset_bboxes=None, ) @@ -158,6 +157,7 @@ def test_run_lightweight_anatomy_builds_single_shard_without_planner_llm( ) assert list(anatomy_data)[:2] == ["version", "toc_hierarchies"] assert "text_lines_preview" not in anatomy_data["page_features"][0] + assert "asset_bboxes" not in anatomy_data["page_features"][0] trace_data = json.loads((output_dir / "trace.json").read_text(encoding="utf-8")) assert "visual_stages" in trace_data["summary"]["budget"] diff --git a/apps/worker/tests/contract/test_profile_agent_protocol_contract.py b/apps/worker/tests/contract/test_profile_agent_protocol_contract.py new file mode 100644 index 000000000..b16233a5e --- /dev/null +++ b/apps/worker/tests/contract/test_profile_agent_protocol_contract.py @@ -0,0 +1,214 @@ +"""Protocol tests: planner next_action and executor finish ownership.""" + +from __future__ import annotations + +import json +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from app.services.document_agent.budget import BudgetTracker +from app.services.document_agent.executor.react_loop import ( + ReActExecutor, + _parse_decision, +) +from app.services.document_agent.manifest import ( + DocumentProfile, + ReflexionDecision, + ToolContext, +) +from app.services.document_agent.planner.planner import _parse_profile_and_decision +from app.services.document_agent.registry import REGISTRY +from app.services.document_agent import tools as _registered_tools # noqa: F401 +from app.services.document_agent.state import AgentBlackboard + + +def test_planner_verdict_now_falls_through_to_ready_to_shard() -> None: + raw = json.dumps( + { + "is_scanned": True, + "category": "Feasibility Study Report", + "routing_category": "generic", + "category_rationale": "scanned prose", + "language": "zh", + "rationale": "scanned PDF not atlas", + "header_y": None, + "footer_y": None, + "next_action": "verdict_now", + "inspect_pages": [], + "grep_query": "", + } + ) + profile, decision = _parse_profile_and_decision(raw) + assert profile.is_scanned is True + assert decision.action == "tool_call" + assert decision.tool_name == "propose.shard_plan" + assert decision.verdict is None + + +def test_planner_ready_to_shard_proposes_shard_plan() -> None: + raw = json.dumps( + { + "is_scanned": False, + "category": "Report", + "routing_category": "generic", + "language": "en", + "rationale": "enough evidence", + "next_action": "ready_to_shard", + } + ) + _profile, decision = _parse_profile_and_decision(raw) + assert decision.action == "tool_call" + assert decision.tool_name == "propose.shard_plan" + + +def test_planner_inspect_more_maps_to_inspect_pages() -> None: + raw = json.dumps( + { + "is_scanned": False, + "category": "Report", + "routing_category": "generic", + "language": "en", + "rationale": "need more pages", + "next_action": "inspect_more", + "inspect_pages": [3, 8], + } + ) + _profile, decision = _parse_profile_and_decision(raw) + assert decision.tool_name == "inspect.pages" + assert decision.tool_args["pages"] == [3, 8] + + +def test_executor_legacy_verdict_now_without_status_becomes_shard() -> None: + decision = _parse_decision( + json.dumps( + { + "action": "verdict_now", + "rationale": "classification done", + } + ) + ) + assert decision.action == "tool_call" + assert decision.tool_name == "propose.shard_plan" + + +def test_executor_legacy_verdict_now_with_abort_status_uses_verdict_tool() -> None: + decision = _parse_decision( + json.dumps( + { + "action": "verdict_now", + "rationale": "cannot profile", + "verdict": {"status": "abort", "rationale": "cannot profile"}, + } + ) + ) + assert decision.action == "tool_call" + assert decision.tool_name == "verdict" + assert decision.tool_args["status"] == "abort" + + +def _seed_pages(blackboard: AgentBlackboard, page_count: int) -> None: + from app.services.document_agent.manifest import PageFeature, PageLabel + + blackboard.page_count = page_count + blackboard.doc_stats = {"page_count": page_count} + blackboard.page_features = [ + PageFeature( + page=page, + raw_text_length=0, + text_density=0.0, + image_coverage=1.0, + image_count=1, + table_count=0, + drawings_count=0, + orientation="portrait", + width=612.0, + height=792.0, + has_asset=True, + is_blank_like=True, + ) + for page in range(1, page_count + 1) + ] + blackboard.page_labels = [ + PageLabel(page=page, kind="normal", confidence=0.9) + for page in range(1, page_count + 1) + ] + + +def test_executor_initial_ready_to_shard_reaches_success_without_abort() -> None: + blackboard = AgentBlackboard() + _seed_pages(blackboard, 4) + blackboard.document_profile = DocumentProfile( + is_scanned=True, + category="Feasibility Study Report", + routing_category="generic", + rationale="scanned PDF not atlas", + ) + from app.services.document_agent.manifest import TocResult + + blackboard.toc_result = TocResult(method="none", notes="no toc") + ctx = ToolContext( + pdf_path="/tmp/scanned.pdf", + job_id="job-scanned", + blackboard=blackboard, + budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), + trace=None, + settings={}, # deterministic executor (no LLM) + ) + initial = ReflexionDecision( + action="tool_call", + rationale="ready", + tool_name="propose.shard_plan", + tool_args={}, + ) + result = ReActExecutor( + ctx, + registry=REGISTRY, + max_rounds=10, + initial_decision=initial, + ).run() + assert result.verdict.status == "success" + assert blackboard.shard_plan is not None + assert len(blackboard.shard_plan.shards) >= 1 + + +def test_executor_empty_initial_tool_falls_through_to_success() -> None: + """Missing tool_name must coerce to propose.shard_plan, not abort.""" + blackboard = AgentBlackboard() + _seed_pages(blackboard, 3) + from app.services.document_agent.manifest import TocResult + + blackboard.toc_result = TocResult(method="none", notes="no toc") + blackboard.document_profile = DocumentProfile( + is_scanned=True, + category="Report", + routing_category="generic", + ) + ctx = ToolContext( + pdf_path="/tmp/scanned.pdf", + job_id="job-legacy", + blackboard=blackboard, + budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), + trace=None, + settings={}, + ) + initial = ReflexionDecision( + action="tool_call", + rationale="stale empty decision", + tool_name=None, + tool_args={}, + ) + result = ReActExecutor( + ctx, + registry=REGISTRY, + max_rounds=10, + initial_decision=initial, + ).run() + assert result.verdict.status == "success" + assert blackboard.shard_plan is not None + assert len(blackboard.shard_plan.shards) == 1 diff --git a/apps/worker/tests/contract/test_toc_confirm_batch_contract.py b/apps/worker/tests/contract/test_toc_confirm_batch_contract.py new file mode 100644 index 000000000..f92992519 --- /dev/null +++ b/apps/worker/tests/contract/test_toc_confirm_batch_contract.py @@ -0,0 +1,126 @@ +"""Contract tests for TOC Phase-1 batched VLM anchor confirmation.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +from app.services.document_agent.budget import BudgetTracker, StageEnvelope +from app.services.document_agent.manifest import TocAnchorPage +from app.services.document_agent.tools import extract_toc_with_boundaries as toc_tool + + +def _anchors(tmp_path: Path, pages: list[int]) -> list[TocAnchorPage]: + out: list[TocAnchorPage] = [] + for page in pages: + png = tmp_path / f"toc_anchor_page_{page}.png" + png.write_bytes(b"fake-png") + out.append( + TocAnchorPage( + page=page, + png_path=str(png), + source="text_scan", + ) + ) + return out + + +def test_iter_chunks_uses_boundary_step_size() -> None: + items = list(range(12)) + chunks = toc_tool._iter_chunks(items, toc_tool.BOUNDARY_STEP_PAGES) # noqa: SLF001 + assert chunks == [ + [0, 1, 2, 3, 4], + [5, 6, 7, 8, 9], + [10, 11], + ] + + +def test_vlm_confirm_anchors_batches_and_merges_partial_failures( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """One confirm chunk may fail; successful chunks still contribute confirmed pages.""" + anchors = _anchors(tmp_path, [5, 11, 40, 44, 55, 78, 97]) + assert toc_tool.BOUNDARY_STEP_PAGES == 5 + # 7 anchors → 2 chunks: [5..55] and [78, 97] + + call_pages: list[list[int]] = [] + + class _FakeClient: + def chat_completion_with_usage(self, **kwargs: Any) -> tuple[str, dict[str, int]]: + content = kwargs["messages"][0]["content"] + pages = [] + for part in content: + if not isinstance(part, dict) or part.get("type") != "text": + continue + text = str(part.get("text") or "") + if text.startswith("\n--- Page ") and text.endswith(" ---"): + pages.append(int(text[len("\n--- Page ") : -len(" ---")])) + call_pages.append(pages) + if 78 in pages: + raise RuntimeError("simulated truncated JSON") + payload = { + "pages": [ + { + "page": page, + "is_toc_start": page in {5, 11, 40}, + "reason": "ok", + } + for page in pages + ] + } + return json.dumps(payload), {"total_tokens": 100} + + monkeypatch.setattr( + "shared.services.ai.llm_overrides.get_vision_client", + lambda requested_model=None: (_FakeClient(), requested_model or "fake-vlm"), + ) + + budget = BudgetTracker( + plan_budget=50_000, + visual_budget=200_000, + visual_stage_envelopes={ + "toc_confirm": StageEnvelope(min_guarantee=0, cap=None), + }, + ) + confirmed, confirm_failed, evidence = toc_tool._vlm_confirm_anchors( # noqa: SLF001 + anchors, + model="fake-vlm", + budget=budget, + ) + + assert sorted(call_pages[0] + call_pages[1]) == [5, 11, 40, 44, 55, 78, 97] + assert {tuple(pages) for pages in call_pages} == { + (5, 11, 40, 44, 55), + (78, 97), + } + assert confirm_failed is False + assert [a.page for a in confirmed] == [5, 11, 40] + assert {e.page_index for e in evidence} == {5, 11, 40, 44, 55, 78, 97} + + +def test_vlm_confirm_anchors_all_chunks_fail_sets_confirm_failed( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + anchors = _anchors(tmp_path, [5, 11]) + + class _FakeClient: + def chat_completion_with_usage(self, **kwargs: Any) -> tuple[str, dict[str, int]]: + raise RuntimeError("boom") + + monkeypatch.setattr( + "shared.services.ai.llm_overrides.get_vision_client", + lambda requested_model=None: (_FakeClient(), requested_model or "fake-vlm"), + ) + + confirmed, confirm_failed, _evidence = toc_tool._vlm_confirm_anchors( # noqa: SLF001 + anchors, + model="fake-vlm", + budget=None, + ) + assert confirmed == [] + assert confirm_failed is True diff --git a/apps/worker/tests/contract/test_toc_phase2_region_concurrency_contract.py b/apps/worker/tests/contract/test_toc_phase2_region_concurrency_contract.py new file mode 100644 index 000000000..a068b0b88 --- /dev/null +++ b/apps/worker/tests/contract/test_toc_phase2_region_concurrency_contract.py @@ -0,0 +1,276 @@ +"""Contract tests for TOC Phase-2 concurrent VLM with serial batch renders.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import gevent +import pytest + +from app.services.document_agent.manifest import TocAnchorPage +from app.services.document_agent.tools import extract_toc_with_boundaries as toc_tool +from app.services.document_agent.tools.vlm_toc_extractor import ( + BatchPageResult, + BatchTocResult, +) + + +def _anchors(tmp_path: Path, pages: list[int]) -> list[TocAnchorPage]: + out: list[TocAnchorPage] = [] + for page in pages: + png = tmp_path / f"toc_anchor_page_{page}.png" + png.write_bytes(b"anchor-png") + out.append( + TocAnchorPage(page=page, png_path=str(png), source="vlm") + ) + return out + + +def _anchor_from_png(png_path: str) -> int: + name = Path(png_path).name + if name.startswith("toc_anchor_page_"): + return int(name[len("toc_anchor_page_") : -len(".png")]) + # toc_a{anchor}_p{page}.png + stem = name.removesuffix(".png") + anchor_part, _sep, _page_part = stem.partition("_p") + assert anchor_part.startswith("toc_a"), png_path + return int(anchor_part[len("toc_a") :]) + + +def _batch_result( + *, + toc_pages: list[int], + non_toc_pages: list[int], + entries: list[dict[str, Any]] | None = None, +) -> BatchTocResult: + page_results: list[BatchPageResult] = [ + BatchPageResult(page=page, is_toc=True, entries=[]) for page in toc_pages + ] + page_results.extend( + BatchPageResult(page=page, is_toc=False, entries=[]) for page in non_toc_pages + ) + if page_results and page_results[-1].is_toc: + page_results[-1] = BatchPageResult( + page=page_results[-1].page, + is_toc=False, + entries=[], + ) + return BatchTocResult( + page_results=page_results, + toc_pages=list(toc_pages), + non_toc_pages=list(non_toc_pages), + all_entries=list(entries or []), + meta={"ok": True}, + ) + + +def _fake_batch_render(worker_fn: Any, *args: Any, **kwargs: Any) -> dict[str, Any]: + """Match ``_render_expand_window_worker`` args: pages list in one spawn.""" + pages = list(args[1]) + output_dir = Path(args[2]) + anchor_page = int(args[4]) + output_dir.mkdir(parents=True, exist_ok=True) + results: list[dict[str, Any]] = [] + for page_num in pages: + png_path = output_dir / f"toc_a{anchor_page}_p{page_num}.png" + png_path.write_bytes(b"png") + results.append({"page": page_num, "png_path": str(png_path)}) + return {"ok": True, "results": results} + + +def test_extract_regions_runs_per_anchor_and_merges_in_page_order( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + confirmed = _anchors(tmp_path, [10, 40, 80]) + calls: list[int] = [] + render_calls: list[list[int]] = [] + + def _tracking_render(worker_fn: Any, *args: Any, **kwargs: Any) -> dict[str, Any]: + render_calls.append(list(args[1])) + return _fake_batch_render(worker_fn, *args, **kwargs) + + def _fake_batch( + *, + page_pngs: list[tuple[int, str]], + model: str, + previous_entries: list[dict[str, Any]] | None = None, + ) -> BatchTocResult: + first_page = page_pngs[0][0] + anchor = _anchor_from_png(page_pngs[0][1]) + calls.append(anchor) + assert previous_entries in (None, []) + return _batch_result( + toc_pages=[first_page], + non_toc_pages=[p for p, _ in page_pngs[1:]], + entries=[ + { + "title": f"Section {anchor}", + "page": first_page, + "level": 1, + } + ], + ) + + monkeypatch.setattr(toc_tool, "run_in_child_process", _tracking_render) + monkeypatch.setattr( + "app.services.document_agent.tools.vlm_toc_extractor.vlm_extract_toc_batch", + _fake_batch, + ) + monkeypatch.setattr( + toc_tool, + "vlm_entries_to_toc_hierarchies", + lambda entries, **kwargs: [ + { + "toc_range": [entries[0]["page"], entries[0]["page"]], + "toc_tree": {entries[0]["title"]: {}}, + } + ], + ) + + regions = toc_tool._extract_regions_for_confirmed_anchors( # noqa: SLF001 + confirmed, + pdf_path="/tmp/doc.pdf", + page_count=100, + output_dir=str(tmp_path / "toc_pages"), + dpi=72, + model="fake-vlm", + ) + + assert [r.anchor_page for r in regions] == [10, 40, 80] + assert all(r.error is None for r in regions) + assert sorted(calls) == [10, 40, 80] + # One spawn per window; start page reused from Phase1 anchor PNG. + assert sorted(render_calls) == [ + [11, 12, 13, 14], + [41, 42, 43, 44], + [81, 82, 83, 84], + ] + + +def test_extract_regions_keeps_success_when_one_anchor_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + confirmed = _anchors(tmp_path, [10, 40]) + + def _fake_batch( + *, + page_pngs: list[tuple[int, str]], + model: str, + previous_entries: list[dict[str, Any]] | None = None, + ) -> BatchTocResult: + anchor = _anchor_from_png(page_pngs[0][1]) + if anchor == 40: + raise RuntimeError("simulated region VLM failure") + first_page = page_pngs[0][0] + return _batch_result( + toc_pages=[first_page], + non_toc_pages=[p for p, _ in page_pngs[1:]], + entries=[{"title": "Intro", "page": first_page, "level": 1}], + ) + + monkeypatch.setattr(toc_tool, "run_in_child_process", _fake_batch_render) + monkeypatch.setattr( + "app.services.document_agent.tools.vlm_toc_extractor.vlm_extract_toc_batch", + _fake_batch, + ) + monkeypatch.setattr( + toc_tool, + "vlm_entries_to_toc_hierarchies", + lambda entries, **kwargs: [ + { + "toc_range": [entries[0]["page"], entries[0]["page"]], + "toc_tree": {entries[0]["title"]: {}}, + } + ], + ) + + regions = toc_tool._extract_regions_for_confirmed_anchors( # noqa: SLF001 + confirmed, + pdf_path="/tmp/doc.pdf", + page_count=100, + output_dir=str(tmp_path / "toc_pages"), + dpi=72, + model="fake-vlm", + ) + + assert regions[0].anchor_page == 10 + assert regions[0].error is None + assert regions[0].entries + assert regions[1].anchor_page == 40 + assert regions[1].error is not None + assert "simulated region VLM failure" in regions[1].error + + +def test_phase2_serial_batch_render_with_concurrent_vlm( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Batch renders never overlap; VLM calls may overlap across anchors.""" + confirmed = _anchors(tmp_path, [10, 40, 80]) + render_active = 0 + render_max = 0 + vlm_active = 0 + vlm_max = 0 + spawn_count = 0 + + def _fake_render(worker_fn: Any, *args: Any, **kwargs: Any) -> dict[str, Any]: + nonlocal render_active, render_max, spawn_count + spawn_count += 1 + render_active += 1 + render_max = max(render_max, render_active) + gevent.sleep(0.02) + result = _fake_batch_render(worker_fn, *args, **kwargs) + render_active -= 1 + return result + + def _fake_batch( + *, + page_pngs: list[tuple[int, str]], + model: str, + previous_entries: list[dict[str, Any]] | None = None, + ) -> BatchTocResult: + nonlocal vlm_active, vlm_max + vlm_active += 1 + vlm_max = max(vlm_max, vlm_active) + gevent.sleep(0.2) + first_page = page_pngs[0][0] + anchor = _anchor_from_png(page_pngs[0][1]) + vlm_active -= 1 + return _batch_result( + toc_pages=[first_page], + non_toc_pages=[p for p, _ in page_pngs[1:]], + entries=[{"title": f"S{anchor}", "page": first_page, "level": 1}], + ) + + monkeypatch.setattr(toc_tool, "run_in_child_process", _fake_render) + monkeypatch.setattr( + "app.services.document_agent.tools.vlm_toc_extractor.vlm_extract_toc_batch", + _fake_batch, + ) + monkeypatch.setattr( + toc_tool, + "vlm_entries_to_toc_hierarchies", + lambda entries, **kwargs: [ + { + "toc_range": [entries[0]["page"], entries[0]["page"]], + "toc_tree": {entries[0]["title"]: {}}, + } + ], + ) + + regions = toc_tool._extract_regions_for_confirmed_anchors( # noqa: SLF001 + confirmed, + pdf_path="/tmp/doc.pdf", + page_count=100, + output_dir=str(tmp_path / "toc_pages"), + dpi=72, + model="fake-vlm", + ) + + assert all(r.error is None for r in regions) + assert render_max == 1 + assert spawn_count == 3 # one batch spawn per anchor window + assert vlm_max >= 2 diff --git a/packages/shared-python/shared/services/storage/zip_package_writer.py b/packages/shared-python/shared/services/storage/zip_package_writer.py index 34ec35477..9d2361834 100644 --- a/packages/shared-python/shared/services/storage/zip_package_writer.py +++ b/packages/shared-python/shared/services/storage/zip_package_writer.py @@ -26,6 +26,9 @@ class ZipPackageWriteRequest: doc_nav: dict[str, Any] | None manifest: dict[str, Any] temp_dir: str | None + # Legacy chunk/docx/md packages may ship toc_hierarchies.json; page_memory + # keeps human TOC as debug-root toc_hierarchy.json and omits the list dump. + include_toc_hierarchies: bool = True @dataclass(frozen=True) @@ -52,7 +55,7 @@ def write(self, request: ZipPackageWriteRequest) -> ZipPackageArtifact: zip_file.writestr("chunks.json", chunks_json.encode("utf-8")) self._write_optional_file(zip_file, request.add_dir, "full.md") - if self._write_optional_file( + if request.include_toc_hierarchies and self._write_optional_file( zip_file, request.add_dir, "toc_hierarchies.json", diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py index 8d6b229a2..c1808002f 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -83,6 +83,7 @@ def generate_zip_package( job_metadata=job_metadata, hierarchy=hierarchy, ) + parse_track = str((job_metadata or {}).get("parse_track") or "") artifact = self._writer.write( ZipPackageWriteRequest( job_id=job_id, @@ -94,6 +95,7 @@ def generate_zip_package( doc_nav=doc_nav, manifest=manifest, temp_dir=temp_dir, + include_toc_hierarchies=parse_track != "page_memory", ) )