From 62b6b081422785f2497ebd2c4140aefec27d60ad Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 21 Jul 2026 16:32:05 +0800 Subject: [PATCH 1/6] feat: enhance page memory serialization and title detection - Added `observed_titles` to the `serialize_page_tags` function for improved title tracking. - Introduced `serialize_scope_skeletons` to handle scope input artifacts, facilitating Stage3 to Stage4 handoff. - Updated `write_scope_artifacts` to conditionally write `tags` and `assets_by_page`, preventing overwrites when set to None. - Enhanced title detection logic in `tag_vlm_titles` to escalate token budgets dynamically when JSON responses are truncated, improving robustness in title extraction. This commit improves the handling of page memory artifacts and enhances the title detection process, ensuring better data integrity and extraction accuracy. --- .../services/page_memory/_serialization.py | 56 ++++- .../services/page_memory/memory_service.py | 14 +- .../app/services/page_memory/page_tagger.py | 204 ++++++++++++------ .../test_page_memory_page_tagger_contract.py | 83 ++++++- 4 files changed, 278 insertions(+), 79 deletions(-) diff --git a/apps/worker/app/services/page_memory/_serialization.py b/apps/worker/app/services/page_memory/_serialization.py index a963076c..87caea10 100644 --- a/apps/worker/app/services/page_memory/_serialization.py +++ b/apps/worker/app/services/page_memory/_serialization.py @@ -149,6 +149,7 @@ def serialize_page_tags(tags: list[Any]) -> list[dict[str, Any]]: "keywords": list(item.keywords), "entities": list(getattr(item, "entities", []) or []), "strategy_used": item.strategy_used, + "observed_titles": list(getattr(item, "observed_titles", []) or []), } for item in tags ] @@ -178,23 +179,68 @@ def serialize_assets(assets_by_page: dict[int, list[Any]]) -> list[dict[str, Any return rows +def serialize_scope_skeletons( + *, + scope_id: str, + start_page: int, + end_page: int, + strategy: str, + skeletons: list[Any], +) -> dict[str, Any]: + """Coarse scope input artifact (Stage3 → Stage4 handoff). + + Closed-closed ``start_page``/``end_page`` plus coarse skeleton rows + (including ``evidence``). Downstream stages read this file; refined + hierarchy lives in ``fine_hierarchy.json``. + """ + start = max(1, int(start_page)) + end = max(start, int(end_page)) + rows = [ + { + "section_path": getattr(item, "section_path", ""), + "title": getattr(item, "title", ""), + "level": int(getattr(item, "level", 0) or 0), + "start_page": int(getattr(item, "start_page", 0) or 0), + "end_page": int(getattr(item, "end_page", 0) or 0), + "parent_path": getattr(item, "parent_path", None), + "evidence": dict(getattr(item, "evidence", {}) or {}), + } + for item in skeletons + ] + return { + "scope_id": scope_id, + "start_page": start, + "end_page": end, + "page_count": end - start + 1, + "strategy": strategy, + "skeleton_count": len(rows), + "skeletons": rows, + } + + def write_scope_artifacts( *, output_dir: str, scope_id: str, scope_manifest_data: dict[str, Any], hierarchy: list[Any], - tags: list[Any], + tags: list[Any] | None = None, assets_by_page: dict[int, list[Any]] | None = None, ) -> None: + """Write per-scope viewing artifacts. + + Always refreshes ``fine_hierarchy.json`` (embeds ``scope`` manifest). + ``tags`` / ``assets_by_page`` of ``None`` leave the existing file untouched + so later stages do not wipe earlier placeholders or results. + """ scope_dir = Path(output_dir) / "scopes" / scope_id - write_json(scope_dir / "scope.json", scope_manifest_data) write_json( scope_dir / "fine_hierarchy.json", serialize_hierarchy_artifact(hierarchy, scope_manifest_data=scope_manifest_data), ) - write_json(scope_dir / "page_tags.json", serialize_page_tags(tags)) - if assets_by_page: + if tags is not None: + write_json(scope_dir / "page_tags.json", serialize_page_tags(tags)) + if assets_by_page is not None: write_json(scope_dir / "assets.json", serialize_assets(assets_by_page)) @@ -208,7 +254,7 @@ def write_top_level_artifacts( root = Path(output_dir) write_json(root / "hierarchy.json", serialize_hierarchy_artifact(hierarchy)) write_json(root / "page_tags.json", serialize_page_tags(tags)) - if assets_by_page: + if assets_by_page is not None: write_json(root / "assets.json", serialize_assets(assets_by_page)) else: (root / "assets.json").unlink(missing_ok=True) diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index f050f631..22008909 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -799,6 +799,7 @@ def _run_hierarchy_scope( }, }, ) + # Refresh hierarchy/tags and write assets without wiping unrelated slots. _write_scope_artifacts( output_dir=output_dir, scope_id=scope.scope_id, @@ -878,34 +879,29 @@ def _record_trace_stage( def _cleanup_page_memory_artifacts(output_dir: str) -> None: root = Path(output_dir) - legacy_files = { + stale_files = { "assets.json", "chunks.json", - "coarse_tag_scope.json", + "coarse_scopes.json", "doc_nav.json", "hierarchy.json", "manifest.json", "node_rows.csv", "node_rows.json", - "page_memory_fine_hierarchy.json", "page_plans.json", "page_rendered.json", "page_tags.json", - "page_tags_after_titles.json", - "page_tags_pre_hierarchy.json", "report.md", - "skeletons.json", - "tag_scope.json", "trace.json", } - for name in legacy_files: + for name in stale_files: path = root / name try: if path.is_file(): path.unlink() except Exception: logger.debug("[page_memory] failed to cleanup artifact {}", path) - for name in ("asset_annotate", "debug", "fine_hierarchy", "images", "pages", "scopes", "tables"): + for name in ("asset_annotate", "debug", "images", "pages", "scopes", "tables"): path = root / name try: if path.is_dir(): diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index 6e2a2aab..4867d6a4 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -49,6 +49,8 @@ class PageTagResult: _MAX_JSON_RETRIES = 1 _DEFAULT_FINE_MIN_PAGES = 4 +# Dense section-start pages can emit long title JSON; escalate only on truncation. +_TITLE_TOKEN_BUDGETS: tuple[int, ...] = (300, 600, 1200) def tag_pages( @@ -367,19 +369,92 @@ def _detect_one( return tag_results +def _completion_tokens(usage: Any) -> int: + if isinstance(usage, dict): + return int(usage.get("completion_tokens") or 0) + return int(getattr(usage, "completion_tokens", 0) or 0) + + +def _title_response_truncated( + raw_response: str, + *, + usage: Any, + max_tokens: int, +) -> bool: + """True when the completion likely hit the budget mid-JSON.""" + if max_tokens > 0 and _completion_tokens(usage) >= max_tokens: + return True + stripped = (raw_response or "").rstrip() + if not stripped: + return False + return not stripped.endswith("}") + + +def _parse_observed_titles( + raw_response: str, + *, + page_index: int, +) -> list[dict[str, Any]]: + data = json.loads(raw_response) + titles_raw = data.get("titles", []) + if not isinstance(titles_raw, list): + return [] + + observed: list[dict[str, Any]] = [] + for item in titles_raw: + if not isinstance(item, dict) or not item.get("text"): + continue + text = str(item["text"]).strip() + is_table = item.get("is_in_table") is True + is_header = item.get("is_in_header_footer") is True + if is_table or is_header: + logger.debug( + "[page_tagger] filtered CoT title on page {}: '{}' (table={}, header={})", + page_index, + text, + is_table, + is_header, + ) + continue + if not text: + continue + prominence = None + try: + prominence = float(item.get("prominence", 0.5)) + except (TypeError, ValueError) as exc: + logger.debug( + "[page_tagger] ignored non-numeric title prominence {}: {}", + item.get("prominence"), + exc, + ) + observed.append( + { + "text": text, + "prominence": prominence, + "is_in_table": is_table, + "is_in_header_footer": is_header, + } + ) + return observed + + def _tag_vlm_titles( page: PageRenderResult, *, model: str, scan_direction: str = "top_to_bottom_left_to_right", ) -> list[dict[str, Any]]: - """Send page PNG to VLM with the title-only prompt and parse results.""" - prompt, temperature, _top_p, max_tokens = build_prompt( + """Send page PNG to VLM with the title-only prompt and parse results. + + Starts at a small completion budget and escalates only when the response + is truncated (budget hit / incomplete JSON that fails to parse). + """ + prompt, temperature, _top_p, _default_max_tokens = build_prompt( "page-memory-vlm-title", "", "", paras={ - "max_tokens": 300, + "max_tokens": _TITLE_TOKEN_BUDGETS[0], "scan_direction": scan_direction, }, ) @@ -407,70 +482,71 @@ def _tag_vlm_titles( client, resolved_model = get_vision_client(requested_model=model) model = resolved_model or model - for attempt in range(_MAX_JSON_RETRIES + 1): - try: - raw_response, usage = client.chat_completion_with_usage( - messages=cast(Any, [{"role": "user", "content": content_parts}]), - model=model, - temperature=temperature, - max_tokens=max_tokens, - response_format={"type": "json_object"}, - usage_task="page_memory.title_detection", - ) + last_truncated = False + for budget_index, max_tokens in enumerate(_TITLE_TOKEN_BUDGETS): + for attempt in range(_MAX_JSON_RETRIES + 1): + try: + raw_response, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=model, + temperature=temperature, + max_tokens=max_tokens, + response_format={"type": "json_object"}, + usage_task="page_memory.title_detection", + ) + except UnavailableException: + raise + except Exception as exc: + logger.warning( + "[page_tagger] title VLM failed for page {}: {}", + page.page_index, + exc, + ) + return [] - data = json.loads(raw_response) - titles_raw = data.get("titles", []) - if not isinstance(titles_raw, list): + try: + observed = _parse_observed_titles( + raw_response, + page_index=page.page_index, + ) + except json.JSONDecodeError: + truncated = _title_response_truncated( + raw_response, + usage=usage, + max_tokens=max_tokens, + ) + if truncated and budget_index + 1 < len(_TITLE_TOKEN_BUDGETS): + last_truncated = True + logger.info( + "[page_tagger] title JSON truncated on page {} " + "(budget={}, completion_tokens={}); escalating", + page.page_index, + max_tokens, + _completion_tokens(usage), + ) + break # next budget + if attempt < _MAX_JSON_RETRIES: + continue + logger.warning( + "[page_tagger] title JSON retry exhausted for page {}", + page.page_index, + ) return [] - observed: list[dict[str, Any]] = [] - for item in titles_raw: - if isinstance(item, dict) and item.get("text"): - text = str(item["text"]).strip() - - is_table = item.get("is_in_table") is True - is_header = item.get("is_in_header_footer") is True - - if is_table or is_header: - logger.debug( - "[page_tagger] filtered CoT title on page {}: '{}' (table={}, header={})", - page.page_index, text, is_table, is_header - ) - continue - - if text: - prominence = None - try: - prominence = float(item.get("prominence", 0.5)) - except (TypeError, ValueError) as exc: - logger.debug( - "[page_tagger] ignored non-numeric title prominence {}: {}", - item.get("prominence"), - exc, - ) - observed.append({ - "text": text, - "prominence": prominence, - "is_in_table": is_table, - "is_in_header_footer": is_header - }) + if last_truncated: + logger.info( + "[page_tagger] title detection recovered on page {} with budget={}", + page.page_index, + max_tokens, + ) return observed + else: + continue - except json.JSONDecodeError: - if attempt < _MAX_JSON_RETRIES: - continue - logger.warning( - "[page_tagger] title JSON retry exhausted for page {}", - page.page_index, - ) - return [] - except UnavailableException: - raise - except Exception as exc: - logger.warning( - "[page_tagger] title VLM failed for page {}: {}", - page.page_index, exc, - ) - return [] - + if last_truncated: + logger.warning( + "[page_tagger] title JSON still truncated on page {} after budgets {}", + page.page_index, + list(_TITLE_TOKEN_BUDGETS), + ) return [] diff --git a/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py b/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py index 22b5104e..ae45c0a5 100644 --- a/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py +++ b/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py @@ -12,7 +12,13 @@ os.environ.setdefault("S3_TEMP_PATH", "/tmp") from app.services.page_memory.page_renderer import PageRenderResult -from app.services.page_memory.page_tagger import PageTagResult, tag_page_titles +from app.services.page_memory.page_tagger import ( + PageTagResult, + _completion_tokens, + _tag_vlm_titles, + _title_response_truncated, + tag_page_titles, +) from shared.core.exceptions.domain_exceptions import UnavailableException @@ -22,6 +28,81 @@ def _write_page_image(tmp_path, page_index: int) -> str: return str(image_path) +def test_title_response_truncated_detects_budget_hit_and_incomplete_json() -> None: + assert _title_response_truncated( + '{"titles":[', + usage={"completion_tokens": 300}, + max_tokens=300, + ) + assert _title_response_truncated( + '{"titles":[{"text":"A"', + usage={"completion_tokens": 120}, + max_tokens=300, + ) + assert not _title_response_truncated( + '{"titles":[]}', + usage={"completion_tokens": 40}, + max_tokens=300, + ) + assert _completion_tokens({"completion_tokens": 12}) == 12 + + +def test_title_detection_escalates_token_budget_on_truncated_json( + monkeypatch, + tmp_path, +) -> None: + truncated = ( + '{\n "titles": [\n' + " {\n" + ' "text": "Section A Governing requirements",\n' + ' "prominence": 1.0,\n' + ' "is_in_table": false,\n' + ) + complete = ( + '{\n "titles": [\n' + " {\n" + ' "text": "Section A Governing requirements",\n' + ' "prominence": 1.0,\n' + ' "is_in_table": false,\n' + ' "is_in_header_footer": false\n' + " }\n" + " ]\n" + "}" + ) + calls: list[int] = [] + + class _FakeClient: + def chat_completion_with_usage(self, **kwargs): + max_tokens = int(kwargs["max_tokens"]) + calls.append(max_tokens) + if max_tokens == 300: + return truncated, {"completion_tokens": 300, "prompt_tokens": 10} + return complete, {"completion_tokens": 180, "prompt_tokens": 10} + + monkeypatch.setattr( + "shared.services.ai.llm_overrides.get_vision_client", + lambda requested_model=None: (_FakeClient(), requested_model or "fake-vlm"), + ) + monkeypatch.setattr( + "app.services.page_memory.page_tagger.build_prompt", + lambda *args, **kwargs: ("prompt", 0.0, 0.01, 300), + ) + + page = PageRenderResult( + page_index=38, + image_path=_write_page_image(tmp_path, 38), + raw_text="", + width=100, + height=200, + is_landscape=False, + ) + observed = _tag_vlm_titles(page, model="fake-vlm") + assert calls == [300, 600] + assert [item["text"] for item in observed] == [ + "Section A Governing requirements" + ] + + def test_title_detection_preserves_page_index_assignment_under_concurrency( monkeypatch, tmp_path, From e1fb17e6e818cc02dd4ff37cca5aa9c5463cc039 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 21 Jul 2026 22:45:35 +0800 Subject: [PATCH 2/6] refactor: remove experimental chart and VLM bbox probing scripts - Deleted `chart_asset_probe.py` and `vlm_bbox_tabula_probe.py` as they are no longer needed. - These scripts were used for experimental purposes related to VLM-driven asset detection and table extraction, but have been superseded by more efficient implementations. This cleanup helps streamline the codebase by removing outdated experimental files. --- apps/worker/experiments/chart_asset_probe.py | 470 ------------------ .../experiments/vlm_bbox_tabula_probe.py | 294 ----------- 2 files changed, 764 deletions(-) delete mode 100644 apps/worker/experiments/chart_asset_probe.py delete mode 100644 apps/worker/experiments/vlm_bbox_tabula_probe.py diff --git a/apps/worker/experiments/chart_asset_probe.py b/apps/worker/experiments/chart_asset_probe.py deleted file mode 100644 index 016e08ca..00000000 --- a/apps/worker/experiments/chart_asset_probe.py +++ /dev/null @@ -1,470 +0,0 @@ -"""Experimental: VLM-driven asset bbox detection + cropping. - -Goal of this experiment ------------------------ -Test whether a VLM can directly locate table/chart/figure bounding boxes on a -*rendered page image* (the same way PAGE-TRACK renders pages), so we can -**crop** those regions out and later hand the crop to a dedicated table -model (e.g. tabular / table-transformer) instead of asking the VLM to -transcribe the whole table verbatim (error-prone + expensive output). - -This script does NOT touch production code. It only: - - 1. Renders selected PDF pages to PNG at a fixed DPI (mirrors PAGE-TRACK, - default 144 DPI) so the pixel<->point mapping is fully controlled. - 2. (VLM) Asks the model only for table/chart/figure regions as normalized - [0,1000] boxes, maps them to pixels, crops, and draws an annotated overlay. - -Outputs (under --out): - pages/page-N.png full page render - crops/page-N_vlm-K_.png VLM crops - crops/page-N_ref-K_.png reference crops from existing chunks - asset_annotate/page_N.png page with VLM(red) + reference(green) boxes - results.json all regions + metadata - report.md human-readable summary - -Run: - cd apps/worker - uv run python experiments/chart_asset_probe.py \ - --pdf "/path/to/doc.pdf" \ - --pages all - # Uses $IMAGE_MODEL (default qwen3.6-flash) unless --model is set. - # Alternate: --model qwen3-vl-32b-instruct (open-weights, local-deployable). - -VLM model guidance (bbox grounding) ------------------------------------------------- -* **Default (cloud):** ``qwen3.6-flash`` via ``$IMAGE_MODEL`` — cheapest, - strong bbox quality on our probe PDFs. -* **Alternate (cloud or self-hosted):** ``qwen3-vl-32b-instruct`` — open - weights, can be deployed locally (vLLM / SGLang / Ollama); DashScope - China pricing (2026-04): input ¥2 / output ¥8 per 1M tokens (see - https://help.aliyun.com/zh/model-studio/model-pricing ). -* Coordinates are requested in a normalized 0-1000 space to be robust - to whatever internal resize the API performs. -""" - -from __future__ import annotations - -import argparse -import base64 -import json -import os -import re -import sys -from dataclasses import dataclass, field, asdict -from pathlib import Path -from typing import Any - -import fitz # PyMuPDF -from PIL import Image, ImageDraw, ImageFont - - -# ── coordinate convention ───────────────────────────────────────────── -# We ask the VLM for boxes in a normalized integer space [0, 1000] for -# BOTH axes, with origin at the top-left of the page image. This is robust -# to whatever internal resize the API performs. -NORM = 1000 -_VALID_KINDS = {"table", "figure"} - -_PROMPT = ( - "You are a precise document layout detector. The attached image is a single " - "rendered PDF page.\n\n" - "Find visually distinct tables and figures that should become reusable " - "document assets. Locate them only — do NOT summarize, transcribe full " - "content, extract keywords, or read data values. Return strict JSON:\n" - "{{\n" - ' "regions": [\n' - " {{\n" - ' "kind": "table|figure",\n' - ' "bbox": [x1, y1, x2, y2],\n' - ' "title": "",\n' - ' "confidence": 0.0\n' - " }}\n" - " ]\n" - "}}\n\n" - "Coordinate system:\n" - "- Treat the page image as a {n}x{n} grid.\n" - "- Origin is the top-left corner.\n" - "- bbox values must be integers in [0, {n}].\n" - "- bbox must tightly include the whole asset: its title, caption, legend, " - "axes, labels, table headers, and footnotes that belong to that asset.\n" - "- Exclude surrounding body paragraphs, page headers, page footers, and " - "page numbers.\n\n" - "Rules:\n" - "- \"table\": data arranged in clear rows and columns — grid lines, cell " - "borders, or strongly aligned cells (data tables, forms, financial tables, " - "appendix tables).\n" - "- \"figure\": any non-table visual asset — bar/line/pie/scatter charts, " - "plots, diagrams, flowcharts, architecture drawings, schematics, or " - "embedded images.\n" - "- Do not mark ordinary paragraphs, bullet lists, title blocks, or loose " - "multi-line text as tables.\n" - "- Do not split a single coherent table or figure into sub-parts.\n" - "- \"title\" is one short label only (single line). Do not duplicate it into " - "other fields and do not write a summary. Use an empty string when there is " - "no visible title or caption.\n" - "- Use confidence 0.0-1.0. Only include assets you can localize.\n" - "- If there are no qualifying assets, return {{\"regions\":[]}}.\n" - "- Return ONLY the JSON object, no markdown fences or explanations." -).format(n=NORM) - - -@dataclass -class Region: - source: str # "vlm" | "reference" - page: int - kind: str # table | figure - bbox_px: list[int] # [x1,y1,x2,y2] in rendered-image pixels - title: str = "" # short asset title / caption (single field, no duplication) - confidence: float = 0.0 - crop_path: str = "" - - -@dataclass -class PageResult: - page: int - width_px: int - height_px: int - width_pt: float - height_pt: float - image_path: str - vlm_regions: list[Region] = field(default_factory=list) - reference_regions: list[Region] = field(default_factory=list) - vlm_error: str = "" - - -# ── rendering ───────────────────────────────────────────────────────── - - -def render_page(page: fitz.Page, dpi: int, out_path: Path) -> tuple[int, int]: - zoom = dpi / 72.0 - mat = fitz.Matrix(zoom, zoom) - pix = page.get_pixmap(matrix=mat, alpha=False) - pix.save(str(out_path)) - return pix.width, pix.height - - -# ── VLM call ────────────────────────────────────────────────────────── - - -def call_vlm(image_path: Path, model: str) -> dict[str, Any]: - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - with open(image_path, "rb") as f: - img_b64 = base64.b64encode(f.read()).decode() - - content_parts = [ - {"type": "text", "text": _PROMPT}, - { - "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{img_b64}"}, - }, - ] - # For this standalone experiment we pass a direct Qwen key when available. - # The production Ali token pool depends on local Redis; direct mode keeps - # the probe runnable on a laptop without changing production behavior. - client = get_openai_client(model=model, api_key=_direct_api_key_for_model(model)) - raw, usage = client.chat_completion_with_usage( - messages=[{"role": "user", "content": content_parts}], - model=model, - temperature=0.0, - max_tokens=1200, - response_format={"type": "json_object"}, - usage_task="experiment.chart_asset_probe", - ) - data = json.loads(raw) - data["_usage"] = usage - return data - - -def _direct_api_key_for_model(model: str) -> str | None: - model_lower = model.lower() - if "qwen" not in model_lower: - return None - single = os.environ.get("ALI_API_KEY", "").strip() - if single: - return single - keys = os.environ.get("ALI_API_KEYS", "").strip() - if not keys: - return None - for item in re.split(r"[,;\s]+", keys): - if item.strip(): - return item.strip() - return None - - -def norm_to_px(box: list[float], w: int, h: int) -> list[int]: - x1, y1, x2, y2 = box - px = [ - int(round(x1 / NORM * w)), - int(round(y1 / NORM * h)), - int(round(x2 / NORM * w)), - int(round(y2 / NORM * h)), - ] - # normalize ordering + clamp - x1, x2 = sorted((px[0], px[2])) - y1, y2 = sorted((px[1], px[3])) - x1 = max(0, min(x1, w)) - x2 = max(0, min(x2, w)) - y1 = max(0, min(y1, h)) - y2 = max(0, min(y2, h)) - return [x1, y1, x2, y2] - - -def load_reference_regions(chunks_path: Path) -> dict[int, list[Region]]: - if not chunks_path.exists(): - raise FileNotFoundError(f"reference chunks not found: {chunks_path}") - payload = json.loads(chunks_path.read_text(encoding="utf-8")) - chunks = payload.get("chunks") if isinstance(payload, dict) else None - if not isinstance(chunks, list): - raise ValueError(f"reference chunks must contain a chunks[] array: {chunks_path}") - - by_page: dict[int, list[Region]] = {} - for chunk in chunks: - if not isinstance(chunk, dict): - continue - metadata = chunk.get("metadata") - if not isinstance(metadata, dict): - continue - bbox = metadata.get("bbox_px") - page_index = metadata.get("page_index") - if not isinstance(bbox, list) or len(bbox) != 4 or page_index is None: - continue - try: - page = int(page_index) - bbox_px = [int(round(float(item))) for item in bbox] - except (TypeError, ValueError): - continue - kind = str(metadata.get("asset_kind") or chunk.get("type") or "asset") - by_page.setdefault(page, []).append( - Region( - source="reference", - page=page, - kind=kind, - bbox_px=bbox_px, - title=str(metadata.get("title") or metadata.get("caption") or ""), - confidence=float(metadata.get("confidence") or 0.0), - ) - ) - return by_page - - -# ── cropping + overlay ──────────────────────────────────────────────── - - -def crop_region(img: Image.Image, region: Region, margin: int, out_path: Path) -> None: - x1, y1, x2, y2 = region.bbox_px - x1 = max(0, x1 - margin) - y1 = max(0, y1 - margin) - x2 = min(img.width, x2 + margin) - y2 = min(img.height, y2 + margin) - if x2 - x1 < 4 or y2 - y1 < 4: - return - img.crop((x1, y1, x2, y2)).save(out_path) - region.crop_path = str(out_path) - - -def draw_overlay(img: Image.Image, page_res: PageResult, out_path: Path) -> None: - canvas = img.convert("RGB").copy() - draw = ImageDraw.Draw(canvas) - try: - font = ImageFont.truetype("/System/Library/Fonts/Supplemental/Arial.ttf", 22) - except Exception: # noqa: BLE001 - font = ImageFont.load_default() - - for r in page_res.reference_regions: - x1, y1, x2, y2 = r.bbox_px - draw.rectangle([x1, y1, x2, y2], outline=(0, 170, 0), width=3) - draw.text((x1 + 2, y1 + 2), f"ref:{r.kind}", fill=(0, 120, 0), font=font) - - for i, r in enumerate(page_res.vlm_regions): - x1, y1, x2, y2 = r.bbox_px - draw.rectangle([x1, y1, x2, y2], outline=(220, 0, 0), width=3) - label = f"vlm:{r.kind} {r.confidence:.2f}" - if r.title: - label += f" | {r.title[:24]}" - draw.text((x1 + 2, max(0, y1 - 24)), label, fill=(200, 0, 0), font=font) - - canvas.save(out_path) - - -# ── page range parsing ──────────────────────────────────────────────── - - -def parse_pages(spec: str, total: int) -> list[int]: - if spec.strip().lower() == "all": - return list(range(1, total + 1)) - pages: set[int] = set() - for part in spec.split(","): - part = part.strip() - if not part: - continue - if "-" in part: - a, b = part.split("-", 1) - pages.update(range(int(a), int(b) + 1)) - else: - pages.add(int(part)) - return sorted(p for p in pages if 1 <= p <= total) - - -# ── main ────────────────────────────────────────────────────────────── - - -def main() -> int: - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("--pdf", required=True) - ap.add_argument("--pages", default="all", help="e.g. 'all', '1-10', '1,3,5'") - ap.add_argument("--dpi", type=int, default=144, help="render DPI (PAGE-TRACK uses 144)") - ap.add_argument( - "--model", - default=os.environ.get("IMAGE_MODEL", ""), - help=( - "VLM model for bbox-only detection. Defaults to $IMAGE_MODEL " - "(qwen3.6-flash). Alternate: qwen3-vl-32b-instruct " - "(open-weights, local-deployable; DashScope ¥2/¥8 per 1M in/out)." - ), - ) - ap.add_argument("--margin", type=int, default=8, help="crop padding px") - ap.add_argument("--no-vlm", action="store_true", help="skip VLM (reference only)") - ap.add_argument( - "--reference-chunks", - default="", - help="existing chunks.json with old asset bbox metadata to overlay in green", - ) - ap.add_argument("--out", default="") - args = ap.parse_args() - - pdf_path = Path(args.pdf) - if not pdf_path.exists(): - print(f"PDF not found: {pdf_path}", file=sys.stderr) - return 2 - if not args.no_vlm and not args.model: - print("No --model and IMAGE_MODEL unset; pass --model or use --no-vlm", - file=sys.stderr) - return 2 - - out_dir = Path(args.out) if args.out else ( - Path.home() / ".knowhere" / "_debug_parse" / pdf_path.stem / "asset_probe" - ) - (out_dir / "pages").mkdir(parents=True, exist_ok=True) - (out_dir / "crops").mkdir(parents=True, exist_ok=True) - (out_dir / "asset_annotate").mkdir(parents=True, exist_ok=True) - - doc = fitz.open(str(pdf_path)) - page_nums = parse_pages(args.pages, doc.page_count) - reference_by_page = ( - load_reference_regions(Path(args.reference_chunks)) - if args.reference_chunks - else {} - ) - print(f"PDF: {pdf_path.name} | {doc.page_count} pages | probing {len(page_nums)} " - f"| dpi={args.dpi} | model={args.model or '(none)'} " - f"| reference={sum(len(items) for items in reference_by_page.values())}") - - results: list[PageResult] = [] - total_tokens = 0 - - for pno in page_nums: - page = doc[pno - 1] - img_path = out_dir / "pages" / f"page-{pno}.png" - w, h = render_page(page, args.dpi, img_path) - pr = PageResult( - page=pno, width_px=w, height_px=h, - width_pt=page.rect.width, height_pt=page.rect.height, - image_path=str(img_path), - ) - pr.reference_regions = list(reference_by_page.get(pno, [])) - print(f"\n[page {pno}] {w}x{h}px") - - if not args.no_vlm: - try: - data = call_vlm(img_path, args.model) - usage = data.pop("_usage", {}) - total_tokens += int(usage.get("total_tokens", 0) or 0) - for k, reg in enumerate(data.get("regions", [])): - box = reg.get("bbox") or reg.get("bbox_norm") - if not box or len(box) != 4: - continue - kind = str(reg.get("kind") or reg.get("type") or "").strip().lower() - if kind not in _VALID_KINDS: - continue - region = Region( - source="vlm", page=pno, - kind=kind, - bbox_px=norm_to_px([float(v) for v in box], w, h), - title=str(reg.get("title") or reg.get("caption") or ""), - confidence=float(reg.get("confidence", 0.0) or 0.0), - ) - pr.vlm_regions.append(region) - print(f" [vlm] {len(pr.vlm_regions)} region(s); " - f"tokens={usage.get('total_tokens', '?')}") - except Exception as exc: # noqa: BLE001 - pr.vlm_error = str(exc) - print(f" [vlm] ERROR: {exc}", file=sys.stderr) - - # crops + overlay - with Image.open(img_path) as im: - for k, r in enumerate(pr.vlm_regions): - crop_region(im, r, args.margin, - out_dir / "crops" / f"page-{pno}_vlm-{k}_{r.kind}.png") - for k, r in enumerate(pr.reference_regions): - crop_region(im, r, args.margin, - out_dir / "crops" / f"page-{pno}_ref-{k}_{r.kind}.png") - if pr.vlm_regions or pr.reference_regions: - draw_overlay(im, pr, out_dir / "asset_annotate" / f"page_{pno}.png") - results.append(pr) - - doc.close() - - # results.json - (out_dir / "results.json").write_text( - json.dumps([asdict(r) for r in results], ensure_ascii=False, indent=2), - encoding="utf-8", - ) - - # report.md - _write_report(out_dir, pdf_path, args, results, total_tokens) - print(f"\nDone. Output: {out_dir}") - print(f" - annotated overlays: {out_dir/'asset_annotate'}") - print(f" - crops: {out_dir/'crops'}") - print(" - results.json / report.md") - return 0 - - -def _write_report(out_dir: Path, pdf_path: Path, args: Any, - results: list[PageResult], total_tokens: int) -> None: - n_vlm = sum(len(r.vlm_regions) for r in results) - n_ref = sum(len(r.reference_regions) for r in results) - lines = [ - f"# Chart/Table Asset Probe — {pdf_path.name}", - "", - f"- pages probed: **{len(results)}**", - f"- dpi: **{args.dpi}**, model: **{args.model or '(no-vlm)'}**", - f"- VLM regions: **{n_vlm}**, reference regions: **{n_ref}**", - f"- total VLM tokens: **{total_tokens}**", - "", - "| page | px | vlm | ref | vlm kinds | vlm titles | vlm error |", - "|-----:|----|----:|----:|-----------|------------|-----------|", - ] - for r in results: - kinds = ",".join(sorted({x.kind for x in r.vlm_regions})) or "-" - titles = "; ".join(x.title for x in r.vlm_regions if x.title) or "-" - err = (r.vlm_error[:40] + "…") if r.vlm_error else "" - lines.append( - f"| {r.page} | {r.width_px}x{r.height_px} | {len(r.vlm_regions)} " - f"| {len(r.reference_regions)} | {kinds} | {titles} | {err} |" - ) - lines += [ - "", - "## How to read", - "- `asset_annotate/page_N.png`: red = new VLM boxes, green = reference boxes from existing chunks.", - "- Judge VLM by: does the red box tightly enclose the table/chart " - + "(incl. caption, excl. body text)? Compare against the green reference.", - "- `crops/`: the actual extracted assets to feed a table model next.", - ] - (out_dir / "report.md").write_text("\n".join(lines), encoding="utf-8") - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/apps/worker/experiments/vlm_bbox_tabula_probe.py b/apps/worker/experiments/vlm_bbox_tabula_probe.py deleted file mode 100644 index 53938bb5..00000000 --- a/apps/worker/experiments/vlm_bbox_tabula_probe.py +++ /dev/null @@ -1,294 +0,0 @@ -"""Experimental: VLM bbox + Tabula PDF table extraction. - -This probe tests the next step after ``chart_asset_probe.py``: - -1. Read only VLM-detected ``kind=table`` regions from ``results.json``. -2. Convert rendered-image pixel boxes back to PDF point coordinates. -3. Pass each area to Tabula (`tabula-py`) against the original PDF text layer. -4. Export candidate DataFrames as HTML/CSV for manual inspection. - -Important: Tabula does not read PNG crops. It needs a text-based PDF, so the -VLM box is used only as a precise `area=[top,left,bottom,right]` constraint. - -Run: - cd apps/worker - uv run --with tabula-py python experiments/vlm_bbox_tabula_probe.py \ - --pdf "/path/to/doc.pdf" \ - --results "/path/to/asset_probe/results.json" -""" - -from __future__ import annotations - -import argparse -import json -import shutil -import subprocess -from dataclasses import asdict, dataclass -from pathlib import Path -from typing import Any - -import pandas as pd - - -@dataclass -class TabulaCandidate: - page: int - region_index: int - mode: str - caption: str - bbox_px: list[int] - area_pt: list[float] - ok: bool - rows: int = 0 - cols: int = 0 - html_path: str = "" - csv_path: str = "" - error: str = "" - - -def _px_box_to_tabula_area( - bbox_px: list[int], - *, - width_px: int, - height_px: int, - width_pt: float, - height_pt: float, - margin_pt: float, -) -> list[float]: - """Convert [x1,y1,x2,y2] px box into Tabula [top,left,bottom,right] points.""" - x1, y1, x2, y2 = bbox_px - left = x1 / width_px * width_pt - right = x2 / width_px * width_pt - top = y1 / height_px * height_pt - bottom = y2 / height_px * height_pt - return [ - round(max(0.0, top - margin_pt), 2), - round(max(0.0, left - margin_pt), 2), - round(min(height_pt, bottom + margin_pt), 2), - round(min(width_pt, right + margin_pt), 2), - ] - - -def _load_vlm_table_regions(results_path: Path) -> list[dict[str, Any]]: - pages = json.loads(results_path.read_text(encoding="utf-8")) - table_regions: list[dict[str, Any]] = [] - for page in pages: - for region_index, region in enumerate(page.get("vlm_regions", [])): - if str(region.get("kind", "")).lower() != "table": - continue - table_regions.append({ - "page": int(page["page"]), - "region_index": region_index, - "caption": str(region.get("caption", "")), - "bbox_px": list(region["bbox_px"]), - "width_px": int(page["width_px"]), - "height_px": int(page["height_px"]), - "width_pt": float(page["width_pt"]), - "height_pt": float(page["height_pt"]), - }) - return table_regions - - -def _is_meaningful_frame(df: pd.DataFrame) -> bool: - if df.empty or df.shape[1] == 0: - return False - non_empty = df.fillna("").astype(str).map(lambda s: bool(s.strip())) - return bool(non_empty.values.sum() >= 2) - - -def _safe_name(page: int, region_index: int, mode: str, candidate_index: int) -> str: - return f"page-{page}_vlm-{region_index}_{mode}-{candidate_index}" - - -def _has_working_java() -> bool: - if shutil.which("java") is None: - return False - try: - result = subprocess.run( - ["java", "-version"], - check=False, - capture_output=True, - text=True, - timeout=10, - ) - except Exception: - return False - combined = f"{result.stdout}\n{result.stderr}" - return result.returncode == 0 and "Unable to locate a Java Runtime" not in combined - - -def _extract_one( - *, - tabula: Any, - pdf_path: Path, - out_dir: Path, - region: dict[str, Any], - mode: str, - area_pt: list[float], -) -> list[TabulaCandidate]: - lattice = mode == "lattice" - stream = mode == "stream" - try: - frames = tabula.read_pdf( - str(pdf_path), - pages=region["page"], - area=area_pt, - guess=False, - lattice=lattice, - stream=stream, - multiple_tables=True, - pandas_options={"header": None}, - ) - except Exception as exc: # noqa: BLE001 - return [ - TabulaCandidate( - page=region["page"], - region_index=region["region_index"], - mode=mode, - caption=region["caption"], - bbox_px=region["bbox_px"], - area_pt=area_pt, - ok=False, - error=str(exc), - ) - ] - - candidates: list[TabulaCandidate] = [] - if not frames: - return [ - TabulaCandidate( - page=region["page"], - region_index=region["region_index"], - mode=mode, - caption=region["caption"], - bbox_px=region["bbox_px"], - area_pt=area_pt, - ok=False, - error="tabula returned no tables", - ) - ] - - for candidate_index, df in enumerate(frames): - name = _safe_name(region["page"], region["region_index"], mode, candidate_index) - candidate = TabulaCandidate( - page=region["page"], - region_index=region["region_index"], - mode=mode, - caption=region["caption"], - bbox_px=region["bbox_px"], - area_pt=area_pt, - ok=_is_meaningful_frame(df), - rows=int(df.shape[0]), - cols=int(df.shape[1]), - ) - if candidate.ok: - html_path = out_dir / "html" / f"{name}.html" - csv_path = out_dir / "csv" / f"{name}.csv" - df.to_html(html_path, index=False, header=False, na_rep="") - df.to_csv(csv_path, index=False, header=False) - candidate.html_path = str(html_path) - candidate.csv_path = str(csv_path) - else: - candidate.error = "empty or near-empty dataframe" - candidates.append(candidate) - return candidates - - -def _write_report(out_dir: Path, candidates: list[TabulaCandidate]) -> None: - ok_count = sum(1 for item in candidates if item.ok) - lines = [ - "# VLM BBox + Tabula Probe", - "", - f"- candidates: **{len(candidates)}**", - f"- successful non-empty tables: **{ok_count}**", - "", - "| page | vlm | mode | ok | shape | caption | error |", - "|-----:|----:|------|----|-------|---------|-------|", - ] - for item in candidates: - error = item.error.replace("\n", " ")[:80] - lines.append( - f"| {item.page} | {item.region_index} | {item.mode} | " - f"{'yes' if item.ok else 'no'} | {item.rows}x{item.cols} | " - f"{item.caption} | {error} |" - ) - lines.append("") - lines.append("Only VLM `kind=table` regions were used; baseline regions were ignored.") - (out_dir / "report.md").write_text("\n".join(lines), encoding="utf-8") - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--pdf", required=True) - parser.add_argument("--results", required=True) - parser.add_argument("--out", default="") - parser.add_argument("--margin-pt", type=float, default=2.0) - args = parser.parse_args() - - pdf_path = Path(args.pdf) - results_path = Path(args.results) - out_dir = ( - Path(args.out) - if args.out - else results_path.parent / "tabula_from_vlm_bbox" - ) - (out_dir / "html").mkdir(parents=True, exist_ok=True) - (out_dir / "csv").mkdir(parents=True, exist_ok=True) - - if not pdf_path.exists(): - raise FileNotFoundError(pdf_path) - if not results_path.exists(): - raise FileNotFoundError(results_path) - if not _has_working_java(): - raise RuntimeError( - "Java runtime is required by tabula-py but was not found. " - "Install a JRE (e.g. OpenJDK/Temurin) and rerun this probe." - ) - - try: - import tabula - except ImportError as exc: - raise RuntimeError( - "tabula-py is required. Run with: uv run --with tabula-py python ..." - ) from exc - - regions = _load_vlm_table_regions(results_path) - print(f"Loaded {len(regions)} VLM table regions from {results_path}") - - candidates: list[TabulaCandidate] = [] - for region in regions: - area_pt = _px_box_to_tabula_area( - region["bbox_px"], - width_px=region["width_px"], - height_px=region["height_px"], - width_pt=region["width_pt"], - height_pt=region["height_pt"], - margin_pt=args.margin_pt, - ) - for mode in ("lattice", "stream"): - extracted = _extract_one( - tabula=tabula, - pdf_path=pdf_path, - out_dir=out_dir, - region=region, - mode=mode, - area_pt=area_pt, - ) - candidates.extend(extracted) - for item in extracted: - print( - f"page={item.page} vlm={item.region_index} " - f"mode={item.mode} ok={item.ok} shape={item.rows}x{item.cols}" - ) - - (out_dir / "results.json").write_text( - json.dumps([asdict(item) for item in candidates], ensure_ascii=False, indent=2), - encoding="utf-8", - ) - _write_report(out_dir, candidates) - print(f"Done. Output: {out_dir}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) From fb88fc4fdfe9e9a9c1feb79f81937acf31666d09 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 21 Jul 2026 22:45:46 +0800 Subject: [PATCH 3/6] feat: enable asset extraction in page memory configuration and update related tests - Changed `asset_extraction_enabled` to `True` in `PageMemoryConfig` to activate asset extraction functionality. - Updated the `page_asset_extraction_enabled` function to return `True`, reflecting the new configuration. - Modified a test to assert that asset extraction is enabled for page memory jobs. - Added `.gitignore` entry for the `experiments` directory in the worker app to exclude local debugging scripts. This commit enhances the asset extraction capabilities in the page memory system, ensuring that the feature is properly enabled and tested. --- .gitignore | 1 + .../contract/test_job_creation_contract.py | 2 +- .../app/services/page_memory/page_assets.py | 2 +- .../unit/test_null_page_parent_locate.py | 211 ---------------- ...test_sort_skeletons_preserves_vlm_order.py | 94 ------- .../worker/tests/unit/test_summary_builder.py | 233 ------------------ .../models/schemas/page_memory_config.py | 2 +- .../shared/services/ai/prompt_service.py | 59 ++--- 8 files changed, 19 insertions(+), 585 deletions(-) delete mode 100644 apps/worker/tests/unit/test_null_page_parent_locate.py delete mode 100644 apps/worker/tests/unit/test_sort_skeletons_preserves_vlm_order.py delete mode 100644 apps/worker/tests/unit/test_summary_builder.py diff --git a/.gitignore b/.gitignore index 428af9f0..859ba378 100644 --- a/.gitignore +++ b/.gitignore @@ -70,6 +70,7 @@ test_*.csv # Local debugging scripts apps/worker/scripts/ +apps/worker/experiments/ apps/worker/start_celery_worker.py apps/worker/start_celery_debug.sh apps/worker/clear_celery_queues.sh diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index a375f515..c5642f0e 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -417,7 +417,7 @@ async def test_should_create_a_v2_page_memory_job_for_pdf_uploads( assert job_metadata["processing_generation"] == "page_memory" assert job_metadata["source_file_name"] == payload["file_name"] assert page_memory_config["max_pages"] == 1500 - assert page_memory_config["asset_extraction_enabled"] is False + assert page_memory_config["asset_extraction_enabled"] is True assert original_request["file_name"] == payload["file_name"] assert "parse_track" not in original_request diff --git a/apps/worker/app/services/page_memory/page_assets.py b/apps/worker/app/services/page_memory/page_assets.py index 7e067a71..05c17fe6 100644 --- a/apps/worker/app/services/page_memory/page_assets.py +++ b/apps/worker/app/services/page_memory/page_assets.py @@ -52,7 +52,7 @@ class PageAsset: def page_asset_extraction_enabled() -> bool: - return False + return True def page_asset_summary_enabled() -> bool: diff --git a/apps/worker/tests/unit/test_null_page_parent_locate.py b/apps/worker/tests/unit/test_null_page_parent_locate.py deleted file mode 100644 index edbc2915..00000000 --- a/apps/worker/tests/unit/test_null_page_parent_locate.py +++ /dev/null @@ -1,211 +0,0 @@ -"""Unit tests for null-page parent locate (compact-strict + self-only emit).""" - -from __future__ import annotations - -import os -from typing import Any -from unittest.mock import MagicMock - -os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") -os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") -os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") -os.environ.setdefault("S3_ACCESS_KEY_ID", "test") -os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") -os.environ.setdefault("S3_TEMP_PATH", "/tmp") - -from app.services.document_agent.structure.hierarchy_locator import ( - TitleMatch, - TitleNode, - locate_title_compact_strict, - resolve_hierarchy_page_ranges, -) -from app.services.page_memory.skeleton_extractor import ( - locate_null_page_parent_overrides, -) - - -def _leaf_match(page: int) -> TitleMatch: - return TitleMatch( - page=page, - confidence=0.9, - source="agent_vlm", - matched_line="", - score=0.9, - candidates=[page], - evidence={}, - ) - - -def _example_tree() -> list[TitleNode]: - return [ - TitleNode( - title="Chapter 1", - level=1, - printed_page=10, - children=[ - TitleNode(title="1.1 Intro", level=2, printed_page=10), - TitleNode(title="1.2 Details", level=2, printed_page=12), - ], - ), - TitleNode( - title="Chapter 2 Overview", - level=1, - printed_page=None, - children=[ - TitleNode(title="2.1 Setup", level=2, printed_page=15), - TitleNode(title="2.2 Results", level=2, printed_page=18), - ], - ), - ] - - -def _example_leaf_overrides() -> dict[tuple[str, ...], TitleMatch]: - return { - ("Chapter 1", "1.1 Intro"): _leaf_match(12), - ("Chapter 1", "1.2 Details"): _leaf_match(14), - ("Chapter 2 Overview", "2.1 Setup"): _leaf_match(17), - ("Chapter 2 Overview", "2.2 Results"): _leaf_match(20), - } - - -def test_compact_strict_matches_cross_line_title() -> None: - match = locate_title_compact_strict( - "Chapter 2 Overview", - scope_pages=[14, 15, 16, 17], - page_texts={ - 14: "1.2 Details body", - 15: "Chapter 2\nOverview\nIntro paragraph", - 16: "more intro", - 17: "2.1 Setup", - }, - ) - assert match is not None - assert match.page == 15 - assert match.evidence.get("accept") == "compact_strict_unique" - - -def test_compact_strict_rejects_ambiguous_pages() -> None: - match = locate_title_compact_strict( - "Chapter 2 Overview", - scope_pages=[14, 15, 16, 17], - page_texts={ - 14: "x", - 15: "Chapter 2\nOverview", - 16: "Chapter 2 Overview again", - 17: "2.1 Setup", - }, - ) - assert match is None - - -def test_resolve_emits_parent_self_only_and_truncates_prev_leaf() -> None: - nodes = _example_tree() - overrides = _example_leaf_overrides() - overrides[("Chapter 2 Overview",)] = TitleMatch( - page=15, - confidence=0.92, - source="anchored", - matched_line="chapter2overview", - score=0.96, - candidates=[15], - evidence={"accept": "compact_strict_unique"}, - ) - page_texts = {page: "" for page in range(12, 21)} - ranges = resolve_hierarchy_page_ranges( - nodes, - page_count=20, - page_texts=page_texts, - body_pages=list(range(12, 21)), - match_overrides=overrides, - ) - by_path = {item.path_titles: item for item in ranges} - - details = by_path[("Chapter 1", "1.2 Details")] - assert details.end_page == 15 - - parent = by_path[("Chapter 2 Overview",)] - assert parent.start_page == 15 - assert parent.end_page == 17 - assert parent.evidence.get("skeleton_kind") == "parent_self_only" - - setup = by_path[("Chapter 2 Overview", "2.1 Setup")] - assert setup.start_page == 17 - - -def test_locate_null_page_parent_overrides_uses_compact_window() -> None: - nodes = _example_tree() - overrides, report = locate_null_page_parent_overrides( - nodes=nodes, - match_overrides=_example_leaf_overrides(), - page_texts={ - 14: "1.2 body", - 15: "Chapter 2\nOverview", - 16: "intro", - 17: "2.1 Setup", - 20: "2.2 Results", - }, - body_pages=list(range(12, 21)), - ctx=None, - ) - parent = overrides[("Chapter 2 Overview",)] - assert parent.page == 15 - assert parent.evidence.get("accept") == "compact_strict_unique" - assert any(row["title"] == "Chapter 2 Overview" and row["page"] == 15 for row in report) - - -def test_locate_null_page_parent_visual_rtl_when_text_ambiguous( - monkeypatch: Any, -) -> None: - nodes = _example_tree() - calls: list[int] = [] - - def _fake_verify( - *, - ctx: Any, - title: str, - candidate_matches: list[TitleMatch], - candidate_page_cap: int, - ) -> dict[str, Any]: - page = candidate_matches[0].page - if title == "Chapter 2 Overview": - calls.append(page) - if title == "Chapter 2 Overview" and page == 15: - return { - "selected_page": 15, - "confidence": 0.8, - "source": "agent_vlm", - "reason": "title heading", - } - return { - "selected_page": None, - "confidence": 0.2, - "source": "agent_vlm", - "reason": "not here", - } - - monkeypatch.setattr( - "app.services.page_memory.skeleton_extractor.verify_section_page_choice", - _fake_verify, - ) - overrides, report = locate_null_page_parent_overrides( - nodes=nodes, - match_overrides=_example_leaf_overrides(), - page_texts={ - 14: "x", - 15: "Chapter 2\nOverview", - 16: "Chapter 2 Overview", - 17: "2.1 Setup", - 20: "2.2", - }, - body_pages=list(range(12, 21)), - ctx=MagicMock(), - ) - parent = overrides[("Chapter 2 Overview",)] - assert parent.page == 15 - assert parent.evidence.get("accept") == "visual_rtl" - assert calls[0] == 17 - assert 15 in calls - assert calls.index(15) > calls.index(17) - entry = next(row for row in report if row["title"] == "Chapter 2 Overview") - assert entry["visual_verify_calls"] >= 1 - assert entry["result"] == "visual_rtl" diff --git a/apps/worker/tests/unit/test_sort_skeletons_preserves_vlm_order.py b/apps/worker/tests/unit/test_sort_skeletons_preserves_vlm_order.py deleted file mode 100644 index 8fea71b4..00000000 --- a/apps/worker/tests/unit/test_sort_skeletons_preserves_vlm_order.py +++ /dev/null @@ -1,94 +0,0 @@ -"""sort_skeletons must not alphabetize same-page VLM/TOC order.""" - -from __future__ import annotations - -from app.services.page_memory._serialization import ( - build_hierarchy_tree, - serialize_hierarchy_artifact, -) -from app.services.page_memory._utils import sort_skeletons -from app.services.page_memory.skeleton_extractor import SectionSkeleton - - -def _skel(path: str, title: str, *, start_page: int, level: int = 3) -> SectionSkeleton: - parent = "/".join(path.split("/")[:-1]) or None - return SectionSkeleton( - section_path=path, - title=title, - level=level, - start_page=start_page, - end_page=start_page, - parent_path=parent, - evidence={}, - ) - - -def test_same_page_siblings_keep_input_order_not_alphabetical() -> None: - base = "doc.pdf/Section D/Part D3 Construction of exits" - # VLM / reading order: Separation (D3D5) before Open access (D3D6) - page_197 = [ - _skel( - f"{base}/Non-fire-isolated stairways and ramps", - "Non-fire-isolated stairways and ramps", - start_page=197, - ), - _skel( - f"{base}/Separation of rising and descending stair flights", - "Separation of rising and descending stair flights", - start_page=197, - ), - _skel( - f"{base}/Open access ramps and balconies", - "Open access ramps and balconies", - start_page=197, - ), - _skel(f"{base}/Smoke lobbies", "Smoke lobbies", start_page=197), - ] - # Later page first in input — page sort must move it after, without - # alphabetizing the same-page block. - mixed = [ - _skel(f"{base}/Landings", "Landings", start_page=201), - *page_197, - ] - titles = [s.title for s in sort_skeletons(mixed)] - assert titles == [ - "Non-fire-isolated stairways and ramps", - "Separation of rising and descending stair flights", - "Open access ramps and balconies", - "Smoke lobbies", - "Landings", - ] - # Alphabetical would put Open before Separation — must not happen - assert titles.index("Separation of rising and descending stair flights") < titles.index( - "Open access ramps and balconies" - ) - - -def test_hierarchy_tree_matches_node_list_order() -> None: - base = "doc.pdf/Part D3" - skeletons = [ - _skel( - f"{base}/Separation of rising and descending stair flights", - "Separation of rising and descending stair flights", - start_page=197, - ), - _skel( - f"{base}/Open access ramps and balconies", - "Open access ramps and balconies", - start_page=197, - ), - ] - tree = build_hierarchy_tree(skeletons) - assert list(tree["Part D3"].keys()) == [ - "Separation of rising and descending stair flights", - "Open access ramps and balconies", - ] - - artifact = serialize_hierarchy_artifact(skeletons) - assert list(artifact["HIERARCHY"]["Part D3"].keys()) == [ - n["title"] for n in artifact["nodes"] - ] - assert [n["title"] for n in artifact["nodes"]] == [ - "Separation of rising and descending stair flights", - "Open access ramps and balconies", - ] diff --git a/apps/worker/tests/unit/test_summary_builder.py b/apps/worker/tests/unit/test_summary_builder.py deleted file mode 100644 index c0685ef2..00000000 --- a/apps/worker/tests/unit/test_summary_builder.py +++ /dev/null @@ -1,233 +0,0 @@ -"""Unit tests for bottom-up doc_nav summary enrichment.""" - -from __future__ import annotations - -from typing import Any, Dict, List - -import pytest - -from app.services.connect_builder.summary_builder import ( - SUMMARY_MAX_LEN, - _deterministic_section_summary, - _llm_summarize, - _recursive_summarize_nav, - build_self_only_lookup, -) - - -def _leaf(title: str, summary: str = "", path: str = "") -> Dict[str, Any]: - node: Dict[str, Any] = {"title": title, "summary": summary, "children": []} - if path: - node["path"] = path - return node - - -def _parent( - title: str, - children: List[Dict[str, Any]], - *, - path: str = "", - summary: str = "", -) -> Dict[str, Any]: - node: Dict[str, Any] = { - "title": title, - "summary": summary, - "children": children, - } - if path: - node["path"] = path - return node - - -class TestDeterministicAssembly: - def test_order_covers_self_only_then_titles(self) -> None: - text = _deterministic_section_summary( - is_top_level=False, - self_only="intro paragraph here", - child_titles=["Alpha", "Beta"], - ) - assert text.startswith("This section covers: ") - assert "intro paragraph here" in text - assert "Alpha, Beta" in text - # self_only before titles - assert text.index("intro paragraph here") < text.index("Alpha, Beta") - - def test_all_child_titles_even_when_summary_empty(self) -> None: - parent = _parent( - "Parent", - [ - _leaf("HasText", summary="body"), - _leaf("NoSummary", summary=""), - ], - path="doc.pdf/Parent", - ) - result = _recursive_summarize_nav( - parent, - use_llm=False, - source_file_name="doc.pdf", - ) - assert "HasText" in result - assert "NoSummary" in result - assert result.startswith("This section covers: ") - - -class TestSelfOnlyLookup: - def test_exact_path_only_excludes_descendants(self) -> None: - chunks = [ - { - "path": "doc.pdf/2.4.4 隐患治理", - "content": "PARENT_INTRO_ONLY", - }, - { - "path": "doc.pdf/2.4.4 隐患治理/清单项A", - "content": "CHILD_BODY_SHOULD_NOT_APPEAR", - }, - ] - lookup = build_self_only_lookup(chunks, source_file_name="doc.pdf") - assert lookup["2.4.4 隐患治理"] == "PARENT_INTRO_ONLY" - assert "CHILD_BODY_SHOULD_NOT_APPEAR" not in lookup["2.4.4 隐患治理"] - assert "2.4.4 隐患治理 / 清单项A" in lookup - - def test_nonleaf_includes_self_only_in_deterministic(self) -> None: - parent = _parent( - "2.4.4 隐患治理", - [ - _leaf("清单项A", summary="a", path="doc.pdf/2.4.4 隐患治理/清单项A"), - _leaf("清单项B", summary="b", path="doc.pdf/2.4.4 隐患治理/清单项B"), - ], - path="doc.pdf/2.4.4 隐患治理", - ) - lookup = build_self_only_lookup( - [{"path": "doc.pdf/2.4.4 隐患治理", "content": "方案包括以下内容："}], - source_file_name="doc.pdf", - ) - result = _recursive_summarize_nav( - parent, - use_llm=False, - self_only_lookup=lookup, - source_file_name="doc.pdf", - ) - assert "方案包括以下内容：" in result - assert "清单项A" in result - assert "清单项B" in result - assert parent.get("self_summary") == "方案包括以下内容：" - - -class TestLlmTrigger: - def test_short_contrib_skips_llm(self, monkeypatch: pytest.MonkeyPatch) -> None: - called = {"n": 0} - - def _boom(**kwargs: Any) -> str: - called["n"] += 1 - return "SHOULD_NOT_USE" - - monkeypatch.setattr( - "app.services.connect_builder.summary_builder._llm_summarize", - _boom, - ) - parent = _parent( - "P", - [_leaf("A", summary="x"), _leaf("B", summary="y")], - path="doc.pdf/P", - ) - result = _recursive_summarize_nav(parent, use_llm=True, source_file_name="doc.pdf") - assert called["n"] == 0 - assert result.startswith("This section covers: ") - assert "A" in result and "B" in result - - def test_long_contrib_calls_llm_with_title_for_empty_summary( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - captured: Dict[str, Any] = {} - - def _fake_llm(**kwargs: Any) -> str: - captured.update(kwargs) - return "LLM_SUMMARY" - - monkeypatch.setattr( - "app.services.connect_builder.summary_builder._llm_summarize", - _fake_llm, - ) - long_a = "A" * (SUMMARY_MAX_LEN + 5) - parent = _parent( - "P", - [ - _leaf("HasSummary", summary=long_a), - _leaf("EmptySummary", summary=""), - ], - path="doc.pdf/P", - ) - lookup = {"P": "SELF_ONLY_INTRO"} - # section path from doc.pdf/P is "P" - result = _recursive_summarize_nav( - parent, - use_llm=True, - self_only_lookup=lookup, - source_file_name="doc.pdf", - ) - assert result == "LLM_SUMMARY" - assert captured["self_only"] == "SELF_ONLY_INTRO" - titles = [t for t, _ in captured["child_rows"]] - contribs = {t: c for t, c in captured["child_rows"]} - assert "EmptySummary" in titles - assert contribs["EmptySummary"] == "EmptySummary" - assert contribs["HasSummary"] == long_a - - def test_single_child_with_self_only_does_not_copy_child( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - monkeypatch.setattr( - "app.services.connect_builder.summary_builder._llm_summarize", - lambda **kwargs: "MERGED", - ) - long_child = "C" * (SUMMARY_MAX_LEN + 1) - parent = _parent( - "P", - [_leaf("OnlyChild", summary=long_child)], - path="doc.pdf/P", - ) - result = _recursive_summarize_nav( - parent, - use_llm=True, - self_only_lookup={"P": "intro"}, - source_file_name="doc.pdf", - ) - assert result == "MERGED" - assert result != long_child - - -class TestPromptPayload: - def test_file_summary_prompt_contains_scope_blocks( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - captured: Dict[str, Any] = {} - - def _fake_client() -> Any: - class _C: - def chat_completion(self, **kwargs: Any) -> str: - captured["messages"] = kwargs.get("messages") - return "ok" - - return _C() - - monkeypatch.setattr( - "shared.services.ai.openai_compatible_client_sync.get_openai_client", - _fake_client, - ) - # Ensure build_prompt path works - out = _llm_summarize( - node_name="Parent", - self_only="intro text", - child_rows=[("ChildA", "summary A"), ("ChildB", "ChildB")], - max_tokens=100, - ) - assert out == "ok" - user = captured["messages"][1]["content"] - assert "SCOPE_TITLE: Parent" in user - assert "SELF_ONLY_CONTENT:" in user - assert "intro text" in user - assert "COVERED_NODES:" in user - assert "[ChildA] summary A" in user - assert "[ChildB] ChildB" in user - # legacy flat blob prompt removed - assert "You will receive summaries of sub-sections" not in user diff --git a/packages/shared-python/shared/models/schemas/page_memory_config.py b/packages/shared-python/shared/models/schemas/page_memory_config.py index a1bc6090..f8134ded 100644 --- a/packages/shared-python/shared/models/schemas/page_memory_config.py +++ b/packages/shared-python/shared/models/schemas/page_memory_config.py @@ -20,7 +20,7 @@ class PageMemoryConfig: hierarchy_model: str | None = None hierarchy_max_tokens: int = 2000 max_heading_depth: int = 6 - asset_extraction_enabled: bool = False + asset_extraction_enabled: bool = True asset_summary_enabled: bool = False asset_model: str = "qwen3.6-flash" asset_max_pages: int | None = None diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index 35323256..27814c2c 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -747,47 +747,7 @@ def build_prompt(task, texts, query, **kwargs): top_p = 0.01 max_tokens = kwargs.get("paras", {}).get("max_tokens", 1200) grid_size = kwargs.get("paras", {}).get("grid_size", 1000) - # Previous production prompt kept for comparison: - # prompt = f"""\ - # You are a precise document layout detector. The attached image is a single PDF - # page screenshot. - # - # Find visually distinct tables, charts, and figures that should become reusable - # document assets. Return strict JSON: - # {{ - # "regions": [ - # {{ - # "kind": "table|chart|figure", - # "bbox": [x1, y1, x2, y2], - # "caption": "", - # "title": "", - # "summary": "<1 sentence searchable summary>", - # "keywords": ["", ""], - # "confidence": 0.0 - # }} - # ] - # }} - # - # Coordinate system: - # - Treat the page image as a {grid_size}x{grid_size} grid. - # - Origin is the top-left corner. - # - bbox values must be integers in [0, {grid_size}]. - # - bbox must tightly include the whole asset: title, caption, legend, axes, - # labels, table headers, and footnotes that are part of the asset. - # - Exclude surrounding body paragraphs, page headers, page footers, and page - # numbers. - # - # Rules: - # - "table": rows/columns of data, forms, financial tables, appendix tables. - # - "chart": plotted data such as bar/line/pie/scatter charts. - # - "figure": distinct diagrams, flowcharts, architecture drawings, embedded images. - # - # - Do not transcribe entire tables. Summarize the topic and extreme values based on main columns or rows. - # - "keywords" must be an array of up to 5 strings in the same language as the visible asset text. - # - Use confidence 0.0-1.0. Only include assets you can localize. - # - If there are no assets, return {{"regions":[]}}. - # - Return ONLY the JSON object, no markdown fences or explanations. - # """ + prompt = f"""\ You are a precise document layout detector. The attached image is a single rendered PDF page. @@ -822,9 +782,20 @@ def build_prompt(task, texts, query, **kwargs): - "figure": any non-table visual asset - bar/line/pie/scatter charts, plots, diagrams, flowcharts, architecture drawings, schematics, or embedded images. - - Do not mark ordinary paragraphs, bullet lists, title blocks, or loose - multi-line text as tables. - - Do not split a single coherent table or figure into sub-parts. + - Prefer one bbox for the whole figure. When multiple visual parts clearly + form one composition (shared caption or a multi-panel explanation of the same concept/process), + return them as a single figure, not separate images. + - Treat a flowchart or process diagram as one figure, including its nodes, + edges, labels, and legend when they belong together. + + - Do NOT extract page backgrounds, watermarks, stamps, or decorative underlays. + - Do NOT extract small logos, icons, bullets, or other scattered decorative + marks that are not standalone informative figures. + - Do NOT extract ornamental digits/letters placed before a heading title as figures. + + - Do NOT mark ordinary paragraphs, bullet lists, title blocks, or loose multi-line text as tables. + - Do NOT split a single coherent table or figure into sub-parts. + - "title" is one short label only (single line). Do not duplicate it into other fields and do not write a summary. Use an empty string when there is no visible title or caption. From c54679ab3a3a0c5dbe65d3d135452db43cf84602 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 24 Jul 2026 12:12:17 +0800 Subject: [PATCH 4/6] feat: integrate asset probing and enhance document profiling - Added `_ensure_asset_probe` method to `ProfileCoordinator` to ensure asset probing occurs during profile generation. - Introduced `has_asset` and `asset_bboxes` fields in `PageFeature` for tracking asset presence and bounding boxes. - Updated `aggregate_doc_stats` to include asset-related statistics when assets have been probed. - Refactored `probe_page_features` to support asset probing alongside existing features. - Modified `planner` and `tagger` logic to accommodate new asset-related data and improve document profiling accuracy. This commit enhances the document agent's ability to extract and utilize asset information, improving overall profiling and analysis capabilities. --- .../document_agent/bootstrap/__init__.py | 12 +- .../bootstrap/aggregate_stats.py | 56 +-- .../document_agent/bootstrap/probe.py | 9 +- .../services/document_agent/coordinator.py | 26 ++ .../app/services/document_agent/manifest.py | 9 +- .../document_agent/planner/planner.py | 62 ++- .../document_agent/planner/prompts.py | 34 +- .../tools/classify_page_kinds.py | 31 -- .../tools/probe_page_features.py | 403 ++++++++++++++++-- .../tools/propose_shard_plan.py | 26 +- .../services/page_memory/memory_service.py | 1 - .../app/services/page_memory/page_plan.py | 32 +- .../app/services/page_memory/page_tagger.py | 95 +---- .../test_doc_profile_anatomy_contract.py | 2 + .../models/schemas/page_memory_config.py | 8 +- .../shared/services/ai/prompt_service.py | 58 +-- 16 files changed, 543 insertions(+), 321 deletions(-) diff --git a/apps/worker/app/services/document_agent/bootstrap/__init__.py b/apps/worker/app/services/document_agent/bootstrap/__init__.py index b8494b99..2a490a4b 100644 --- a/apps/worker/app/services/document_agent/bootstrap/__init__.py +++ b/apps/worker/app/services/document_agent/bootstrap/__init__.py @@ -2,6 +2,14 @@ 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 +from app.services.document_agent.bootstrap.probe import ( + probe_page_assets, + probe_page_features, +) -__all__ = ["aggregate_doc_stats", "classify_page_kinds", "probe_page_features"] +__all__ = [ + "aggregate_doc_stats", + "classify_page_kinds", + "probe_page_assets", + "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 index f0739edf..96a11942 100644 --- a/apps/worker/app/services/document_agent/bootstrap/aggregate_stats.py +++ b/apps/worker/app/services/document_agent/bootstrap/aggregate_stats.py @@ -8,28 +8,22 @@ from app.services.document_agent.manifest import PageFeature, ToolContext, ToolResult +# Coarse sampling extrema are text-only. Low text / density already proxy +# chart-heavy or asset-heavy pages; table/drawing counts are not used to +# nominate extrema pages for VLM coarse classification. 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", } @@ -97,38 +91,30 @@ def aggregate_doc_stats(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: deduped_extrema = sorted(set(extrema_pages)) landscape_pages = sum(1 for feature in features if feature.orientation == "landscape") - scan_like_pages = sum( - 1 - for feature in features - if feature.raw_text_length < 50 and feature.image_coverage >= 0.5 - ) - image_heavy_pages = sum( - 1 for feature in features if feature.image_coverage >= 0.35 - ) - table_signal_pages = sum( - 1 - for feature in features - if feature.table_count > 0 or feature.drawings_count >= 25 - ) - doc_shape = { + doc_shape: dict[str, Any] = { "page_count": page_count, "landscape_pages": landscape_pages, "landscape_ratio": round(landscape_pages / page_count, 4) if page_count else 0.0, - "scan_like_pages": scan_like_pages, - "scan_like_ratio": round(scan_like_pages / page_count, 4) - if page_count - else 0.0, - "image_heavy_pages": image_heavy_pages, - "image_heavy_ratio": round(image_heavy_pages / page_count, 4) - if page_count - else 0.0, - "table_signal_pages": table_signal_pages, - "table_signal_ratio": round(table_signal_pages / page_count, 4) - if page_count - else 0.0, } + if ctx.blackboard.global_signals.get("assets_probed"): + scan_like_pages = sum( + 1 + for feature in features + if feature.raw_text_length < 50 and feature.image_coverage >= 0.5 + ) + asset_pages = sum(1 for feature in features if feature.has_asset) + doc_shape.update( + { + "scan_like_pages": scan_like_pages, + "scan_like_ratio": round(scan_like_pages / page_count, 4) + if page_count + else 0.0, + "asset_pages": asset_pages, + "asset_ratio": round(asset_pages / page_count, 4) if page_count else 0.0, + } + ) ctx.blackboard.doc_stats = stats ctx.blackboard.extrema_pages = deduped_extrema ctx.blackboard.global_signals["doc_stats"] = stats diff --git a/apps/worker/app/services/document_agent/bootstrap/probe.py b/apps/worker/app/services/document_agent/bootstrap/probe.py index 14877b1b..591dfbd9 100644 --- a/apps/worker/app/services/document_agent/bootstrap/probe.py +++ b/apps/worker/app/services/document_agent/bootstrap/probe.py @@ -1,5 +1,8 @@ -"""Bootstrap wrapper for deterministic page probing.""" +"""Bootstrap wrappers for deterministic page probing.""" -from app.services.document_agent.tools.probe_page_features import probe_page_features +from app.services.document_agent.tools.probe_page_features import ( + probe_page_assets, + probe_page_features, +) -__all__ = ["probe_page_features"] +__all__ = ["probe_page_assets", "probe_page_features"] diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 4630ff43..2f25f6a1 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -10,6 +10,7 @@ from app.services.document_agent.bootstrap import ( aggregate_doc_stats, classify_page_kinds, + probe_page_assets, probe_page_features, ) from app.services.document_agent.budget import BudgetTracker, StageEnvelope @@ -151,6 +152,7 @@ def _run_coarse(self) -> DocumentProfile: profile, _initial_decision, _planner_result = self._propose_profile( actor="planner:coarse" ) + self._ensure_asset_probe() return profile def _run_structural(self) -> PageAnatomyMap: @@ -164,6 +166,7 @@ def _run_structural(self) -> PageAnatomyMap: profile, initial_decision, _planner_result = self._propose_profile( actor="planner" ) + self._ensure_asset_probe() executor_result = ReActExecutor( self.ctx, registry=REGISTRY, @@ -200,6 +203,7 @@ def _run_lightweight_anatomy( self.state = DocumentAgentState.RUNNING if not self.blackboard.page_features: self._run_bootstrap() + self._ensure_asset_probe() if self.blackboard.toc_result is None: if self._toc_profile_enabled(): self.blackboard.toc_result = TocResult( @@ -284,6 +288,28 @@ def _run_bootstrap(self) -> None: raise RuntimeError(result.error or f"{tool_name} failed") self.round_index += 1 + def _ensure_asset_probe(self) -> None: + if self.blackboard.global_signals.get("assets_probed"): + return + if not self.blackboard.page_features: + raise RuntimeError("page_features missing; run text bootstrap first") + for tool_name, handler in ( + ("probe.page_assets", probe_page_assets), + ("aggregate.doc_stats", aggregate_doc_stats), + ): + 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 _toc_result_requires_strict_retry(self) -> bool: toc_result = self.blackboard.toc_result return bool( diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index 8861f631..01762066 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -7,7 +7,7 @@ from typing import Any, Literal -PageKind = Literal["normal", "table_heavy", "image_heavy", "low_content", "landscape"] +PageKind = Literal["normal", "landscape"] TocFailureKind = Literal["none", "confirm_failed", "rejected_all", "degraded"] ReflexionAction = Literal["tool_call", "verdict_now"] @@ -26,7 +26,10 @@ class PageFeature: orientation: Literal["portrait", "landscape"] width: float 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 text_lines_preview: list[str] = field(default_factory=list) def to_dict(self) -> dict[str, Any]: @@ -52,6 +55,10 @@ class DocumentProfile: category_rationale: str = "" language: str = "unknown" rationale: str = "" + # Content-band margins as fractions of page height (top origin, y down). + # header_y: lowest header line among sample pages; footer_y: highest footer. + header_y: float | None = None + footer_y: float | None = None def to_dict(self) -> dict[str, Any]: return asdict(self) diff --git a/apps/worker/app/services/document_agent/planner/planner.py b/apps/worker/app/services/document_agent/planner/planner.py index 533cc473..7e5e1ce0 100644 --- a/apps/worker/app/services/document_agent/planner/planner.py +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -26,18 +26,6 @@ "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.", } @@ -56,11 +44,9 @@ def _feature_rows(ctx: ToolContext, pages: list[int]) -> list[dict[str, Any]]: "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, + "width": feature.width, + "height": feature.height, "is_blank_like": feature.is_blank_like, } ) @@ -78,6 +64,11 @@ def _segment_sample(candidates: list[int], count: int) -> list[int]: return [candidates[round(index * step)] for index in range(count)] +# Coarse VLM budget: extrema first, then front/mid/back fill, hard cap 10. +_COARSE_SAMPLE_CAP = 10 +_COARSE_SEGMENT_QUOTAS = (2, 2, 2) # front, middle, back + + def _sample_pages( page_count: int, extrema_pages: list[int], @@ -85,9 +76,16 @@ def _sample_pages( ) -> list[int]: """Select representative pages for VLM profiling. + Strategy (cap 10): + 1. Text extrema first (length/density min+max, ≤4 unique). + 2. Fill remaining slots with 2/2/2 stratified samples from + front/middle/back of the non-extrema pool. + 3. Hard truncate to ``_COARSE_SAMPLE_CAP``. + Args: page_count: Total number of pages. - extrema_pages: Pages with statistical extrema (min/max text, tables, etc.). + extrema_pages: Pages with text extrema (min/max raw_text_length / + text_density). Low-text extrema already surface chart/asset pages. exclude_pages: Pages to skip entirely (e.g. TOC pages already detected by the TOC pipeline). These inflate text-density metrics without adding profiling value. @@ -98,21 +96,34 @@ def _sample_pages( extrema = [page for page in extrema_pages if 1 <= page <= page_count and page not in skip] pool = [page for page in range(1, page_count + 1) if page not in set(extrema) and page not in skip] if not pool: - return sorted(set(extrema)) + return sorted(set(extrema))[:_COARSE_SAMPLE_CAP] third = max(len(pool) // 3, 1) front = pool[:third] middle = pool[third : third * 2] back = pool[third * 2 :] + front_n, middle_n, back_n = _COARSE_SEGMENT_QUOTAS sampled = ( - _segment_sample(front, 4) - + _segment_sample(middle or pool, 3) - + _segment_sample(back or pool, 3) + _segment_sample(front, front_n) + + _segment_sample(middle or pool, middle_n) + + _segment_sample(back or pool, back_n) ) ordered = [] for page in extrema + sampled: if page not in ordered: ordered.append(page) - return ordered[:20] + return ordered[:_COARSE_SAMPLE_CAP] + + +def _parse_margin_ratio(value: Any) -> float | None: + if value is None or value == "": + return None + try: + ratio = float(value) + except (TypeError, ValueError): + return None + if 0.0 <= ratio <= 1.0: + return ratio + return None def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDecision]: @@ -130,6 +141,11 @@ def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDec is_scanned = raw_is_scanned.strip().lower() in {"true", "yes", "1", "scanned"} else: is_scanned = bool(raw_is_scanned) + header_y = _parse_margin_ratio(data.get("header_y")) + footer_y = _parse_margin_ratio(data.get("footer_y")) + if header_y is not None and footer_y is not None and header_y >= footer_y: + header_y = None + footer_y = None profile = DocumentProfile( is_scanned=is_scanned, category=category or "unknown document", @@ -137,6 +153,8 @@ def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDec category_rationale=str(data.get("category_rationale") or ""), language=str(data.get("language") or "unknown"), rationale=str(data.get("rationale") or ""), + header_y=header_y, + footer_y=footer_y, ) next_action = str(data.get("next_action") or "ready_to_shard") tool_name: str | None = None diff --git a/apps/worker/app/services/document_agent/planner/prompts.py b/apps/worker/app/services/document_agent/planner/prompts.py index 7bb1f316..4ca93ef8 100644 --- a/apps/worker/app/services/document_agent/planner/prompts.py +++ b/apps/worker/app/services/document_agent/planner/prompts.py @@ -1,20 +1,26 @@ """Prompts for the document profile planner.""" PLANNER_INSTRUCTIONS = ( - "You are a document profile agent. Use global page-feature statistics, " - "optional TOC/H1 evidence, and page screenshots to classify the PDF. Return strict " - "JSON only with keys: is_scanned, category, routing_category, " - "category_rationale, language, rationale, next_action, inspect_pages, grep_query. " - "category is a concise semantic document type, at most 5 English words, such " - "as Financial Prospectus, Technical Manual, Corporate Policy, Research Report, " - "Engineering Atlas, or Scanned Handbook. routing_category must be one of " - "atlas, scanned, slides, generic. Set routing_category=atlas only for " - "engineering drawing collections, construction standard atlases, or page sets " - "whose primary unit is a drawing/detail sheet rather than prose. 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." + "You are a document profile agent. Use page-feature statistics, optional " + "TOC/H1 evidence, and the provided page screenshots to classify the PDF. " + "Return strict JSON only with keys: is_scanned, category, routing_category, " + "category_rationale, language, rationale, header_y, footer_y, next_action, " + "inspect_pages, grep_query. " + "category is a concise semantic document type in at most 5 English words. " + "routing_category must be one of atlas, scanned, slides, generic. " + "Set routing_category=atlas only when pages are primarily drawing/detail " + "sheets rather than prose. " + "header_y and footer_y are document-level horizontal content-margin lines " + "as fractions of page height in [0, 1], origin at the top with y increasing " + "downward. From the sample pages shown: header_y is the lowest header line " + "you observe (largest y) when any header is present, otherwise null; " + "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. " + "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/tools/classify_page_kinds.py b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py index fb60765d..923a816c 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 @@ -11,17 +11,6 @@ def _label_feature(feature: PageFeature) -> PageLabel: page = feature.page - if ( - feature.raw_text_length < 80 - and feature.image_coverage < 0.02 - and feature.drawings_count < 5 - ): - return PageLabel( - page=page, - kind="low_content", - confidence=0.78, - evidence={"signal": "low_text_image_drawings"}, - ) if feature.orientation == "landscape": return PageLabel( page=page, @@ -29,23 +18,6 @@ def _label_feature(feature: PageFeature) -> PageLabel: confidence=0.78, evidence={"width": feature.width, "height": feature.height}, ) - if feature.image_coverage >= 0.35 and feature.raw_text_length < 250: - return PageLabel( - page=page, - kind="image_heavy", - confidence=0.84, - evidence={"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, - }, - ) return PageLabel(page=page, kind="normal", confidence=0.65, evidence={}) @@ -67,9 +39,6 @@ def classify_page_kinds(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: "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 []), } ) 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 678557fa..f4e3ff8a 100644 --- a/apps/worker/app/services/document_agent/tools/probe_page_features.py +++ b/apps/worker/app/services/document_agent/tools/probe_page_features.py @@ -1,9 +1,11 @@ -"""Full-page structural probing.""" +"""Full-page structural probing: text pass, then optional asset pass.""" from __future__ import annotations import gc +import math import time +from collections import defaultdict from typing import Any from app.services.document_agent.manifest import PageFeature, ToolContext, ToolResult @@ -14,6 +16,16 @@ ) from loguru import logger +# Cluster nearby drawing strokes into one figure region (PDF points). +_FIGURE_CLUSTER_GAP = 18.0 +# Drop clusters smaller than this many strokes (noise). +_FIGURE_CLUSTER_MIN_PATHS = 3 +# Drop aggregated figures smaller than this area (PDF pt²); catches most logos. +_FIGURE_MIN_AREA = 5000.0 +# Near-full-page sparse stroke frames are treated as borders, not figures. +_FIGURE_FULLPAGE_AREA_RATIO = 0.92 +_FIGURE_FULLPAGE_MAX_PATHS = 25 + def _rect_area(rect: Any) -> float: width = max(float(getattr(rect, "width", 0.0) or 0.0), 0.0) @@ -21,12 +33,185 @@ def _rect_area(rect: Any) -> float: 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 +def _clip_bbox( + x0: float, + y0: float, + x1: float, + y1: float, + *, + clip: Any, +) -> list[float]: + return [ + round(max(x0, float(clip.x0)), 2), + round(max(y0, float(clip.y0)), 2), + round(min(x1, float(clip.x1)), 2), + round(min(y1, float(clip.y1)), 2), + ] + + +def _valid_bbox(box: list[float]) -> bool: + return len(box) == 4 and box[2] > box[0] and box[3] > box[1] + + +def _outside_content_band( + y0: float, + y1: float, + *, + page_h: float, + header_y: float | None, + footer_y: float | None, +) -> bool: + """True if any vertical edge of an aggregated figure leaves the content band. + + ``header_y`` / ``footer_y`` are fractions of page height (top origin, y down). + """ + if header_y is not None and y0 < header_y * page_h: + return True + if footer_y is not None and y1 > footer_y * page_h: + return True + return False + + +def _rects_near( + a: tuple[float, float, float, float], + b: tuple[float, float, float, float], + gap: float, +) -> bool: + ax0, ay0, ax1, ay1 = a + bx0, by0, bx1, by1 = b + return not ( + ax1 + gap < bx0 + or bx1 + gap < ax0 + or ay1 + gap < by0 + or by1 + gap < ay0 + ) + + +def _cluster_indices(rects: list[tuple[float, float, float, float]], gap: float) -> list[list[int]]: + """Union-find cluster of nearby AABBs via spatial grid (avg ~O(n)).""" + n = len(rects) + if n == 0: + return [] + parent = list(range(n)) + + def find(i: int) -> int: + while parent[i] != i: + parent[i] = parent[parent[i]] + i = parent[i] + return i + + def union(i: int, j: int) -> None: + ri, rj = find(i), find(j) + if ri != rj: + parent[rj] = ri + + cell = max(gap, 1.0) + grid: dict[tuple[int, int], list[int]] = defaultdict(list) + + def cell_range( + x0: float, y0: float, x1: float, y1: float + ) -> tuple[int, int, int, int]: + return ( + math.floor(x0 / cell), + math.floor(y0 / cell), + math.floor(x1 / cell), + math.floor(y1 / cell), + ) + + for i, (x0, y0, x1, y1) in enumerate(rects): + ix0, iy0, ix1, iy1 = cell_range(x0, y0, x1, y1) + for ix in range(ix0, ix1 + 1): + for iy in range(iy0, iy1 + 1): + grid[(ix, iy)].append(i) + + for i, (x0, y0, x1, y1) in enumerate(rects): + ix0, iy0, ix1, iy1 = cell_range(x0 - gap, y0 - gap, x1 + gap, y1 + gap) + seen: set[int] = set() + for ix in range(ix0, ix1 + 1): + for iy in range(iy0, iy1 + 1): + for j in grid.get((ix, iy), ()): + if j <= i or j in seen: + continue + seen.add(j) + if _rects_near(rects[i], rects[j], gap): + union(i, j) + + groups: dict[int, list[int]] = defaultdict(list) + for i in range(n): + groups[find(i)].append(i) + return list(groups.values()) + +def _figure_bboxes_from_drawings( + drawings: list[Any], + *, + page_rect: Any, + page_area: float, + header_y: float | None, + footer_y: float | None, +) -> list[dict[str, Any]]: + """Cluster drawing strokes, then filter aggregated figure bboxes only.""" + page_h = float(getattr(page_rect, "height", 0.0) or 0.0) + kept_rects: list[tuple[float, float, float, float]] = [] + for drawing in drawings: + if not isinstance(drawing, dict): + continue + rect = drawing.get("rect") + if rect is None: + continue + x0 = float(getattr(rect, "x0", 0.0) or 0.0) + y0 = float(getattr(rect, "y0", 0.0) or 0.0) + x1 = float(getattr(rect, "x1", 0.0) or 0.0) + y1 = float(getattr(rect, "y1", 0.0) or 0.0) + if x1 <= x0 or y1 <= y0: + continue + kept_rects.append((x0, y0, x1, y1)) + + if not kept_rects: + return [] + + figures: list[dict[str, Any]] = [] + for idxs in _cluster_indices(kept_rects, _FIGURE_CLUSTER_GAP): + if len(idxs) < _FIGURE_CLUSTER_MIN_PATHS: + continue + xs0 = [kept_rects[i][0] for i in idxs] + ys0 = [kept_rects[i][1] for i in idxs] + xs1 = [kept_rects[i][2] for i in idxs] + ys1 = [kept_rects[i][3] for i in idxs] + box = _clip_bbox(min(xs0), min(ys0), max(xs1), max(ys1), clip=page_rect) + if not _valid_bbox(box): + continue + area = max(box[2] - box[0], 0.0) * max(box[3] - box[1], 0.0) + if area < _FIGURE_MIN_AREA: + continue + if _outside_content_band( + box[1], box[3], page_h=page_h, header_y=header_y, footer_y=footer_y + ): + continue + area_ratio = area / max(page_area, 1.0) + if ( + area_ratio > _FIGURE_FULLPAGE_AREA_RATIO + and len(idxs) < _FIGURE_FULLPAGE_MAX_PATHS + ): + continue + figures.append({"kind": "figure", "bbox": box}) + return figures + + +def _probe_visual_assets( + page: Any, + page_area: float, + *, + header_y: float | None, + footer_y: float | None, +) -> dict[str, Any]: + """Collect counts + coarse asset bboxes from images / tables / drawings.""" + image_area = 0.0 + bboxes: list[dict[str, Any]] = [] + seen_image_rects: set[tuple[float, float, float, float]] = set() + page_rect = page.rect + images = page.get_images(full=True) or [] - seen: set[tuple[float, float, float, float]] = set() + image_count = len(images) for image in images: xref = image[0] try: @@ -34,56 +219,88 @@ def _image_coverage(page: Any, page_area: float) -> tuple[float, int]: 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), + box = _clip_bbox( + float(getattr(rect, "x0", 0.0) or 0.0), + float(getattr(rect, "y0", 0.0) or 0.0), + float(getattr(rect, "x1", 0.0) or 0.0), + float(getattr(rect, "y1", 0.0) or 0.0), + clip=page_rect, ) - if key in seen: + if not _valid_bbox(box): continue - seen.add(key) - area += _rect_area(rect) - return min(area / page_area, 1.0), len(images) - + key = tuple(box) + if key in seen_image_rects: + continue + seen_image_rects.add(key) + image_area += max(box[2] - box[0], 0.0) * max(box[3] - box[1], 0.0) + bboxes.append({"kind": "image", "bbox": box}) -def _table_count(page: Any) -> int: + table_count = 0 try: finder = page.find_tables() - return len(getattr(finder, "tables", []) or []) + tables = getattr(finder, "tables", []) or [] + table_count = len(tables) + for table in tables: + raw = getattr(table, "bbox", None) + if raw is None: + continue + box = _clip_bbox( + float(raw[0]), + float(raw[1]), + float(raw[2]), + float(raw[3]), + clip=page_rect, + ) + if _valid_bbox(box): + bboxes.append({"kind": "table", "bbox": box}) except Exception: - return 0 + table_count = 0 + try: + drawings = page.get_drawings() or [] + except Exception: + drawings = [] + drawings_count = len(drawings) + bboxes.extend( + _figure_bboxes_from_drawings( + drawings, + page_rect=page_rect, + page_area=page_area, + header_y=header_y, + footer_y=footer_y, + ) + ) -def _probe_one(page: Any, page_number: int) -> dict[str, Any]: + coverage = min(image_area / page_area, 1.0) if page_area > 0 else 0.0 + return { + "image_coverage": round(coverage, 4), + "image_count": image_count, + "table_count": table_count, + "drawings_count": drawings_count, + "asset_bboxes": bboxes or None, + } + + +def _probe_text_one(page: Any, page_number: int) -> dict[str, Any]: rect = page.rect area = max(_rect_area(rect), 1.0) text = page.get_text() or "" raw_text_length = len(text.strip()) - 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, + "is_blank_like": raw_text_length < 50, "text_lines_preview": top_lines(text, max_lines=30), } @worker -def _probe_worker(queue, pdf_path: str) -> None: +def _probe_text_worker(queue, pdf_path: str) -> None: import pymupdf # type: ignore[import] features: list[dict[str, Any]] = [] @@ -92,7 +309,7 @@ def _probe_worker(queue, pdf_path: str) -> None: 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)) + features.append(_probe_text_one(doc[idx], idx + 1)) finally: try: doc.close() @@ -102,23 +319,54 @@ def _probe_worker(queue, pdf_path: str) -> None: queue.put({"ok": True, "page_count": page_count, "features": features}) +@worker +def _probe_assets_worker( + queue, + pdf_path: str, + header_y: float | None, + footer_y: float | None, +) -> None: + import pymupdf # type: ignore[import] + + assets: list[dict[str, Any]] = [] + try: + doc = pymupdf.open(pdf_path) + for idx in range(int(doc.page_count)): + page = doc[idx] + area = max(_rect_area(page.rect), 1.0) + visual = _probe_visual_assets( + page, area, header_y=header_y, footer_y=footer_y + ) + assets.append({"page": idx + 1, **visual}) + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "assets": assets}) + + def probe_page_features(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + """Text-only page probe used before coarse VLM classification.""" start = time.monotonic() try: - result = run_in_child_process(_probe_worker, ctx.pdf_path, timeout=300) + result = run_in_child_process(_probe_text_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), + image_coverage=0.0, + image_count=0, + table_count=0, + drawings_count=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), + has_asset=False, is_blank_like=bool(item.get("is_blank_like")), + asset_bboxes=None, text_lines_preview=list(item.get("text_lines_preview") or []), ) for item in (result.get("features") or []) @@ -126,7 +374,8 @@ def probe_page_features(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: 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) + ctx.blackboard.global_signals["assets_probed"] = False + logger.info("[document_agent] probed text on {} pages", ctx.blackboard.page_count) return ToolResult( status="ok", payload={"page_count": ctx.blackboard.page_count}, @@ -138,3 +387,81 @@ def probe_page_features(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: error=str(exc), latency_ms=int((time.monotonic() - start) * 1000), ) + + +def probe_page_assets(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + """Asset coarse extraction after coarse VLM; drawings honor margin band.""" + start = time.monotonic() + if not ctx.blackboard.page_features: + return ToolResult( + status="error", + error="page_features missing; run text probe first", + latency_ms=int((time.monotonic() - start) * 1000), + ) + profile = ctx.blackboard.document_profile + header_y = getattr(profile, "header_y", None) if profile else None + footer_y = getattr(profile, "footer_y", None) if profile else None + try: + result = run_in_child_process( + _probe_assets_worker, + ctx.pdf_path, + header_y, + footer_y, + timeout=300, + ) + by_page = { + int(item["page"]): item for item in (result.get("assets") or []) + } + 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 + updated.append( + PageFeature( + page=feature.page, + raw_text_length=feature.raw_text_length, + text_density=feature.text_density, + image_coverage=float(visual.get("image_coverage") or 0.0), + image_count=int(visual.get("image_count") or 0), + table_count=int(visual.get("table_count") or 0), + drawings_count=int(visual.get("drawings_count") or 0), + orientation=feature.orientation, + width=feature.width, + 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 + ), + text_lines_preview=feature.text_lines_preview, + ) + ) + ctx.blackboard.page_features = sorted(updated, key=lambda f: f.page) + ctx.blackboard.global_signals["assets_probed"] = True + ctx.blackboard.global_signals["content_margins"] = { + "header_y": header_y, + "footer_y": footer_y, + } + logger.info( + "[document_agent] probed assets on {} pages (header_y={}, footer_y={})", + ctx.blackboard.page_count, + header_y, + footer_y, + ) + return ToolResult( + status="ok", + payload={ + "page_count": ctx.blackboard.page_count, + "header_y": header_y, + "footer_y": footer_y, + }, + 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_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index 714ba999..1c2737db 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 @@ -651,20 +651,20 @@ def _deterministic_no_toc_plan( *, page_count: int, max_pages: int, - low_content_pages: list[int], + blank_pages: list[int], ) -> tuple[list[tuple[int, str, str, float]], str]: - """Deterministic shard plan using low-content pages as split candidates.""" + """Deterministic shard plan using blank-like pages as split candidates.""" cuts: list[tuple[int, str, str, float]] = [] previous = 0 while page_count - previous > max_pages: target = previous + max_pages - # Look for a low-content page near the max boundary + # Look for a blank-like page near the max boundary eligible = [ - p for p in low_content_pages if previous + (max_pages - 20) < p <= target + p for p in blank_pages if previous + (max_pages - 20) < p <= target ] if eligible: chosen = max(eligible) - cuts.append((chosen, "blank_separator", f"low-content page at {chosen}", 0.5)) + cuts.append((chosen, "blank_separator", f"blank-like page at {chosen}", 0.5)) previous = chosen else: cut_page = previous + max_pages @@ -673,10 +673,10 @@ def _deterministic_no_toc_plan( return cuts, "too_large" -def _get_low_content_pages(ctx: ToolContext) -> list[int]: - """Extract low-content page numbers from page labels.""" - labels = ctx.blackboard.page_labels or [] - return sorted(label.page for label in labels if label.kind == "low_content") +def _get_blank_pages(ctx: ToolContext) -> list[int]: + """Extract blank-like page numbers from page features.""" + features = ctx.blackboard.page_features or [] + return sorted(feature.page for feature in features if feature.is_blank_like) @register_tool( @@ -781,14 +781,14 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: ) rationale = "Deterministic chapter plan (no LLM)." else: - # Path B: No TOC — purely deterministic using low-content pages - low_content_pages = _get_low_content_pages(ctx) + # Path B: No TOC — purely deterministic using blank-like pages + blank_pages = _get_blank_pages(ctx) cuts, reason = _deterministic_no_toc_plan( page_count=page_count, max_pages=max_pages, - low_content_pages=low_content_pages, + blank_pages=blank_pages, ) - rationale = "Deterministic plan from low-content page boundaries (no TOC)." + rationale = "Deterministic plan from blank-like page boundaries (no TOC)." shards = _cuts_to_shards(cuts, page_count) enabled = len(shards) > 1 diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 22008909..4460fd1d 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -737,7 +737,6 @@ def _run_hierarchy_scope( page_count=page_count, page_labels=page_labels, page_features=page_features, - tag_mode=page_memory_config.tag_mode, ) final_page_set = set(final_pages) plans = [plan for plan in plans if plan.page_index in final_page_set] diff --git a/apps/worker/app/services/page_memory/page_plan.py b/apps/worker/app/services/page_memory/page_plan.py index 0cbdb69c..234e2807 100644 --- a/apps/worker/app/services/page_memory/page_plan.py +++ b/apps/worker/app/services/page_memory/page_plan.py @@ -1,14 +1,10 @@ """Page processing plan: maps page labels to tagging strategy. Reads ``PageLabel.kind`` + ``PageFeature`` and assigns each page -one of three strategies: +one of two strategies: - ``vlm_lite`` — send the page image to VLM for summary/entities -- ``text_only`` — use raw text via page-memory-text-tag (no VLM call) - ``skip_tagging`` — blank-like page, preserve image only - -Set ``tag_mode=text`` to force all non-skip pages to ``text_only``. -Default is ``vlm``. """ from __future__ import annotations @@ -23,7 +19,6 @@ class PageProcessingStrategy(str, Enum): """Tagging strategy for a single page.""" VLM_LITE = "vlm_lite" - TEXT_ONLY = "text_only" SKIP_TAGGING = "skip_tagging" @@ -41,15 +36,12 @@ def derive_page_processing_plan( page_count: int, page_labels: list[PageLabel], page_features: list[PageFeature], - tag_mode: str = "vlm", ) -> list[PagePlan]: """Assign a processing strategy to every page. Rules: - - ``low_content`` + ``is_blank_like`` → ``skip_tagging`` - - ``tag_mode=text`` → ``text_only`` for all remaining pages - - ``table_heavy`` with sufficient raw text → ``text_only`` - - everything else (normal / image_heavy / landscape) → ``vlm_lite`` + - ``is_blank_like`` → ``skip_tagging`` + - everything else → ``vlm_lite`` Returns one ``PagePlan`` per page (1-indexed), ordered by page_index. """ @@ -60,33 +52,21 @@ def derive_page_processing_plan( for page in range(1, page_count + 1): label = label_map.get(page) feature = feature_map.get(page) - strategy, reason = _classify_page(label, feature, tag_mode=tag_mode) + strategy, reason = _classify_page(label, feature) plans.append(PagePlan(page_index=page, strategy=strategy, reason=reason)) return plans -# ── minimum raw text length for table_heavy → text_only ────────────── -_TABLE_TEXT_THRESHOLD = 200 - def _classify_page( label: PageLabel | None, feature: PageFeature | None, - *, - tag_mode: str, ) -> tuple[PageProcessingStrategy, str]: """Determine the strategy for a single page.""" kind = label.kind if label else "normal" is_blank = feature.is_blank_like if feature else False - text_len = feature.raw_text_length if feature else 0 - - if kind == "low_content" and is_blank: - return PageProcessingStrategy.SKIP_TAGGING, "low_content + blank_like" - - if tag_mode.strip().lower() == "text": - return PageProcessingStrategy.TEXT_ONLY, f"text_mode_override kind={kind}" - if kind == "table_heavy" and text_len >= _TABLE_TEXT_THRESHOLD: - return PageProcessingStrategy.TEXT_ONLY, f"table_heavy with {text_len} chars" + if is_blank: + return PageProcessingStrategy.SKIP_TAGGING, "blank_like" return PageProcessingStrategy.VLM_LITE, f"kind={kind}" diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index 4867d6a4..2204ddbd 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -2,8 +2,6 @@ For ``vlm_lite`` pages, sends the page PNG to the VLM and expects a JSON response with ``summary`` and ``keywords``. -For ``text_only`` pages, calls the existing ``summary-full`` LLM prompt -to extract summary + keywords from raw text. For ``skip_tagging`` pages, content is preserved but summary is omitted. Title detection (``tag_page_titles``) extracts outline-level headings in @@ -102,15 +100,12 @@ def _tag_one(page: PageRenderResult) -> PageTagResult: if strategy == PageProcessingStrategy.SKIP_TAGGING: return _tag_skip(page) - if strategy == PageProcessingStrategy.TEXT_ONLY: - return _tag_text_only(page) - if not model: logger.warning( - "[page_tagger] no VLM model configured for page {}; using text_only", + "[page_tagger] no VLM model configured for page {}; skipping tag", page.page_index, ) - return _tag_text_only(page) + return _tag_skip(page) return _tag_vlm_lite(page, model=model) @@ -121,10 +116,9 @@ def _tag_one(page: PageRenderResult) -> PageTagResult: results = [cast(PageTagResult, g.value) for g in greenlets] vlm_calls = sum(1 for r in results if r.strategy_used == "vlm_lite") logger.info( - "[page_tagger] tagged {} pages ({} VLM calls, {} text_only, {} skipped, {} failed) concurrency={}", + "[page_tagger] tagged {} pages ({} VLM calls, {} skipped, {} failed) concurrency={}", len(results), vlm_calls, - sum(1 for r in results if r.strategy_used == "text_only"), sum(1 for r in results if r.strategy_used == "skip_tagging"), len(greenlets) - len(results), resolved_max_concurrent, @@ -146,85 +140,6 @@ def _tag_skip(page: PageRenderResult) -> PageTagResult: ) -def _tag_text_only( - page: PageRenderResult, -) -> PageTagResult: - """Extract summary + entities from raw page text via page-memory-text-tag. - - Uses the same output spec ({summary, entities}) as the VLM path so - downstream consumers need no special-casing. Returns an EMPTY-marked - result when the page has no extractable text. - """ - raw = page.raw_text.strip() - if not raw: - return PageTagResult( - page_index=page.page_index, - summary="EMPTY", - keywords=[], - strategy_used="text_only", - ) - - model = os.environ.get("NORMOL_MODEL", "deepseek-v4-flash") - text_input = raw[:4000] # cap to avoid token overflow - prompt, temperature, top_p, max_tokens = build_prompt( - "page-memory-text-tag", - "", - "", - paras={"max_tokens": 600, "page_text": text_input}, - ) - - try: - from shared.services.ai.llm_overrides import get_text_client - - client, model = get_text_client(requested_model=model) - raw_response, usage = client.chat_completion_with_usage( - messages=cast(Any, [{"role": "user", "content": prompt}]), - model=model, - temperature=temperature, - top_p=top_p, - max_tokens=max_tokens, - response_format={"type": "json_object"}, - usage_task="page_memory.text_tag", - ) - except UnavailableException: - raise - except Exception as exc: - logger.warning( - "[page_tagger] text_only LLM failed for page {}: {}", - page.page_index, - exc, - ) - return PageTagResult( - page_index=page.page_index, - summary="", - keywords=[], - strategy_used="text_only", - ) - - try: - data = json.loads(raw_response) - except json.JSONDecodeError: - return PageTagResult( - page_index=page.page_index, - summary="", - keywords=[], - strategy_used="text_only", - ) - - entities_raw = data.get("entities") or [] - entities = [ - e for e in entities_raw - if isinstance(e, dict) and e.get("text") - ] - return PageTagResult( - page_index=page.page_index, - summary=str(data.get("summary") or "").strip(), - keywords=[str(e["text"]) for e in entities], - entities=entities, - strategy_used="text_only", - ) - - def _tag_vlm_lite( page: PageRenderResult, *, @@ -233,10 +148,10 @@ def _tag_vlm_lite( """Send page PNG to the VLM via the unified engine.""" if not page.image_path or not os.path.exists(page.image_path): logger.warning( - "[page_tagger] no PNG for page {}; text_only fallback", + "[page_tagger] no PNG for page {}; skipping tag", page.page_index, ) - return _tag_text_only(page) + return _tag_skip(page) result = summarize( mode="page", 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 b45a26af..8901c127 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -56,7 +56,9 @@ def _page_feature(page: int = 1) -> PageFeature: orientation="portrait", width=72.0, height=72.0, + has_asset=False, is_blank_like=False, + asset_bboxes=None, text_lines_preview=["Section 1"], ) diff --git a/packages/shared-python/shared/models/schemas/page_memory_config.py b/packages/shared-python/shared/models/schemas/page_memory_config.py index f8134ded..95c2122d 100644 --- a/packages/shared-python/shared/models/schemas/page_memory_config.py +++ b/packages/shared-python/shared/models/schemas/page_memory_config.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import asdict, dataclass -from typing import Literal, Self +from typing import Self @dataclass(frozen=True) @@ -15,7 +15,6 @@ class PageMemoryConfig: tag_concurrency: int = 4 title_detection_concurrency: int = 3 node_assembly_concurrency: int = 3 - tag_mode: Literal["vlm", "text"] = "vlm" fine_min_pages: int = 4 hierarchy_model: str | None = None hierarchy_max_tokens: int = 2000 @@ -60,10 +59,6 @@ def from_mapping(cls, value: object) -> Self: return cls.default() default = cls.default() - tag_mode = str(value.get("tag_mode", default.tag_mode)).strip().lower() - resolved_tag_mode: Literal["vlm", "text"] = ( - "text" if tag_mode == "text" else "vlm" - ) return cls( max_pages=_as_int(value.get("max_pages"), default.max_pages), scope_concurrency=_as_int( @@ -82,7 +77,6 @@ def from_mapping(cls, value: object) -> Self: value.get("node_assembly_concurrency"), default.node_assembly_concurrency, ), - tag_mode=resolved_tag_mode, fine_min_pages=_as_int( value.get("fine_min_pages"), default.fine_min_pages, diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index 27814c2c..f5152cab 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -483,36 +483,6 @@ def build_prompt(task, texts, query, **kwargs): - Return ONLY the JSON object, with no markdown fences or extra text. """ - elif task == "page-memory-text-tag": - temperature = 0 - top_p = 0.01 - max_tokens = kwargs.get("paras", {}).get("max_tokens", 600) - entity_line = _entity_instruction() - page_text = kwargs.get("paras", {}).get("page_text", "") - prompt = f"""\ - You are annotating a single document page for a document memory system. - The following is the extracted text from the page: - \"\"\" - {page_text} - \"\"\" - - Return one strict JSON object with exactly these keys: - - {{ - "summary": "", - "entities": [{{"text": "", "type": ""}}] - }} - - Rules: - - "summary": describe the main content in a few sentences, in the same - language as the text. If the text contains a table, state its topic - and key columns; if it describes a figure or chart, describe what it - depicts and any standout values. If the text is empty or carries no - meaningful content, set summary to an empty string. - {entity_line} - - Return ONLY the JSON object, with no markdown fences or extra text. - """ - elif task == "page-memory-vlm-title": temperature = 0 top_p = 0.01 @@ -776,15 +746,27 @@ def build_prompt(task, texts, query, **kwargs): numbers. Rules: - - "table": data arranged in clear rows and columns - grid lines, cell - borders, or strongly aligned cells (data tables, forms, financial tables, - appendix tables). - - "figure": any non-table visual asset - bar/line/pie/scatter charts, - plots, diagrams, flowcharts, architecture drawings, schematics, or embedded - images. + - "table": ONLY a conventional data table with visible grid and strongly regular cell alignment. + Require ALL of: + (1) an explicit header row and/or header column that labels the fields. + (2) one intact axis-aligned rectangular footprint: all four corners of + the table body are present, every data row spans that full width, and + the cell grid fills the rectangle without cutouts, protruding corner + panels, or L-shaped outlines. + typical table cases: forms, financial tables, data rows with field headers. + + - Do NOT mark as "table": process boards, flowchart-like matrices, cards + connected by arrows, multi-column visual layouts, comparison panels, or + any region whose meaning depends on icons/arrows/color blocks rather than + plain headered cells. Those must be "figure". + - If unsure whether it meets this bar, prefer "figure". + + - "figure": any non-table visual asset - charts, plots, diagrams, + flowcharts, architecture drawings, schematics, embedded images, and the + table-like visuals excluded above. - Prefer one bbox for the whole figure. When multiple visual parts clearly - form one composition (shared caption or a multi-panel explanation of the same concept/process), - return them as a single figure, not separate images. + form one composition (shared caption or a multi-panel explanation of the + same concept/process), return them as a single figure, not separate images. - Treat a flowchart or process diagram as one figure, including its nodes, edges, labels, and legend when they belong together. From 9a395b0aa6bc6cf1353e954d327bbddc57819655 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 24 Jul 2026 12:13:48 +0800 Subject: [PATCH 5/6] fix: tighten image bbox key typing for pyright Co-authored-by: Cursor --- .../app/services/document_agent/tools/probe_page_features.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 f4e3ff8a..ddc26d48 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 @@ -228,7 +228,7 @@ def _probe_visual_assets( ) if not _valid_bbox(box): continue - key = tuple(box) + key = (box[0], box[1], box[2], box[3]) if key in seen_image_rects: continue seen_image_rects.add(key) From c00aef5315bbbb0f6f28a26d29100151131d59a6 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 24 Jul 2026 17:22:33 +0800 Subject: [PATCH 6/6] refactor: streamline document agent structure and enhance profiling capabilities - Removed the `ProfileAgent` class to simplify the document agent architecture. - Updated `ProfileCoordinator` methods to include a `skip_shard_plan` parameter, allowing for more flexible profiling without LLM shard decisions. - Enhanced `PageAnatomyMap` to include `toc_hierarchies` for improved TOC handling. - Cleaned up unused fields and methods related to asset probing and text previews, ensuring a more efficient codebase. This commit improves the overall structure and efficiency of the document agent, facilitating better profiling and asset management. --- .../app/services/document_agent/__init__.py | 14 - .../services/document_agent/coordinator.py | 43 ++-- .../app/services/document_agent/manifest.py | 3 +- .../services/document_agent/profile_agent.py | 60 ----- .../tools/classify_page_kinds.py | 1 - .../tools/extract_toc_with_boundaries.py | 10 - .../tools/probe_page_features.py | 4 - .../orchestration/format_adapters.py | 4 + .../document_parser/profiling/doc_profiler.py | 29 ++- .../services/page_memory/memory_service.py | 21 +- .../test_doc_profile_anatomy_contract.py | 241 +++++++++++++++--- .../test_page_memory_asset_java_contract.py | 43 ++++ .../shared/core/config/storage.py | 8 +- 13 files changed, 321 insertions(+), 160 deletions(-) delete mode 100644 apps/worker/app/services/document_agent/profile_agent.py diff --git a/apps/worker/app/services/document_agent/__init__.py b/apps/worker/app/services/document_agent/__init__.py index 59e0afc9..6baae446 100644 --- a/apps/worker/app/services/document_agent/__init__.py +++ b/apps/worker/app/services/document_agent/__init__.py @@ -1,7 +1,5 @@ """Page anatomy agent for hierarchy-first PDF profiling.""" -from typing import TYPE_CHECKING - from app.services.document_agent.manifest import ( PageAnatomyMap, PageFeature, @@ -9,21 +7,9 @@ ShardPlan, ) -if TYPE_CHECKING: - from app.services.document_agent.profile_agent import ProfileAgent - - -def __getattr__(name: str): - if name == "ProfileAgent": - from app.services.document_agent.profile_agent import ProfileAgent - - return ProfileAgent - raise AttributeError(name) - __all__ = [ "PageAnatomyMap", "PageFeature", "PageLabel", - "ProfileAgent", "ShardPlan", ] diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 2f25f6a1..2f8c699f 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -97,13 +97,6 @@ def __init__( self.round_index = 0 self._planner_cache: tuple[DocumentProfile, Any, ToolResult] | None = None - def run(self) -> PageAnatomyMap: - try: - return self._run_structural() - except Exception as exc: - self._record_failure(exc) - raise - def run_coarse(self) -> DocumentProfile: try: return self._run_coarse() @@ -111,9 +104,9 @@ def run_coarse(self) -> DocumentProfile: self._record_failure(exc) raise - def run_structural(self) -> PageAnatomyMap: + def run_structural(self, *, skip_shard_plan: bool = False) -> PageAnatomyMap: try: - return self._run_structural() + return self._run_structural(skip_shard_plan=skip_shard_plan) except Exception as exc: self._record_failure(exc) raise @@ -155,7 +148,7 @@ def _run_coarse(self) -> DocumentProfile: self._ensure_asset_probe() return profile - def _run_structural(self) -> PageAnatomyMap: + def _run_structural(self, *, skip_shard_plan: bool = False) -> PageAnatomyMap: self.state = DocumentAgentState.RUNNING if not self.blackboard.page_features: self._run_bootstrap() @@ -167,14 +160,21 @@ def _run_structural(self) -> PageAnatomyMap: actor="planner" ) self._ensure_asset_probe() - 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 skip_shard_plan: + # Page-memory oversized path never consumes shard_plan; only + # build_anatomy_map's invariant needs a non-empty plan. + self._apply_single_shard_placeholder() + else: + 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}" + ) anatomy = build_anatomy_map(self.ctx) self._persist_ready_anatomy(anatomy) return anatomy @@ -218,9 +218,7 @@ def _run_lightweight_anatomy( # it. Populate a single-shard placeholder to skip the LLM shard # decision + H2 refinement (kept global for chunk-track oversized # MinerU sharding). - self.blackboard.shard_plan = single_shard_plan( - self.blackboard.page_count - ) + self._apply_single_shard_placeholder() else: result = REGISTRY.dispatch("propose.shard_plan", self.ctx, {}) self.trace.record_step( @@ -238,6 +236,9 @@ def _run_lightweight_anatomy( self._persist_ready_anatomy(anatomy) return anatomy + def _apply_single_shard_placeholder(self) -> None: + self.blackboard.shard_plan = single_shard_plan(self.blackboard.page_count) + def _persist_ready_anatomy(self, anatomy: PageAnatomyMap) -> None: persist_result = persist_anatomy_map(self.ctx, {}) self.trace.record_step( diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index 01762066..d078f9d5 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -30,7 +30,6 @@ class PageFeature: is_blank_like: bool # PDF-space boxes for detected assets; None when none were extracted. asset_bboxes: list[dict[str, Any]] | None = None - text_lines_preview: list[str] = field(default_factory=list) def to_dict(self) -> dict[str, Any]: return asdict(self) @@ -236,6 +235,7 @@ class PageAnatomyMap: def to_dict(self) -> dict[str, Any]: return { "version": self.version, + "toc_hierarchies": self.toc_hierarchies, "job_id": self.job_id, "file_path": self.file_path, "page_count": self.page_count, @@ -247,7 +247,6 @@ def to_dict(self) -> dict[str, Any]: "document_profile": self.document_profile.to_dict() if self.document_profile else None, - "toc_hierarchies": self.toc_hierarchies, "toc_page_offset": self.toc_page_offset, "global_signals": dict(self.global_signals), "trace_summary": dict(self.trace_summary), diff --git a/apps/worker/app/services/document_agent/profile_agent.py b/apps/worker/app/services/document_agent/profile_agent.py deleted file mode 100644 index c56e18c1..00000000 --- a/apps/worker/app/services/document_agent/profile_agent.py +++ /dev/null @@ -1,60 +0,0 @@ -"""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 DocumentProfile, 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() - - def run_coarse( - self, - file_path: str, - job_id: str, - *, - output_dir: str | None = None, - db: Any | None = None, - ) -> DocumentProfile: - 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_coarse() 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 923a816c..d3a11d52 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,7 +39,6 @@ def classify_page_kinds(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: "confidence": label.confidence, "evidence": label.evidence, "raw_text_length": feature.raw_text_length if feature else None, - "text_preview": (feature.text_lines_preview[:4] if feature else []), } ) return ToolResult( 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 af237901..a2de3d29 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 @@ -447,16 +447,6 @@ def extract_toc_with_boundaries( "batch_meta": batch_meta, } - # 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) - # Build toc_ranges from confirmed TOC pages for summary toc_ranges_out: list[list[int]] = [] if toc_hierarchies: 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 ddc26d48..6e21bb8e 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 @@ -9,7 +9,6 @@ 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_parser.formats.pdf.pymupdf_subprocess import ( run_in_child_process, worker, @@ -295,7 +294,6 @@ def _probe_text_one(page: Any, page_number: int) -> dict[str, Any]: "width": round(float(rect.width), 2), "height": round(float(rect.height), 2), "is_blank_like": raw_text_length < 50, - "text_lines_preview": top_lines(text, max_lines=30), } @@ -367,7 +365,6 @@ def probe_page_features(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: has_asset=False, is_blank_like=bool(item.get("is_blank_like")), asset_bboxes=None, - text_lines_preview=list(item.get("text_lines_preview") or []), ) for item in (result.get("features") or []) ] @@ -435,7 +432,6 @@ def probe_page_assets(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: if isinstance(visual.get("asset_bboxes"), list) else None ), - text_lines_preview=feature.text_lines_preview, ) ) ctx.blackboard.page_features = sorted(updated, key=lambda f: f.page) diff --git a/apps/worker/app/services/document_parser/orchestration/format_adapters.py b/apps/worker/app/services/document_parser/orchestration/format_adapters.py index 7b3446ec..5a26a45b 100644 --- a/apps/worker/app/services/document_parser/orchestration/format_adapters.py +++ b/apps/worker/app/services/document_parser/orchestration/format_adapters.py @@ -138,6 +138,10 @@ def parse(self, session: ParseSession) -> ParseOutput: @dataclass(frozen=True) class PptxParseAdapter: # Deprecated: prefer page_memory track for PPTX (parse_track="page_memory"). + # TODO(pptx): convert PPTX→PDF first, then reuse the standard PDF PROFILE + # path instead of a separate PPTX PROFILE; keep page_memory input schema + # stable (fields may expand later). page_memory/normalizer.py already + # converts before PROFILE. document_format: object def parse(self, session: ParseSession) -> ParseOutput: diff --git a/apps/worker/app/services/document_parser/profiling/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py index 1dbf6dab..a0a3077b 100644 --- a/apps/worker/app/services/document_parser/profiling/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -42,13 +42,14 @@ def profile_document( filename: File name (used to infer type) job_id: Parse job id for profile trace artifacts output_dir: Parser output directory - skip_shard_plan: When True, the lightweight anatomy stage skips the - LLM shard decision (+ H2 refinement) and populates a single-shard - placeholder instead. Used by the page-memory track, which never - consumes the shard plan. Chunk-track keeps the default (False). + skip_shard_plan: When True, lightweight and structural anatomy skip + LLM/ReAct shard planning and populate a single-shard placeholder. + Used by the page-memory track, which never consumes the shard plan. + Chunk-track keeps the default (False) so oversized MinerU sharding + still receives a real plan. oversized_policy: Controls oversized PDF admission. ``chunk`` applies - the legacy MinerU shard gate, while ``page_memory`` lets the - page-memory track continue to structural profiling. + the MinerU shard gate, while ``page_memory`` lets the page-memory + track continue to structural profiling. Returns: ParserDocumentProfile @@ -112,7 +113,13 @@ def _profile_pdf_with_db( ) -> ParserDocumentProfile: profile_job_id = job_id or filename agent_output_dir = os.path.join(output_dir, "_doc_agent") if output_dir else None - page_toc_enabled = settings.PDF_PROFILE_TOC_ENABLED + # Page-memory sections are anchored on the TOC (page-based VLM TOC pipeline), + # so TOC profiling is mandatory for that track regardless of the global + # PDF_PROFILE_TOC_ENABLED flag (which only gates the optional chunk-track + # TOC profiling that can otherwise fall back to MinerU markdown headings). + page_toc_enabled = ( + oversized_policy == "page_memory" or settings.PDF_PROFILE_TOC_ENABLED + ) coordinator = ProfileCoordinator( pdf_path=file_path, job_id=profile_job_id, @@ -152,7 +159,9 @@ def _profile_pdf_with_db( raise_if_oversized_pdf_not_supported(page_count=profile.page_count) if not profile.is_atlas: try: - profile.anatomy = coordinator.run_structural() + profile.anatomy = coordinator.run_structural( + skip_shard_plan=skip_shard_plan + ) profile.toc = _map_toc_profile(coordinator) except Exception as exc: if oversized_policy == "page_memory": @@ -165,8 +174,12 @@ def _profile_pdf_with_db( original_exception=exc, ) from exc else: + # TODO(page_memory): oversized atlas skips anatomy, so coarse + # has_asset / page_features never reach page_memory (Root fallback). profile.toc = _map_toc_profile(coordinator) else: + # TODO(page_memory): non-oversized atlas skips anatomy; coarse + # has_asset / page_features never reach page_memory via profile.anatomy. if not profile.is_atlas: profile.anatomy = coordinator.run_lightweight_anatomy( skip_shard_plan=skip_shard_plan diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 4460fd1d..0c09005b 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -771,10 +771,13 @@ def _run_hierarchy_scope( assets_by_page: dict[int, list[Any]] = {} if asset_extraction_enabled and asset_max_pages > 0: + asset_rendered = _select_rendered_pages_with_assets( + rendered, page_features + ) with stage_timer("page_memory.assets", page_count=asset_max_pages): assets_by_page = extract_page_assets_from_renders( pdf_path=pdf_path, - rendered_pages=rendered, + rendered_pages=asset_rendered, output_dir=output_dir, model_name=page_memory_config.asset_model, budget=None, @@ -818,6 +821,22 @@ def _run_hierarchy_scope( ) +# ── helpers ─────────────────────────────────────────────────────────── + + +def _select_rendered_pages_with_assets( + rendered: list[Any], + page_features: list[Any], +) -> list[Any]: + """Keep only rendered pages that coarse profile marked ``has_asset``.""" + asset_pages = { + int(getattr(feature, "page", 0) or 0) + for feature in page_features + if getattr(feature, "has_asset", False) + } + return [item for item in rendered if item.page_index in asset_pages] + + # ── whole_doc builder (PR3, unchanged) ──────────────────────────────── 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 8901c127..839b6165 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -59,10 +59,33 @@ def _page_feature(page: int = 1) -> PageFeature: has_asset=False, is_blank_like=False, asset_bboxes=None, - text_lines_preview=["Section 1"], ) +def _seed_preprobed_pages( + coordinator: ProfileCoordinator, + *, + page_count: int, + pages: list[int] | None = None, +) -> None: + """Seed text-bootstrap state for coordinator tests without a real PDF. + + Marks assets as already probed so `_ensure_asset_probe` does not spawn a + PyMuPDF child against a missing fixture path. + """ + probed_pages = pages or list(range(1, page_count + 1)) + coordinator.blackboard.page_count = page_count + coordinator.blackboard.page_features = [_page_feature(page) for page in probed_pages] + coordinator.blackboard.page_labels = [ + PageLabel(page=page, kind="normal", confidence=1.0) for page in probed_pages + ] + coordinator.blackboard.doc_stats = {"page_count": page_count} + coordinator.blackboard.global_signals["page_kind_counts"] = { + "normal": page_count + } + coordinator.blackboard.global_signals["assets_probed"] = True + + def test_toc_anchor_text_scan_matches_full_page_and_cross_line_keywords() -> None: late_lines = [f"body line {idx}" for idx in range(60)] + ["目录"] split_lines = ["Table of", "Con", "tents"] @@ -114,14 +137,7 @@ def test_run_lightweight_anatomy_builds_single_shard_without_planner_llm( output_dir=str(output_dir), settings={"shard_threshold": 200}, ) - coordinator.blackboard.page_count = 2 - coordinator.blackboard.page_features = [_page_feature(1), _page_feature(2)] - coordinator.blackboard.page_labels = [ - PageLabel(page=1, kind="normal", confidence=1.0), - PageLabel(page=2, kind="normal", confidence=1.0), - ] - coordinator.blackboard.doc_stats = {"page_count": 2} - coordinator.blackboard.global_signals["page_kind_counts"] = {"normal": 2} + _seed_preprobed_pages(coordinator, page_count=2) coordinator.blackboard.document_profile = DocumentProfile( is_scanned=False, category="Research Report", @@ -137,10 +153,65 @@ def test_run_lightweight_anatomy_builds_single_shard_without_planner_llm( assert anatomy.shard_plan.shards[0].page_end == 2 assert anatomy.toc_result.method == "none" assert (output_dir / "anatomy_map.json").exists() + anatomy_data = json.loads( + (output_dir / "anatomy_map.json").read_text(encoding="utf-8") + ) + assert list(anatomy_data)[:2] == ["version", "toc_hierarchies"] + assert "text_lines_preview" 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"] +def test_run_coarse_runs_asset_probe_after_planner(monkeypatch, tmp_path: Path) -> None: + coordinator = ProfileCoordinator( + pdf_path=str(tmp_path / "doc.pdf"), + job_id="job-asset-probe-after-coarse", + output_dir=str(tmp_path / "profile"), + ) + (tmp_path / "profile").mkdir() + coordinator.blackboard.page_count = 2 + coordinator.blackboard.page_features = [_page_feature(1), _page_feature(2)] + coordinator.blackboard.page_labels = [ + PageLabel(page=1, kind="normal", confidence=1.0), + PageLabel(page=2, kind="normal", confidence=1.0), + ] + coordinator.blackboard.doc_stats = {"page_count": 2} + coordinator.blackboard.global_signals["page_kind_counts"] = {"normal": 2} + + calls: list[str] = [] + + def fake_propose(_self): + calls.append("planner") + return ( + DocumentProfile( + is_scanned=False, + category="Research Report", + routing_category=PdfRoutingCategory.GENERIC.value, + ), + None, + ToolResult(status="ok", payload={}), + ) + + def fake_probe_page_assets(ctx, _args): + calls.append("probe.page_assets") + ctx.blackboard.global_signals["assets_probed"] = True + return ToolResult(status="ok", payload={"page_count": 2}) + + def fake_aggregate(ctx, _args): + calls.append("aggregate.doc_stats") + return ToolResult(status="ok", payload={}) + + monkeypatch.setattr(coordinator_module.ProfilePlanner, "propose", fake_propose) + monkeypatch.setattr(coordinator_module, "probe_page_assets", fake_probe_page_assets) + monkeypatch.setattr(coordinator_module, "aggregate_doc_stats", fake_aggregate) + + profile = coordinator.run_coarse() + + assert profile.category == "Research Report" + assert calls == ["planner", "probe.page_assets", "aggregate.doc_stats"] + assert coordinator.blackboard.global_signals["assets_probed"] is True + + def test_parse_run_recorder_doc_profile_uses_final_anatomy_toc() -> None: recorder = ParseRunRecorder(job_id="job-doc-profile") recorder.set_doc_profile( @@ -236,14 +307,7 @@ def test_run_structural_retries_transient_confirm_failed_toc_result( output_dir=str(tmp_path / "profile"), ) (tmp_path / "profile").mkdir() - coordinator.blackboard.page_count = 3 - coordinator.blackboard.page_features = [_page_feature(1), _page_feature(2)] - coordinator.blackboard.page_labels = [ - PageLabel(page=1, kind="normal", confidence=1.0), - PageLabel(page=2, kind="normal", confidence=1.0), - ] - coordinator.blackboard.doc_stats = {"page_count": 3} - coordinator.blackboard.global_signals["page_kind_counts"] = {"normal": 3} + _seed_preprobed_pages(coordinator, page_count=3, pages=[1, 2]) coordinator.blackboard.document_profile = DocumentProfile( is_scanned=False, category="Prospectus", @@ -325,6 +389,71 @@ def run(self): assert anatomy.toc_result.toc_pages == [17] +def test_run_structural_skip_shard_plan_uses_placeholder_without_executor( + monkeypatch, + tmp_path: Path, +) -> None: + coordinator = ProfileCoordinator( + pdf_path=str(tmp_path / "oversized.pdf"), + job_id="job-structural-skip-shard", + output_dir=str(tmp_path / "profile"), + ) + (tmp_path / "profile").mkdir() + _seed_preprobed_pages(coordinator, page_count=4) + coordinator.blackboard.document_profile = DocumentProfile( + is_scanned=False, + category="Prospectus", + routing_category=PdfRoutingCategory.GENERIC.value, + ) + coordinator.blackboard.toc_result = TocResult( + toc_pages=[1], + method="vlm_batch", + notes="ok", + ) + coordinator.blackboard.toc_hierarchies = [ + {"toc_range": [1, 1], "toc_range_unit": "page", "toc_tree": {}} + ] + + monkeypatch.setattr( + coordinator, + "_run_toc_extraction_pipeline", + lambda: (_ for _ in ()).throw( + AssertionError("existing TOC should not be re-extracted") + ), + ) + monkeypatch.setattr( + coordinator, + "_persist_ready_anatomy", + lambda _anatomy: None, + ) + monkeypatch.setattr( + coordinator_module.ProfilePlanner, + "propose", + lambda self: ( + coordinator.blackboard.document_profile, + None, + ToolResult(status="ok", payload={}), + ), + ) + + class BoomExecutor: + def __init__(self, *_args, **_kwargs) -> None: + raise AssertionError("ReActExecutor must not run when skip_shard_plan") + + def run(self): # pragma: no cover + raise AssertionError("unreachable") + + monkeypatch.setattr(coordinator_module, "ReActExecutor", BoomExecutor) + + anatomy = coordinator.run_structural(skip_shard_plan=True) + + assert anatomy.shard_plan.enabled is False + assert len(anatomy.shard_plan.shards) == 1 + assert anatomy.shard_plan.shards[0].page_start == 1 + assert anatomy.shard_plan.shards[0].page_end == 4 + assert anatomy.toc_result.toc_pages == [1] + + def test_run_structural_trusts_rejected_all_toc_and_fails_open( monkeypatch, tmp_path: Path, @@ -335,14 +464,7 @@ def test_run_structural_trusts_rejected_all_toc_and_fails_open( output_dir=str(tmp_path / "profile"), ) (tmp_path / "profile").mkdir() - coordinator.blackboard.page_count = 3 - coordinator.blackboard.page_features = [_page_feature(1), _page_feature(2)] - coordinator.blackboard.page_labels = [ - PageLabel(page=1, kind="normal", confidence=1.0), - PageLabel(page=2, kind="normal", confidence=1.0), - ] - coordinator.blackboard.doc_stats = {"page_count": 3} - coordinator.blackboard.global_signals["page_kind_counts"] = {"normal": 3} + _seed_preprobed_pages(coordinator, page_count=3, pages=[1, 2]) coordinator.blackboard.document_profile = DocumentProfile( is_scanned=False, category="Prospectus", @@ -428,14 +550,7 @@ def test_run_coarse_runs_toc_before_planner_for_oversized_and_reuses_planner( settings={"toc_before_coarse": True}, ) (tmp_path / "profile").mkdir() - coordinator.blackboard.page_count = 3 - coordinator.blackboard.page_features = [_page_feature(1), _page_feature(2)] - coordinator.blackboard.page_labels = [ - PageLabel(page=1, kind="normal", confidence=1.0), - PageLabel(page=2, kind="normal", confidence=1.0), - ] - coordinator.blackboard.doc_stats = {"page_count": 3} - coordinator.blackboard.global_signals["page_kind_counts"] = {"normal": 3} + _seed_preprobed_pages(coordinator, page_count=3, pages=[1, 2]) calls: list[str] = [] @@ -660,6 +775,57 @@ def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): assert profile.anatomy is fake_anatomy +def test_page_memory_forces_toc_profiling_despite_kill_switch( + monkeypatch, + tmp_path: Path, +) -> None: + fake_anatomy = object() + init_settings: list[dict[str, object]] = [] + + class FakeCoordinator: + def __init__(self, **kwargs) -> None: + self.calls: list[str] = [] + init_settings.append(kwargs["settings"]) + self.blackboard = SimpleNamespace( + page_count=2, + doc_stats={"page_count": 2}, + global_signals={}, + toc_result=None, + toc_hierarchies=None, + ) + + def run_coarse(self) -> DocumentProfile: + self.calls.append("run_coarse") + self.blackboard.toc_result = TocResult(method="none") + return DocumentProfile( + is_scanned=False, + category="Research Report", + routing_category=PdfRoutingCategory.GENERIC.value, + ) + + def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): + self.calls.append("run_lightweight_anatomy") + return fake_anatomy + + monkeypatch.setattr(doc_profiler, "ProfileCoordinator", FakeCoordinator) + monkeypatch.setattr(doc_profiler.settings, "MAX_PDF_PAGE_LIMIT", 200) + # Global kill switch OFF; page_memory track must still enable TOC. + monkeypatch.setattr(doc_profiler.settings, "PDF_PROFILE_TOC_ENABLED", False) + + profile = profile_document( + str(tmp_path / "standard.pdf"), + "standard.pdf", + job_id="job-page-memory-toc-forced", + output_dir=str(tmp_path), + skip_shard_plan=True, + oversized_policy="page_memory", + ) + + assert init_settings[0]["toc_profile_enabled"] is True + assert init_settings[0]["toc_before_coarse"] is True + assert profile.anatomy is fake_anatomy + + def test_page_memory_profile_bypasses_chunk_oversized_gate( monkeypatch, tmp_path: Path, @@ -687,8 +853,9 @@ def run_coarse(self) -> DocumentProfile: routing_category=PdfRoutingCategory.GENERIC.value, ) - def run_structural(self): + def run_structural(self, *, skip_shard_plan: bool = False): self.calls.append("run_structural") + self.skip_shard_plan = skip_shard_plan return fake_anatomy monkeypatch.setattr(doc_profiler, "ProfileCoordinator", FakeCoordinator) @@ -700,11 +867,13 @@ def run_structural(self): "oversized.pdf", job_id="job-page-memory-oversized", output_dir=str(tmp_path), + skip_shard_plan=True, oversized_policy="page_memory", ) assert profile.anatomy is fake_anatomy assert fake_instances[0].calls == ["run_coarse", "run_structural"] + assert fake_instances[0].skip_shard_plan is True def test_standard_pdf_profile_maps_page_toc_evidence( @@ -822,10 +991,10 @@ def run_coarse(self) -> DocumentProfile: routing_category=PdfRoutingCategory.ATLAS.value, ) - def run_structural(self): + def run_structural(self, *, skip_shard_plan: bool = False): raise AssertionError("oversized atlas should not run structural anatomy") - def run_lightweight_anatomy(self): + def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): raise AssertionError("oversized atlas should not run lightweight anatomy") monkeypatch.setattr(doc_profiler, "ProfileCoordinator", FakeCoordinator) diff --git a/apps/worker/tests/contract/test_page_memory_asset_java_contract.py b/apps/worker/tests/contract/test_page_memory_asset_java_contract.py index 78f13946..3d809c3d 100644 --- a/apps/worker/tests/contract/test_page_memory_asset_java_contract.py +++ b/apps/worker/tests/contract/test_page_memory_asset_java_contract.py @@ -2,6 +2,7 @@ import json import os +from types import SimpleNamespace os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") @@ -175,3 +176,45 @@ def test_page_assets_adds_java_home_to_path(monkeypatch, tmp_path) -> None: assert page_assets._has_working_java() is True # noqa: SLF001 assert os.environ["PATH"].split(os.pathsep)[0] == str(java_bin) + + +def test_select_rendered_pages_with_assets_keeps_only_has_asset_pages() -> None: + from app.services.page_memory.memory_service import ( + _select_rendered_pages_with_assets, + ) + + rendered = [ + PageRenderResult( + page_index=1, + image_path="/tmp/1.png", + raw_text="", + width=10, + height=10, + is_landscape=False, + ), + PageRenderResult( + page_index=2, + image_path="/tmp/2.png", + raw_text="", + width=10, + height=10, + is_landscape=False, + ), + PageRenderResult( + page_index=3, + image_path="/tmp/3.png", + raw_text="", + width=10, + height=10, + is_landscape=False, + ), + ] + page_features = [ + SimpleNamespace(page=1, has_asset=False), + SimpleNamespace(page=2, has_asset=True), + SimpleNamespace(page=3, has_asset=True), + ] + + selected = _select_rendered_pages_with_assets(rendered, page_features) + + assert [item.page_index for item in selected] == [2, 3] diff --git a/packages/shared-python/shared/core/config/storage.py b/packages/shared-python/shared/core/config/storage.py index 1706b2d2..f8a31c39 100644 --- a/packages/shared-python/shared/core/config/storage.py +++ b/packages/shared-python/shared/core/config/storage.py @@ -82,9 +82,11 @@ class StorageConfig(BaseModel): PDF_PROFILE_TOC_ENABLED: bool = Field( default=False, description=( - "Enable page-owned PDF TOC profiling during parser-entry DOC_PROFILE. " - "When disabled, PDF parsing treats documents as no-TOC and does not " - "fall back to Markdown TOC detection." + "Enable page-owned PDF TOC profiling for the CHUNK track during " + "parser-entry DOC_PROFILE. When disabled, chunk-track PDF parsing " + "treats documents as no-TOC and does not fall back to Markdown TOC " + "detection. NOTE: the page-memory track always forces TOC profiling " + "on regardless of this flag, since its sections are TOC-anchored." ), ) MINERU_SHARD_CONCURRENCY: int = Field(