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/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/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..2f8c699f 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 @@ -96,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() @@ -110,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 @@ -151,9 +145,10 @@ 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: + def _run_structural(self, *, skip_shard_plan: bool = False) -> PageAnatomyMap: self.state = DocumentAgentState.RUNNING if not self.blackboard.page_features: self._run_bootstrap() @@ -164,14 +159,22 @@ def _run_structural(self) -> PageAnatomyMap: profile, initial_decision, _planner_result = self._propose_profile( actor="planner" ) - 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}") + self._ensure_asset_probe() + 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 @@ -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( @@ -214,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( @@ -234,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( @@ -284,6 +289,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..d078f9d5 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,8 +26,10 @@ class PageFeature: orientation: Literal["portrait", "landscape"] width: float height: float + has_asset: bool is_blank_like: bool - text_lines_preview: list[str] = field(default_factory=list) + # PDF-space boxes for detected assets; None when none were extracted. + asset_bboxes: list[dict[str, Any]] | None = None def to_dict(self) -> dict[str, Any]: return asdict(self) @@ -52,6 +54,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) @@ -229,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, @@ -240,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/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/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 fb60765d..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 @@ -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,10 +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 []), } ) 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 678557fa..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 @@ -1,19 +1,30 @@ -"""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 -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, ) 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 +32,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 +218,87 @@ 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 = (box[0], box[1], box[2], box[3]) + 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, - "text_lines_preview": top_lines(text, max_lines=30), + "is_blank_like": raw_text_length < 50, } @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 +307,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,31 +317,62 @@ 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")), - text_lines_preview=list(item.get("text_lines_preview") or []), + asset_bboxes=None, ) for item in (result.get("features") or []) ] ctx.blackboard.page_features = sorted(features, key=lambda f: f.page) ctx.blackboard.page_count = int(result.get("page_count") or len(features)) ctx.blackboard.global_signals["total_pages"] = ctx.blackboard.page_count - logger.info("[document_agent] probed {} pages", ctx.blackboard.page_count) + 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 +384,80 @@ 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 + ), + ) + ) + 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/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/_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..0c09005b 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] @@ -772,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, @@ -799,6 +801,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, @@ -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) ──────────────────────────────── @@ -878,34 +897,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_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/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 6e2a2aab..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 @@ -49,6 +47,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( @@ -100,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) @@ -119,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, @@ -144,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, *, @@ -231,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", @@ -367,19 +284,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 +397,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/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()) 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..839b6165 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -56,11 +56,36 @@ def _page_feature(page: int = 1) -> PageFeature: orientation="portrait", width=72.0, height=72.0, + has_asset=False, is_blank_like=False, - text_lines_preview=["Section 1"], + asset_bboxes=None, ) +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"] @@ -112,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", @@ -135,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( @@ -234,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", @@ -323,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, @@ -333,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", @@ -426,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] = [] @@ -658,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, @@ -685,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) @@ -698,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( @@ -820,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/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, 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/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( 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..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,12 +15,11 @@ 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 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 @@ -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 35323256..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 @@ -747,47 +717,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. @@ -816,15 +746,38 @@ 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. - - 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. + - "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. + - 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.